smp: fix proxy reconnection to relay after restart (#1806)

* tests: add SMP proxy relay reconnection tests

Reproduces the proxy failing to reconnect to a destination relay when the
sender disconnects mid-connection (empty session var left in smpClients).

* fix: bracket session var creation to drop it on interrupt

getSessVar inserts an empty session var that the connect path then fills with
putTMVar. If the connecting thread is killed by an async exception before that
fill (a proxy worker on client disconnect, an agent worker on cancel), the empty
var was left in the map forever and every later request for that server blocked
on it until timing out (permanent PCEResponseTimeout).

Wrap get-or-create with withGetSessVar (bracketOnError) at the call sites, so the
cleanup is established where the var is created and covers the whole connect: on
interrupt before fill the still-empty var is dropped and the next request
reconnects. This closes the window between getSessVar and the fill that a handler
installed inside the connect function cannot cover.

* test: cover session var leak on interrupted connect

UtilTests: tryAllErrors rethrows ThreadKilled/StackOverflow (the mechanism
that skips putTMVar). SMPProxyTests: agent client reconnection after a
cancelled connect, plus a control proving the stalling relay alone does not
cause the failure; refine the relay reconnection tests.

* refactor

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>

src/Simplex/Messaging/Agent/Client.hs

@@ -674,8 +674,7 @@ getSMPServerClient :: AgentClient -> NetworkRequestMode -> SMPTransportSession -
 getSMPServerClient c@AgentClient {active, smpClients, workerSeq} nm tSess = do
   unlessM (readTVarIO active) $ throwE INACTIVE
   ts <- liftIO getCurrentTime
-  atomically (getSessVar workerSeq tSess smpClients ts)
-    >>= either newClient (waitForProtocolClient c nm tSess smpClients)
+  withGetSessVar workerSeq tSess smpClients ts newClient (waitForProtocolClient c nm tSess smpClients)
   where
     newClient v = do
       prs <- liftIO TM.emptyIO
@@ -686,29 +685,25 @@ getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq
   unlessM (readTVarIO active) $ throwE INACTIVE
   proxySrv <- maybe (getNextServer c userId proxySrvs [destSrv]) pure proxySrv_
   ts <- liftIO getCurrentTime
-  atomically (getClientVar proxySrv ts) >>= \(tSess, auth, v) ->
-    either (newProxyClient tSess auth ts) (waitForProxyClient tSess auth) v
+  (tSess, auth) <- atomically $ proxiedTransportSession proxySrv
+  withGetSessVar workerSeq tSess smpClients ts (newProxyClient tSess auth ts) (waitForProxyClient tSess auth)
   where
-    getClientVar :: SMPServerWithAuth -> UTCTime -> STM (SMPTransportSession, Maybe SMP.BasicAuth, Either SMPClientVar SMPClientVar)
-    getClientVar proxySrv ts = do
+    proxiedTransportSession :: SMPServerWithAuth -> STM (SMPTransportSession, Maybe SMP.BasicAuth)
+    proxiedTransportSession proxySrv = do
       ProtoServerWithAuth srv auth <- TM.lookup destSess smpProxiedRelays >>= maybe (TM.insert destSess proxySrv smpProxiedRelays $> proxySrv) pure
-      let tSess = (userId, srv, qId)
-      (tSess,auth,) <$> getSessVar workerSeq tSess smpClients ts
+      pure ((userId, srv, qId), auth)
     newProxyClient :: SMPTransportSession -> Maybe SMP.BasicAuth -> UTCTime -> SMPClientVar -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
     newProxyClient tSess auth ts v = do
       prs <- liftIO TM.emptyIO
-      -- we do not need to check if it is a new proxied relay session,
-      -- as the client is just created and there are no sessions yet
-      rv <- atomically $ either id id <$> getSessVar workerSeq destSrv prs ts
       clnt <- smpConnectClient c nm tSess prs v
-      (clnt,) <$> newProxiedRelay clnt auth rv
+      -- the relay var is always new (the client is just created and has no sessions yet)
+      sess <- withGetSessVar workerSeq destSrv prs ts (newProxiedRelay clnt auth) (waitForProxiedRelay tSess)
+      pure (clnt, sess)
     waitForProxyClient :: SMPTransportSession -> Maybe SMP.BasicAuth -> SMPClientVar -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
     waitForProxyClient tSess auth v = do
       clnt@(SMPConnectedClient _ prs) <- waitForProtocolClient c nm tSess smpClients v
       ts <- liftIO getCurrentTime
-      sess <-
-        atomically (getSessVar workerSeq destSrv prs ts)
-          >>= either (newProxiedRelay clnt auth) (waitForProxiedRelay tSess)
+      sess <- withGetSessVar workerSeq destSrv prs ts (newProxiedRelay clnt auth) (waitForProxiedRelay tSess)
       pure (clnt, sess)
     newProxiedRelay :: SMPConnectedClient -> Maybe SMP.BasicAuth -> ProxiedRelayVar -> AM (Either AgentErrorType ProxiedRelay)
     newProxiedRelay (SMPConnectedClient smp prs) proxyAuth rv =
@@ -846,10 +841,7 @@ getNtfServerClient :: AgentClient -> NetworkRequestMode -> NtfTransportSession -
 getNtfServerClient c@AgentClient {active, ntfClients, workerSeq, proxySessTs, presetDomains} nm tSess@(_, srv, _) = do
   unlessM (readTVarIO active) $ throwE INACTIVE
   ts <- liftIO getCurrentTime
-  atomically (getSessVar workerSeq tSess ntfClients ts)
-    >>= either
-      (newProtocolClient c tSess ntfClients connectClient)
-      (waitForProtocolClient c nm tSess ntfClients)
+  withGetSessVar workerSeq tSess ntfClients ts (newProtocolClient c tSess ntfClients connectClient) (waitForProtocolClient c nm tSess ntfClients)
   where
     connectClient :: NtfClientVar -> AM NtfClient
     connectClient v = do
@@ -870,10 +862,7 @@ getXFTPServerClient :: AgentClient -> XFTPTransportSession -> AM XFTPClient
 getXFTPServerClient c@AgentClient {active, xftpClients, workerSeq, proxySessTs, presetDomains} tSess@(_, srv, _) = do
   unlessM (readTVarIO active) $ throwE INACTIVE
   ts <- liftIO getCurrentTime
-  atomically (getSessVar workerSeq tSess xftpClients ts)
-    >>= either
-      (newProtocolClient c tSess xftpClients connectClient)
-      (waitForProtocolClient c NRMBackground tSess xftpClients)
+  withGetSessVar workerSeq tSess xftpClients ts (newProtocolClient c tSess xftpClients connectClient) (waitForProtocolClient c NRMBackground tSess xftpClients)
   where
     connectClient :: XFTPClientVar -> AM XFTPClient
     connectClient v = do

src/Simplex/Messaging/Client/Agent.hs

@@ -205,11 +205,8 @@ getSMPServerClient' ca srv = snd <$> getSMPServerClient'' ca srv
 getSMPServerClient'' :: SMPClientAgent p -> SMPServer -> ExceptT SMPClientError IO (OwnServer, SMPClient)
 getSMPServerClient'' ca@SMPClientAgent {agentCfg, smpClients, smpSessions, workerSeq} srv = do
   ts <- liftIO getCurrentTime
-  atomically (getClientVar ts) >>= either (ExceptT . newSMPClient) waitForSMPClient
+  withGetSessVar workerSeq srv smpClients ts (ExceptT . newSMPClient) waitForSMPClient
   where
-    getClientVar :: UTCTime -> STM (Either SMPClientVar SMPClientVar)
-    getClientVar = getSessVar workerSeq srv smpClients
-
     waitForSMPClient :: SMPClientVar -> ExceptT SMPClientError IO (OwnServer, SMPClient)
     waitForSMPClient v = do
       let ProtocolClientConfig {networkConfig = NetworkConfig {tcpConnectTimeout}} = smpCfg agentCfg

src/Simplex/Messaging/Notifications/Server.hs

@@ -660,13 +660,14 @@ guardPushProvider token action =
 getOrCreatePushWorker :: NtfPushServer -> (Maybe T.Text, PushProvider) -> OwnServer -> M (TBQueue (NtfTknRec, PushNotification))
 getOrCreatePushWorker s@NtfPushServer {pushWorkers, pushWorkerSeq, pushQSize} key@(srvHost_, _) isOwn = do
   ts <- liftIO getCurrentTime
-  atomically (getSessVar pushWorkerSeq key pushWorkers ts) >>= \case
-    Left v -> do
+  withGetSessVar' pushWorkerSeq key pushWorkers ts createWorker existingWorker
+  where
+    createWorker v = do
       q <- liftIO $ newTBQueueIO pushQSize
       tId <- mkWeakThreadId =<< forkIO (runPushWorker s srvHost_ isOwn q)
       atomically $ putTMVar (sessionVar v) PushWorker {workerQ = q, workerThreadId = tId}
       pure q
-    Right v -> workerQ <$> atomically (readTMVar $ sessionVar v)
+    existingWorker v = workerQ <$> atomically (readTMVar $ sessionVar v)
 
 runPushWorker :: NtfPushServer -> Maybe T.Text -> OwnServer -> TBQueue (NtfTknRec, PushNotification) -> M ()
 runPushWorker s srvHost_ isOwn q = forever $ do

src/Simplex/Messaging/Session.hs

@@ -6,14 +6,20 @@ module Simplex.Messaging.Session
   ( SessionVar (..),
     getSessVar,
     removeSessVar,
+    withGetSessVar,
+    withGetSessVar',
     tryReadSessVar,
   ) where
 
 import Control.Concurrent.STM
+import Control.Monad.Except (ExceptT (..), runExceptT)
+import Control.Monad.IO.Class (liftIO)
+import Control.Monad.IO.Unlift (MonadUnliftIO)
 import Data.Time (UTCTime)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (($>>=))
+import Simplex.Messaging.Util (whenM, ($>>=))
+import UnliftIO.Exception (bracketOnError)
 
 data SessionVar a = SessionVar
   { sessionVar :: TMVar a,
@@ -38,5 +44,31 @@ removeSessVar v sessKey vs =
     Just v' | sessionVarId v == sessionVarId v' -> TM.delete sessKey vs
     _ -> pure ()
 
+-- | Get or create a session var and route to onNew (newly created) or onExisting. The new-var
+-- branch is bracketed from the point of creation: if it is interrupted before filling the var
+-- (e.g. an async exception during connect), the still-empty var is dropped from the map so the
+-- next request creates a fresh session instead of blocking on a var that will never be filled.
+-- A thrown ExceptT error is a normal result (the var keeps the error it was filled with) - only
+-- an interrupting exception drops the empty var.
+withGetSessVar ::
+  (Ord k, MonadUnliftIO m) =>
+  TVar Int -> k -> TMap k (SessionVar a) -> UTCTime ->
+  (SessionVar a -> ExceptT e m b) -> (SessionVar a -> ExceptT e m b) -> ExceptT e m b
+withGetSessVar sessSeq sessKey vs ts onNew onExisting =
+  ExceptT $ withGetSessVar' sessSeq sessKey vs ts (runExceptT . onNew) (runExceptT . onExisting)
+
+-- | withGetSessVar for actions in the underlying monad (without ExceptT).
+withGetSessVar' ::
+  (Ord k, MonadUnliftIO m) =>
+  TVar Int -> k -> TMap k (SessionVar a) -> UTCTime ->
+  (SessionVar a -> m b) -> (SessionVar a -> m b) -> m b
+withGetSessVar' sessSeq sessKey vs ts onNew onExisting =
+  bracketOnError
+    (liftIO $ atomically $ getSessVar sessSeq sessKey vs ts)
+    (either (liftIO . atomically . dropEmptySessVar) (\_ -> pure ()))
+    (either onNew onExisting)
+  where
+    dropEmptySessVar v = whenM (isEmptyTMVar $ sessionVar v) $ removeSessVar v sessKey vs
+
 tryReadSessVar :: Ord k => k -> TMap k (SessionVar a) -> STM (Maybe a)
 tryReadSessVar sessKey vs = TM.lookup sessKey vs $>>= (tryReadTMVar . sessionVar)