limit concurrency when resubscribing, method to resubscribe if not subscribed (#368)

src/Simplex/Messaging/Agent.hs

@@ -43,6 +43,7 @@ module Simplex.Messaging.Agent
     acceptContact,
     rejectContact,
     subscribeConnection,
+    resubscribeConnection,
     sendMessage,
     ackMessage,
     suspendConnection,
@@ -139,6 +140,9 @@ rejectContact c = withAgentEnv c .: rejectContact' c
 subscribeConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
 subscribeConnection c = withAgentEnv c . subscribeConnection' c
 
+resubscribeConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
+resubscribeConnection c = withAgentEnv c . resubscribeConnection' c
+
 -- | Send message to the connection (SEND command)
 sendMessage :: AgentErrorMonad m => AgentClient -> ConnId -> MsgBody -> m AgentMsgId
 sendMessage c = withAgentEnv c .: sendMessage' c
@@ -355,6 +359,12 @@ subscribeConnection' c connId =
     SomeConn _ (RcvConnection _ rq) -> subscribeQueue c rq connId
     SomeConn _ (ContactConnection _ rq) -> subscribeQueue c rq connId
 
+resubscribeConnection' :: forall m. AgentMonad m => AgentClient -> ConnId -> m ()
+resubscribeConnection' c connId =
+  unlessM
+    (atomically $ hasActiveSubscription c connId)
+    (subscribeConnection' c connId)
+
 -- | Send message to the connection (SEND command) in Reader monad
 sendMessage' :: forall m. AgentMonad m => AgentClient -> ConnId -> MsgBody -> m AgentMsgId
 sendMessage' c connId msg =
@@ -398,7 +408,7 @@ resumeMsgDelivery c connId sq@SndQueue {server, sndId} = do
     withStore (`getPendingMsgs` connId)
       >>= queuePendingMsgs c connId sq
   where
-    queueDelivering qKey = atomically $ isJust <$> TM.lookup qKey (smpQueueMsgDeliveries c)
+    queueDelivering qKey = atomically $ TM.member qKey (smpQueueMsgDeliveries c)
     connQueued = atomically $ isJust <$> TM.lookupInsert connId True (connMsgsQueued c)
 
 queuePendingMsgs :: AgentMonad m => AgentClient -> ConnId -> SndQueue -> [InternalId] -> m ()

src/Simplex/Messaging/Agent/Client.hs

@@ -38,6 +38,7 @@ module Simplex.Messaging.Agent.Client
     deleteQueue,
     logServer,
     removeSubscription,
+    hasActiveSubscription,
   )
 where
 
@@ -55,7 +56,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 (catMaybes, isNothing)
+import Data.Maybe (catMaybes)
 import Data.Text.Encoding
 import Data.Word (Word16)
 import Simplex.Messaging.Agent.Env.SQLite
@@ -75,7 +76,7 @@ import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Util (bshow, catchAll_, ifM, liftEitherError, liftError, tryError, unlessM, whenM)
 import Simplex.Messaging.Version
 import System.Timeout (timeout)
-import UnliftIO (async, forConcurrently)
+import UnliftIO (async, pooledForConcurrentlyN)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -200,14 +201,15 @@ getSMPServerClient c@AgentClient {active, smpClients, msgQ} srv = do
         reconnectClient `catchError` const loop
 
     reconnectClient :: m ()
-    reconnectClient =
+    reconnectClient = do
+      n <- asks $ resubscriptionConcurrency . config
       withAgentLock c . withClient c srv $ \smp -> do
         cs <- atomically $ mapM readTVar =<< TM.lookup srv (pendingSubscrSrvrs c)
-        conns <- forConcurrently (maybe [] M.toList cs) $ \sub@(connId, _) ->
+        conns <- pooledForConcurrentlyN n (maybe [] M.toList cs) $ \sub@(connId, _) ->
           ifM
-            (atomically $ isNothing <$> TM.lookup connId (subscrConns c))
-            (subscribe_ smp sub `catchError` handleError connId)
+            (atomically $ hasActiveSubscription c connId)
             (pure $ Just connId)
+            (subscribe_ smp sub `catchError` handleError connId)
         liftIO . unless (null conns) . notifySub "" . UP srv $ catMaybes conns
       where
         subscribe_ :: SMPClient -> (ConnId, RcvQueue) -> ExceptT ProtocolClientError IO (Maybe ConnId)
@@ -408,6 +410,9 @@ addSubscription c rq@RcvQueue {server} connId = atomically $ do
   addSubs_ (subscrSrvrs c) rq connId
   removePendingSubscription c server connId
 
+hasActiveSubscription :: AgentClient -> ConnId -> STM Bool
+hasActiveSubscription c connId = TM.member connId (subscrConns c)
+
 addPendingSubscription :: AgentClient -> RcvQueue -> ConnId -> STM ()
 addPendingSubscription = addSubs_ . pendingSubscrSrvrs
 

src/Simplex/Messaging/Agent/Env/SQLite.hs

@@ -49,6 +49,7 @@ data AgentConfig = AgentConfig
     ntfCfg :: ProtocolClientConfig,
     reconnectInterval :: RetryInterval,
     helloTimeout :: NominalDiffTime,
+    resubscriptionConcurrency :: Int,
     caCertificateFile :: FilePath,
     privateKeyFile :: FilePath,
     certificateFile :: FilePath
@@ -78,6 +79,7 @@ defaultAgentConfig =
       ntfCfg = defaultClientConfig {defaultTransport = ("443", transport @TLS)},
       reconnectInterval = defaultReconnectInterval,
       helloTimeout = 2 * nominalDay,
+      resubscriptionConcurrency = 16,
       -- CA certificate private key is not needed for initialization
       -- ! we do not generate these
       caCertificateFile = "/etc/opt/simplex-agent/ca.crt",

src/Simplex/Messaging/Server.hs

@@ -343,7 +343,7 @@ client clnt@Client {subscriptions, ntfSubscriptions, rcvQ, sndQ} Server {subscri
 
         subscribeNotifications :: m (Transmission BrokerMsg)
         subscribeNotifications = atomically $ do
-          whenM (isNothing <$> TM.lookup queueId ntfSubscriptions) $ do
+          unlessM (TM.member queueId ntfSubscriptions) $ do
             writeTBQueue ntfSubscribedQ (queueId, clnt)
             TM.insert queueId () ntfSubscriptions
           pure ok