refactor server and store

src/ConnStore/STM.hs

@@ -31,14 +31,12 @@ newConnStore = newTVar ConnStoreData {connections = M.empty, senders = M.empty}
 instance MonadUnliftIO m => MonadConnStore STMConnStore m where
   addConn :: STMConnStore -> RecipientKey -> m (Either ErrorType Connection)
   addConn store rKey = atomically $ do
-    db <- readTVar store
     let c@Connection {recipientId = rId, senderId = sId} = mkConnection rKey
-        db' =
-          db
-            { connections = M.insert rId c (connections db),
-              senders = M.insert sId rId (senders db)
-            }
-    writeTVar store db'
+    modifyTVar store $ \db ->
+      db
+        { connections = M.insert rId c (connections db),
+          senders = M.insert sId rId (senders db)
+        }
     return $ Right c
 
   getConn :: STMConnStore -> Sing (p :: Party) -> ConnId -> m (Either ErrorType Connection)
@@ -52,23 +50,26 @@ instance MonadUnliftIO m => MonadConnStore STMConnStore m where
   getConn _ SBroker _ =
     return $ Left INTERNAL
 
-  secureConn store rId sKey = updateConnections store rId $ \db c ->
-    case senderKey c of
-      Just _ -> (Left AUTH, db)
-      _ -> (Right (), db {connections = M.insert rId c {senderKey = Just sKey} (connections db)})
+  secureConn store rId sKey =
+    updateConnections store rId $ \db c ->
+      case senderKey c of
+        Just _ -> (Left AUTH, db)
+        _ -> (Right (), db {connections = M.insert rId c {senderKey = Just sKey} (connections db)})
 
   suspendConn :: STMConnStore -> RecipientId -> m (Either ErrorType ())
-  suspendConn store rId = updateConnections store rId $ \db c ->
-    (Right (), db {connections = M.insert rId c {status = ConnOff} (connections db)})
+  suspendConn store rId =
+    updateConnections store rId $ \db c ->
+      (Right (), db {connections = M.insert rId c {status = ConnOff} (connections db)})
 
   deleteConn :: STMConnStore -> RecipientId -> m (Either ErrorType ())
-  deleteConn store rId = updateConnections store rId $ \db c ->
-    ( Right (),
-      db
-        { connections = M.delete rId (connections db),
-          senders = M.delete (senderId c) (senders db)
-        }
-    )
+  deleteConn store rId =
+    updateConnections store rId $ \db c ->
+      ( Right (),
+        db
+          { connections = M.delete rId (connections db),
+            senders = M.delete (senderId c) (senders db)
+          }
+      )
 
 updateConnections ::
   MonadUnliftIO m =>

src/MsgStore/STM.hs

@@ -34,9 +34,9 @@ instance MonadUnliftIO m => MonadMsgStore STMMsgStore MsgQueue m where
         return q
 
   delMsgQueue :: STMMsgStore -> RecipientId -> m ()
-  delMsgQueue store rId = atomically $ do
-    m <- messages <$> readTVar store
-    writeTVar store . MsgStoreData $ M.delete rId m
+  delMsgQueue store rId = atomically . modifyTVar store $
+    \(MsgStoreData ms) ->
+      MsgStoreData $ M.delete rId ms
 
 instance MonadUnliftIO m => MonadMsgQueue MsgQueue m where
   writeMsg :: MsgQueue -> Message -> m ()

src/Server.hs

@@ -26,6 +26,7 @@ import Transmission
 import Transport
 import UnliftIO.Async
 import UnliftIO.Concurrent
+import qualified UnliftIO.Exception as E
 import UnliftIO.IO
 import UnliftIO.STM
 
@@ -113,26 +114,15 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
       case cmd of
         Cmd SBroker END -> unsubscribeConn >> return (connId, cmd)
         Cmd SBroker _ -> return (connId, cmd)
-        Cmd SSender (SEND msgBody) ->
-          getConn st SSender connId
-            >>= fmap (mkSigned connId) . either (return . ERR) (storeMessage msgBody)
+        Cmd SSender (SEND msgBody) -> sendMessage st msgBody
         Cmd SRecipient command -> case command of
           CONN rKey -> createConn st rKey
           SUB -> subscribeConn connId
-          ACK -> deliverMessage tryDelPeekMsg -- TODO? sending ACK without message loses the message
+          ACK -> deliverMessage tryDelPeekMsg connId -- TODO? sending ACK without message loses the message
           KEY sKey -> okResponse <$> secureConn st connId sKey
           OFF -> okResponse <$> suspendConn st connId
           DEL -> okResponse <$> deleteConn st connId
       where
-        ok :: Signed
-        ok = (connId, Cmd SBroker OK)
-
-        mkSigned :: ConnId -> Command 'Broker -> Signed
-        mkSigned cId command = (cId, Cmd SBroker command)
-
-        okResponse :: Either ErrorType () -> Signed
-        okResponse = mkSigned connId . either ERR (const OK)
-
         createConn :: MonadConnStore s m => s -> RecipientKey -> m Signed
         createConn st rKey =
           addConn st rKey >>= \case
@@ -148,17 +138,20 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
             when (M.notMember rId cs) $ do
               writeTBQueue subscribedQ (rId, clnt)
               writeTVar connections $ M.insert rId (Left ()) cs
-          deliverMessage tryPeekMsg
+          deliverMessage tryPeekMsg rId
 
         unsubscribeConn :: m ()
         unsubscribeConn = do
           cs <- readTVarIO connections
+          atomically . writeTVar connections $ M.delete connId cs
           case M.lookup connId cs of
-            Nothing -> return ()
-            Just (Left ()) -> atomically $ writeTVar connections $ M.delete connId cs
-            Just (Right threadId) -> do
-              killThread threadId
-              atomically $ writeTVar connections $ M.delete connId cs
+            Just (Right threadId) -> killThread threadId
+            _ -> return ()
+
+        sendMessage :: MonadConnStore s m => s -> MsgBody -> m Signed
+        sendMessage st msgBody =
+          getConn st SSender connId
+            >>= fmap (mkSigned connId) . either (return . ERR) (storeMessage msgBody)
 
         storeMessage :: MsgBody -> Connection -> m (Command 'Broker)
         storeMessage msgBody c = case status c of
@@ -170,25 +163,37 @@ client clnt@Client {connections, rcvQ, sndQ} Server {subscribedQ} =
             return OK
           ConnOff -> return $ ERR AUTH
 
-        deliverMessage :: (MsgQueue -> m (Maybe Message)) -> m Signed
-        deliverMessage tryPeek = do
+        deliverMessage :: (MsgQueue -> m (Maybe Message)) -> RecipientId -> m Signed
+        deliverMessage tryPeek rId = do
           ms <- asks msgStore
-          q <- getMsgQueue ms connId
+          q <- getMsgQueue ms rId
           tryPeek q >>= \case
-            Just Message {msgId, ts, msgBody} ->
-              return . mkSigned connId $ MSG msgId ts msgBody
-            Nothing -> do
-              cs <- readTVarIO connections
-              case M.lookup connId cs of
-                Nothing -> return ok
-                Just (Right _) -> return ok
-                Just (Left ()) -> do
-                  threadId <- forkIO $ subscriber q
-                  atomically . writeTVar connections $ M.insert connId (Right threadId) cs
-                  return ok
+            Just msg -> return $ msgResponse msg
+            Nothing -> forkSubscriber q rId
 
-        subscriber :: MsgQueue -> m ()
-        subscriber q = do
-          Message {msgId, ts, msgBody} <- peekMsg q
-          -- TODO refactor with deliver
-          atomically $ writeTBQueue sndQ $ mkSigned connId $ MSG msgId ts msgBody
+        forkSubscriber :: MsgQueue -> RecipientId -> m Signed
+        forkSubscriber q rId = do
+          cs <- readTVarIO connections
+          case M.lookup rId cs of
+            Just (Left ()) -> do
+              E.bracket
+                (forkIO subscriber)
+                (\_ -> trackSubscriber $ Left ())
+                (trackSubscriber . Right)
+              return ok
+            _ -> return ok
+          where
+            trackSubscriber sThrd = atomically . modifyTVar connections $ M.insert rId sThrd
+            subscriber = peekMsg q >>= atomically . writeTBQueue sndQ . msgResponse
+
+        ok :: Signed
+        ok = (connId, Cmd SBroker OK)
+
+        mkSigned :: ConnId -> Command 'Broker -> Signed
+        mkSigned cId command = (cId, Cmd SBroker command)
+
+        okResponse :: Either ErrorType () -> Signed
+        okResponse = mkSigned connId . either ERR (const OK)
+
+        msgResponse :: Message -> Signed
+        msgResponse Message {msgId, ts, msgBody} = mkSigned connId $ MSG msgId ts msgBody