core: fix memory leak in SMP reconnects

src/Simplex/Messaging/Agent.hs

@@ -14,7 +14,7 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# OPTIONS_GHC -fno-warn-ambiguous-fields #-}
 
 -- |
 -- Module      : Simplex.Messaging.Agent
@@ -117,6 +117,7 @@
   )
 where
 
+import Control.Concurrent.STM (retry)
 import Control.Logger.Simple (logError, logInfo, showText)
 import Control.Monad
 import Control.Monad.Except
@@ -138,6 +139,7 @@
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
 import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing, mapMaybe)
+import qualified Data.Set as S
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Time.Clock
@@ -179,7 +181,6 @@
 import Simplex.RemoteControl.Client
 import Simplex.RemoteControl.Invitation
 import Simplex.RemoteControl.Types
-import UnliftIO.Async (race_)
 import UnliftIO.Concurrent (forkFinally, forkIO, threadDelay)
 import UnliftIO.STM
 
@@ -200,7 +201,16 @@
       pure c
     runAgentThreads c
       | backgroundMode = subscriber c
-      | otherwise = raceAny_ [subscriber c, runNtfSupervisor c, cleanupManager c]
+      | otherwise = raceAny_ [subscriber c, smpSupervisor c, runNtfSupervisor c, cleanupManager c]
+
+smpSupervisor :: AgentMonad' m => AgentClient -> m ()
+smpSupervisor c@AgentClient {smpSubRequests} = forever $ do
+  subs <- atomically $ do
+    subs <- readTVar smpSubRequests
+    if S.null subs then retry else writeTVar smpSubRequests S.empty
+    pure subs
+  mapM_ (resubscribeSMPSession c) subs
+  threadDelay 1000000 -- try to aggregate reconnect requests in the Set
 
 disconnectAgentClient :: MonadUnliftIO m => AgentClient -> m ()
 disconnectAgentClient c@AgentClient {agentEnv = Env {ntfSupervisor = ns, xftpAgent = xa}} = do
@@ -504,7 +514,7 @@
 
 -- | Runs an SMP agent instance that receives commands and sends responses via 'TBQueue's.
 runAgentClient :: AgentMonad' m => AgentClient -> m ()
-runAgentClient c = race_ (subscriber c) (client c)
+runAgentClient c = raceAny_ [subscriber c, client c, smpSupervisor c]
 
 client :: forall m. AgentMonad' m => AgentClient -> m ()
 client c@AgentClient {rcvQ, subQ} = forever $ do
@@ -730,7 +740,7 @@
   AgentConfig {smpClientVRange, smpAgentVRange} <- asks config
   pure $
     (,)
-      <$> (qUri `compatibleVersion` smpClientVRange) 
+      <$> (qUri `compatibleVersion` smpClientVRange)
       <*> (crAgentVRange `compatibleVersion` smpAgentVRange pqSup)
 
 versionPQSupport_ :: VersionSMPA -> Maybe CR.VersionE2E -> PQSupport

src/Simplex/Messaging/Agent/Client.hs

@@ -30,6 +30,7 @@
     closeAgentClient,
     closeProtocolServerClients,
     reconnectServerClients,
+    resubscribeSMPSession,
     closeXFTPServerClient,
     runSMPServerTest,
     runXFTPServerTest,
@@ -286,6 +287,7 @@
     deleteLock :: Lock,
     -- smpSubWorkers for SMP servers sessions
     smpSubWorkers :: TMap SMPTransportSession (SessionVar (Async ())),
+    smpSubRequests :: TVar (Set SMPTransportSession),
     agentStats :: TMap AgentStatsKey (TVar Int),
     clientId :: Int,
     agentEnv :: Env
@@ -425,6 +427,7 @@
   invLocks <- TM.empty
   deleteLock <- createLock
   smpSubWorkers <- TM.empty
+  smpSubRequests <- newTVar mempty
   agentStats <- TM.empty
   return
     AgentClient
@@ -458,6 +461,7 @@
         invLocks,
         deleteLock,
         smpSubWorkers,
+        smpSubRequests,
         agentStats,
         clientId,
         agentEnv
@@ -516,16 +520,15 @@
     -- make it expensive to check for pending subscriptions.
     newClient v =
       newProtocolClient c tSess smpClients connectClient v
-        `catchAgentError` \e -> resubscribeSMPSession c tSess >> throwError e
+        `catchAgentError` \e -> submitSMPResubscribe c tSess >> throwError e
     connectClient :: SMPClientVar -> m SMPClient
     connectClient v = do
       cfg <- getClientConfig c smpCfg
       g <- asks random
-      u <- askUnliftIO
-      liftEitherError (protocolClientError SMP $ B.unpack $ strEncode srv) (getProtocolClient g tSess cfg (Just msgQ) $ clientDisconnected u v)
+      liftEitherError (protocolClientError SMP $ B.unpack $ strEncode srv) $ getProtocolClient g tSess cfg (Just msgQ) (clientDisconnected v)
 
-    clientDisconnected :: UnliftIO m -> SMPClientVar -> SMPClient -> IO ()
-    clientDisconnected u v client = do
+    clientDisconnected :: SMPClientVar -> SMPClient -> IO ()
+    clientDisconnected v client = do
       removeClientAndSubs >>= serverDown
       logInfo . decodeUtf8 $ "Agent disconnected from " <> showServer srv
       where
@@ -548,11 +551,18 @@
           unless (null conns) $ notifySub "" $ DOWN srv conns
           unless (null qs) $ do
             atomically $ mapM_ (releaseGetLock c) qs
-            unliftIO u $ resubscribeSMPSession c tSess
+            submitSMPResubscribe c tSess
 
         notifySub :: forall e. AEntityI e => ConnId -> ACommand 'Agent e -> IO ()
         notifySub connId cmd = atomically $ writeTBQueue (subQ c) ("", connId, APC (sAEntity @e) cmd)
 
+-- | Schedule SMP resubscription job
+--
+-- Attempting to resubscribe directly in exception handlers results in untraceable resource leaking.
+-- A supervisor agent thread starts workers to process session resubscription requests instead allowing the handler to finish quickly.
+submitSMPResubscribe :: MonadIO m => AgentClient -> SMPTransportSession -> m ()
+submitSMPResubscribe c = atomically . modifyTVar' (smpSubRequests c) . S.insert
+
 resubscribeSMPSession :: AgentMonad' m => AgentClient -> SMPTransportSession -> m ()
 resubscribeSMPSession c@AgentClient {smpSubWorkers} tSess =
   atomically getWorkerVar >>= mapM_ (either newSubWorker (\_ -> pure ()))
@@ -1066,19 +1076,18 @@
   atomically $ do
     modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId qs'))
     RQ.batchAddQueues (pendingSubs c) qs'
-  u <- askUnliftIO
   -- only "checked" queues are subscribed
-  (errs <>) <$> sendTSessionBatches "SUB" 90 id (subscribeQueues_ u) c qs'
+  (errs <>) <$> sendTSessionBatches "SUB" 90 id subscribeQueues_ c qs'
   where
     checkQueue rq = do
       prohibited <- atomically $ hasGetLock c rq
       pure $ if prohibited then Left (rq, Left $ CMD PROHIBITED) else Right rq
-    subscribeQueues_ :: UnliftIO m -> SMPClient -> NonEmpty RcvQueue -> IO (BatchResponses SMPClientError ())
-    subscribeQueues_ u smp qs' = do
+    subscribeQueues_ :: SMPClient -> NonEmpty RcvQueue -> IO (BatchResponses SMPClientError ())
+    subscribeQueues_ smp qs' = do
       rs <- sendBatch subscribeSMPQueues smp qs'
       mapM_ (uncurry $ processSubResult c) rs
-      when (any temporaryClientError . lefts . map snd $ L.toList rs) . unliftIO u $
-        resubscribeSMPSession c (transportSession' smp)
+      when (any temporaryClientError . lefts . map snd $ L.toList rs) $
+        submitSMPResubscribe c (transportSession' smp)
       pure rs
 
 type BatchResponses e r = (NonEmpty (RcvQueue, Either e r))
@@ -1515,7 +1524,7 @@
   where
     statsKey = AgentStatsKey {userId, host = strEncode $ clientTransportHost pc, clientTs = strEncode $ clientSessionTs pc, cmd, res}
 
 userServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> TMap UserId (NonEmpty (ProtoServerWithAuth p))
 userServers c = case protocolTypeI @p of
   SPSMP -> smpServers c
   SPXFTP -> xftpServers c