batch all connections in DOWN/UP agent messages (#363)

src/Simplex/Messaging/Agent/Client.hs

@@ -55,7 +55,7 @@ import qualified Data.ByteString.Char8 as B
 import Data.List.NonEmpty (NonEmpty)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
-import Data.Maybe (isNothing)
+import Data.Maybe (catMaybes, isNothing)
 import Data.Text.Encoding
 import Data.Word (Word16)
 import Simplex.Messaging.Agent.Env.SQLite
@@ -72,10 +72,10 @@ import Simplex.Messaging.Protocol (BrokerMsg, ErrorType, ProtocolServer (..), Qu
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (bshow, liftEitherError, liftError, tryError, whenM)
+import Simplex.Messaging.Util (bshow, ifM, liftEitherError, liftError, tryError)
 import Simplex.Messaging.Version
 import System.Timeout (timeout)
-import UnliftIO (async, forConcurrently_)
+import UnliftIO (async, forConcurrently)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -181,7 +181,8 @@ getSMPServerClient c@AgentClient {smpClients, msgQ} srv =
 
     serverDown :: UnliftIO m -> Map ConnId RcvQueue -> IO ()
     serverDown u cs = unless (M.null cs) $ do
-      mapM_ (notifySub DOWN) $ M.keysSet cs
+      let conns = M.keys cs
+      unless (null conns) . notifySub "" $ DOWN srv conns
       unliftIO u reconnectServer
 
     reconnectServer :: m ()
@@ -199,26 +200,30 @@ getSMPServerClient c@AgentClient {smpClients, msgQ} srv =
     reconnectClient =
       withAgentLock c . withClient c srv $ \smp -> do
         cs <- atomically $ mapM readTVar =<< TM.lookup srv (pendingSubscrSrvrs c)
-        forConcurrently_ (maybe [] M.toList cs) $ \sub@(connId, _) ->
-          whenM (atomically $ isNothing <$> TM.lookup connId (subscrConns c)) $
-            subscribe_ smp sub `catchError` handleError connId
+        conns <- forConcurrently (maybe [] M.toList cs) $ \sub@(connId, _) ->
+          ifM
+            (atomically $ isNothing <$> TM.lookup connId (subscrConns c))
+            (subscribe_ smp sub `catchError` handleError connId)
+            (pure $ Just connId)
+        liftIO . unless (null conns) . notifySub "" . UP srv $ catMaybes conns
       where
-        subscribe_ :: SMPClient -> (ConnId, RcvQueue) -> ExceptT ProtocolClientError IO ()
+        subscribe_ :: SMPClient -> (ConnId, RcvQueue) -> ExceptT ProtocolClientError IO (Maybe ConnId)
         subscribe_ smp (connId, rq@RcvQueue {rcvPrivateKey, rcvId}) = do
           subscribeSMPQueue smp rcvPrivateKey rcvId
           addSubscription c rq connId
-          liftIO $ notifySub UP connId
+          pure $ Just connId
 
-        handleError :: ConnId -> ProtocolClientError -> ExceptT ProtocolClientError IO ()
+        handleError :: ConnId -> ProtocolClientError -> ExceptT ProtocolClientError IO (Maybe ConnId)
         handleError connId = \case
           e@PCEResponseTimeout -> throwError e
           e@PCENetworkError -> throwError e
           e -> do
-            liftIO $ notifySub (ERR $ protocolClientError SMP e) connId
+            liftIO . notifySub connId . ERR $ protocolClientError SMP e
             atomically $ removePendingSubscription c srv connId
+            pure Nothing
 
-    notifySub :: ACommand 'Agent -> ConnId -> IO ()
-    notifySub cmd connId = atomically $ writeTBQueue (subQ c) ("", connId, cmd)
+    notifySub :: ConnId -> ACommand 'Agent -> IO ()
+    notifySub connId cmd = atomically $ writeTBQueue (subQ c) ("", connId, cmd)
 
 getNtfServerClient :: forall m. AgentMonad m => AgentClient -> NtfServer -> m NtfClient
 getNtfServerClient c@AgentClient {ntfClients} srv =

src/Simplex/Messaging/Agent/Protocol.hs

@@ -210,8 +210,8 @@ data ACommand (p :: AParty) where
   CON :: ACommand Agent -- notification that connection is established
   SUB :: ACommand Client
   END :: ACommand Agent
-  DOWN :: ACommand Agent
-  UP :: ACommand Agent
+  DOWN :: SMPServer -> [ConnId] -> ACommand Agent
+  UP :: SMPServer -> [ConnId] -> ACommand Agent
   SEND :: MsgBody -> ACommand Client
   MID :: AgentMsgId -> ACommand Agent
   SENT :: AgentMsgId -> ACommand Agent
@@ -817,8 +817,8 @@ commandP =
     <|> "INFO " *> infoCmd
     <|> "SUB" $> ACmd SClient SUB
     <|> "END" $> ACmd SAgent END
-    <|> "DOWN" $> ACmd SAgent DOWN
-    <|> "UP" $> ACmd SAgent UP
+    <|> "DOWN " *> downsResp
+    <|> "UP " *> upsResp
     <|> "SEND " *> sendCmd
     <|> "MID " *> msgIdResp
     <|> "SENT " *> sentResp
@@ -840,12 +840,15 @@ commandP =
     acptCmd = ACmd SClient .: ACPT <$> A.takeTill (== ' ') <* A.space <*> A.takeByteString
     rjctCmd = ACmd SClient . RJCT <$> A.takeByteString
     infoCmd = ACmd SAgent . INFO <$> A.takeByteString
+    downsResp = ACmd SAgent .: DOWN <$> strP <* A.space <*> connections
+    upsResp = ACmd SAgent .: UP <$> strP <* A.space <*> connections
     sendCmd = ACmd SClient . SEND <$> A.takeByteString
     msgIdResp = ACmd SAgent . MID <$> A.decimal
     sentResp = ACmd SAgent . SENT <$> A.decimal
     msgErrResp = ACmd SAgent .: MERR <$> A.decimal <* A.space <*> strP
     message = ACmd SAgent .: MSG <$> msgMetaP <* A.space <*> A.takeByteString
     ackCmd = ACmd SClient . ACK <$> A.decimal
+    connections = strP `A.sepBy'` (A.char ',')
     msgMetaP = do
       integrity <- strP
       recipient <- " R=" *> partyMeta A.decimal
@@ -872,8 +875,8 @@ serializeCommand = \case
   INFO cInfo -> "INFO " <> serializeBinary cInfo
   SUB -> "SUB"
   END -> "END"
-  DOWN -> "DOWN"
-  UP -> "UP"
+  DOWN srv conns -> B.unwords ["DOWN", strEncode srv, connections conns]
+  UP srv conns -> B.unwords ["UP", strEncode srv, connections conns]
   SEND msgBody -> "SEND " <> serializeBinary msgBody
   MID mId -> "MID " <> bshow mId
   SENT mId -> "SENT " <> bshow mId
@@ -888,6 +891,8 @@ serializeCommand = \case
   where
     showTs :: UTCTime -> ByteString
     showTs = B.pack . formatISO8601Millis
+    connections :: [ConnId] -> ByteString
+    connections = B.intercalate "," . map strEncode
     serializeMsgMeta :: MsgMeta -> ByteString
     serializeMsgMeta MsgMeta {integrity, recipient = (rmId, rTs), broker = (bmId, bTs), sndMsgId} =
       B.unwords
@@ -939,6 +944,8 @@ tGet party h = liftIO (tGetRaw h) >>= tParseLoadBody
       ACPT {} -> Right cmd
       -- ERROR response does not always have connId
       ERR _ -> Right cmd
+      DOWN {} -> Right cmd
+      UP {} -> Right cmd
       -- other responses must have connId
       _
         | B.null connId -> Left $ CMD NO_CONN

src/Simplex/Messaging/Protocol.hs

@@ -388,6 +388,8 @@ type SMPServer = ProtocolServer
 pattern SMPServer :: HostName -> ServiceName -> C.KeyHash -> ProtocolServer
 pattern SMPServer host port keyHash = ProtocolServer host port keyHash
 
+{-# COMPLETE SMPServer #-}
+
 -- | SMP server location and transport key digest (hash).
 data ProtocolServer = ProtocolServer
   { host :: HostName,