make new smp client attempt to reconnect on network error (#315)

* make new smp client attempt to reconnect on network error

* reconnect async wip

* remove comments

* trigger tryReconnectClient

* add subscription before trying to take smp client

* move atomically modify asyncClients

* remove logs

* corrections

* don't put error into new smp client TMVar

* UnliftIO async

* re-work

* test for connection to the server that was initially down

* fix tests

* additionally test that bob receives the message after re-connect

* make test independent of event order

Co-authored-by: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>

src/Simplex/Messaging/Agent/Client.hs

@@ -37,7 +37,7 @@ module Simplex.Messaging.Agent.Client
 where
 
 import Control.Concurrent (forkIO)
-import Control.Concurrent.Async (Async, async, uninterruptibleCancel)
+import Control.Concurrent.Async (Async, uninterruptibleCancel)
 import Control.Concurrent.STM (stateTVar)
 import Control.Logger.Simple
 import Control.Monad.Except
@@ -62,8 +62,10 @@ import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Protocol (QueueId, QueueIdsKeys (..), SndPublicVerifyKey)
 import qualified Simplex.Messaging.Protocol as SMP
-import Simplex.Messaging.Util (bshow, liftEitherError, liftError, liftIOEither, tryError)
+import Simplex.Messaging.Util (bshow, liftEitherError, liftError, tryError)
 import Simplex.Messaging.Version
+import System.Timeout (timeout)
+import UnliftIO (async)
 import UnliftIO.Exception (Exception, IOException)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
@@ -76,11 +78,13 @@ data AgentClient = AgentClient
     msgQ :: TBQueue SMPServerTransmission,
     smpClients :: TVar (Map SMPServer SMPClientVar),
     subscrSrvrs :: TVar (Map SMPServer (Map ConnId RcvQueue)),
+    pendingSubscrSrvrs :: TVar (Map SMPServer (Map ConnId RcvQueue)),
     subscrConns :: TVar (Map ConnId SMPServer),
     connMsgsQueued :: TVar (Map ConnId Bool),
     smpQueueMsgQueues :: TVar (Map (ConnId, SMPServer, SMP.SenderId) (TQueue InternalId)),
     smpQueueMsgDeliveries :: TVar (Map (ConnId, SMPServer, SMP.SenderId) (Async ())),
     reconnections :: TVar [Async ()],
+    asyncClients :: TVar [Async ()],
     clientId :: Int,
     agentEnv :: Env,
     smpSubscriber :: Async (),
@@ -95,14 +99,16 @@ newAgentClient agentEnv = do
   msgQ <- newTBQueue qSize
   smpClients <- newTVar M.empty
   subscrSrvrs <- newTVar M.empty
+  pendingSubscrSrvrs <- newTVar M.empty
   subscrConns <- newTVar M.empty
   connMsgsQueued <- newTVar M.empty
   smpQueueMsgQueues <- newTVar M.empty
   smpQueueMsgDeliveries <- newTVar M.empty
   reconnections <- newTVar []
+  asyncClients <- newTVar []
   clientId <- stateTVar (clientCounter agentEnv) $ \i -> (i + 1, i + 1)
   lock <- newTMVar ()
-  return AgentClient {rcvQ, subQ, msgQ, smpClients, subscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, reconnections, clientId, agentEnv, smpSubscriber = undefined, lock}
+  return AgentClient {rcvQ, subQ, msgQ, smpClients, subscrSrvrs, pendingSubscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, reconnections, asyncClients, clientId, agentEnv, smpSubscriber = undefined, lock}
 
 -- | Agent monad with MonadReader Env and MonadError AgentErrorType
 type AgentMonad m = (MonadUnliftIO m, MonadReader Env m, MonadError AgentErrorType m)
@@ -133,20 +139,39 @@ getSMPServerClient c@AgentClient {smpClients, msgQ} srv =
       pure smpVar
 
     waitForSMPClient :: TMVar (Either AgentErrorType SMPClient) -> m SMPClient
-    waitForSMPClient = liftIOEither . atomically . readTMVar
+    waitForSMPClient smpVar = do
+      SMPClientConfig {tcpTimeout} <- asks $ smpCfg . config
+      smpClient_ <- liftIO $ tcpTimeout `timeout` atomically (readTMVar smpVar)
+      liftEither $ case smpClient_ of
+        Just (Right smpClient) -> Right smpClient
+        Just (Left e) -> Left e
+        Nothing -> Left $ BROKER TIMEOUT
 
     newSMPClient :: TMVar (Either AgentErrorType SMPClient) -> m SMPClient
-    newSMPClient smpVar =
-      tryError connectClient >>= \r -> case r of
-        Right smp -> do
-          logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv
-          atomically $ putTMVar smpVar r
-          pure smp
-        Left e -> do
-          atomically $ do
-            putTMVar smpVar r
-            modifyTVar smpClients $ M.delete srv
-          throwError e
+    newSMPClient smpVar = tryConnectClient pure tryConnectAsync
+      where
+        tryConnectClient :: (SMPClient -> m a) -> m () -> m a
+        tryConnectClient successAction retryAction =
+          tryError connectClient >>= \r -> case r of
+            Right smp -> do
+              logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv
+              atomically $ putTMVar smpVar r
+              successAction smp
+            Left e -> do
+              if e == BROKER NETWORK || e == BROKER TIMEOUT
+                then retryAction
+                else atomically $ do
+                  putTMVar smpVar (Left e)
+                  modifyTVar smpClients $ M.delete srv
+              throwError e
+        tryConnectAsync :: m ()
+        tryConnectAsync = do
+          a <- async connectAsync
+          atomically $ modifyTVar (asyncClients c) (a :)
+        connectAsync :: m ()
+        connectAsync = do
+          ri <- asks $ reconnectInterval . config
+          withRetryInterval ri $ \loop -> void $ tryConnectClient (const reconnectClient) loop
 
     connectClient :: m SMPClient
     connectClient = do
@@ -169,33 +194,43 @@ getSMPServerClient c@AgentClient {smpClients, msgQ} srv =
       cs <- M.lookup srv <$> readTVar (subscrSrvrs c)
       modifyTVar (subscrSrvrs c) $ M.delete srv
       modifyTVar (subscrConns c) $ maybe id (deleteKeys . M.keysSet) cs
+      mapM_ (modifyTVar (pendingSubscrSrvrs c) . addPendingSubs) cs
       return cs
       where
+        addPendingSubs :: Map ConnId RcvQueue -> Map SMPServer (Map ConnId RcvQueue) -> Map SMPServer (Map ConnId RcvQueue)
+        addPendingSubs cs = M.alter (Just . addSubs cs) srv
+        addSubs cs = maybe cs (M.union cs)
         deleteKeys :: Ord k => Set k -> Map k a -> Map k a
         deleteKeys ks m = S.foldr' M.delete m ks
 
     serverDown :: UnliftIO m -> Map ConnId RcvQueue -> IO ()
     serverDown u cs = unless (M.null cs) $ do
       mapM_ (notifySub DOWN) $ M.keysSet cs
-      a <- async . unliftIO u $ tryReconnectClient cs
+      unliftIO u reconnectServer
+
+    reconnectServer :: m ()
+    reconnectServer = do
+      a <- async tryReconnectClient
       atomically $ modifyTVar (reconnections c) (a :)
 
-    tryReconnectClient :: Map ConnId RcvQueue -> m ()
-    tryReconnectClient cs = do
+    tryReconnectClient :: m ()
+    tryReconnectClient = do
       ri <- asks $ reconnectInterval . config
       withRetryInterval ri $ \loop ->
-        reconnectClient cs `catchError` const loop
+        reconnectClient `catchError` const loop
 
-    reconnectClient :: Map ConnId RcvQueue -> m ()
-    reconnectClient cs = do
+    reconnectClient :: m ()
+    reconnectClient = do
       withAgentLock c . withSMP c srv $ \smp -> do
         subs <- readTVarIO $ subscrConns c
-        forM_ (M.toList cs) $ \(connId, rq@RcvQueue {rcvPrivateKey, rcvId}) ->
+        cs <- M.lookup srv <$> readTVarIO (pendingSubscrSrvrs c)
+        forM_ (maybe [] M.toList cs) $ \(connId, rq@RcvQueue {rcvPrivateKey, rcvId}) ->
           when (isNothing $ M.lookup connId subs) $ do
             subscribeSMPQueue smp rcvPrivateKey rcvId
               `catchError` \case
-                SMPServerError e -> liftIO $ notifySub (ERR $ SMP e) connId
-                e -> throwError e
+                e@SMPResponseTimeout -> throwError e
+                e@SMPNetworkError -> throwError e
+                e -> liftIO $ notifySub (ERR $ smpClientError e) connId
             addSubscription c rq connId
             liftIO $ notifySub UP connId
 
@@ -206,6 +241,7 @@ closeAgentClient :: MonadUnliftIO m => AgentClient -> m ()
 closeAgentClient c = liftIO $ do
   closeSMPServerClients c
   cancelActions $ reconnections c
+  cancelActions $ asyncClients c
   cancelActions $ smpQueueMsgDeliveries c
 
 closeSMPServerClients :: AgentClient -> IO ()
@@ -294,26 +330,40 @@ newRcvQueue_ a c srv = do
 
 subscribeQueue :: AgentMonad m => AgentClient -> RcvQueue -> ConnId -> m ()
 subscribeQueue c rq@RcvQueue {server, rcvPrivateKey, rcvId} connId = do
-  withLogSMP c server rcvId "SUB" $ \smp ->
-    subscribeSMPQueue smp rcvPrivateKey rcvId
-  addSubscription c rq connId
+  addPendingSubscription c rq connId
+  withLogSMP c server rcvId "SUB" $ \smp -> do
+    liftIO (runExceptT $ subscribeSMPQueue smp rcvPrivateKey rcvId) >>= \case
+      Left e -> do
+        atomically . when (e /= SMPNetworkError && e /= SMPResponseTimeout) $
+          removePendingSubscription c server connId
+        throwError e
+      Right _ -> addSubscription c rq connId
 
 addSubscription :: MonadUnliftIO m => AgentClient -> RcvQueue -> ConnId -> m ()
 addSubscription c rq@RcvQueue {server} connId = atomically $ do
   modifyTVar (subscrConns c) $ M.insert connId server
-  modifyTVar (subscrSrvrs c) $ M.alter (Just . addSub) server
-  where
-    addSub :: Maybe (Map ConnId RcvQueue) -> Map ConnId RcvQueue
-    addSub (Just cs) = M.insert connId rq cs
-    addSub _ = M.singleton connId rq
+  addSubs_ (subscrSrvrs c) rq connId
+  removePendingSubscription c server connId
 
-removeSubscription :: AgentMonad m => AgentClient -> ConnId -> m ()
-removeSubscription AgentClient {subscrConns, subscrSrvrs} connId = atomically $ do
-  cs <- readTVar subscrConns
-  writeTVar subscrConns $ M.delete connId cs
-  mapM_
-    (modifyTVar subscrSrvrs . M.alter (>>= delSub))
-    (M.lookup connId cs)
+addPendingSubscription :: MonadUnliftIO m => AgentClient -> RcvQueue -> ConnId -> m ()
+addPendingSubscription c rq connId =
+  atomically $ addSubs_ (pendingSubscrSrvrs c) rq connId
+
+addSubs_ :: TVar (Map SMPServer (Map ConnId RcvQueue)) -> RcvQueue -> ConnId -> STM ()
+addSubs_ ss rq@RcvQueue {server} connId = modifyTVar ss $ M.alter (Just . addSub) server
+  where
+    addSub = maybe (M.singleton connId rq) (M.insert connId rq)
+
+removeSubscription :: MonadUnliftIO m => AgentClient -> ConnId -> m ()
+removeSubscription c@AgentClient {subscrConns} connId = atomically $ do
+  server_ <- stateTVar subscrConns $ \cs -> (M.lookup connId cs, M.delete connId cs)
+  mapM_ (\server -> removeSubs_ (subscrSrvrs c) server connId) server_
+
+removePendingSubscription :: AgentClient -> SMPServer -> ConnId -> STM ()
+removePendingSubscription c = removeSubs_ (pendingSubscrSrvrs c)
+
+removeSubs_ :: TVar (Map SMPServer (Map ConnId RcvQueue)) -> SMPServer -> ConnId -> STM ()
+removeSubs_ ss server connId = modifyTVar ss $ M.update delSub server
   where
     delSub :: Map ConnId RcvQueue -> Maybe (Map ConnId RcvQueue)
     delSub cs =

tests/AgentTests.hs

@@ -62,9 +62,11 @@ agentTests (ATransport t) = do
       smpAgentTest3_1_1 $ testSubscription t
     it "should send notifications to client when server disconnects" $
       smpAgentServerTest $ testSubscrNotification t
-  describe "Message delivery" $ do
+  describe "Message delivery and server reconnection" $ do
     it "should deliver messages after losing server connection and re-connecting" $
       smpAgentTest2_2_2_needs_server $ testMsgDeliveryServerRestart t
+    it "should connect to the server when server goes up if it initially was down" $
+      smpAgentTestN [] $ testServerConnectionAfterError t
     it "should deliver pending messages after agent restarting" $
       smpAgentTest1_1_1 $ testMsgDeliveryAgentRestart t
     it "should concurrently deliver messages to connections without blocking" $
@@ -300,6 +302,43 @@ testMsgDeliveryServerRestart t alice bob = do
   where
     withServer test' = withSmpServerStoreLogOn (ATransport t) testPort2 (const test') `shouldReturn` ()
 
+testServerConnectionAfterError :: forall c. Transport c => TProxy c -> [c] -> IO ()
+testServerConnectionAfterError t _ = do
+  withAgent1 $ \bob -> do
+    withAgent2 $ \alice -> do
+      withServer $ do
+        connect (bob, "bob") (alice, "alice")
+
+      bob <# ("", "alice", DOWN)
+      alice <# ("", "bob", DOWN)
+      alice #: ("1", "bob", "SEND 5\nhello") #> ("1", "bob", MID 4)
+      alice #:# "nothing else delivered before the server is restarted"
+      bob #:# "nothing else delivered before the server is restarted"
+
+  withAgent1 $ \bob -> do
+    withAgent2 $ \alice -> do
+      bob #: ("1", "alice", "SUB") #> ("1", "alice", ERR (BROKER NETWORK))
+      alice #: ("1", "bob", "SUB") #> ("1", "bob", ERR (BROKER NETWORK))
+      withServer $ do
+        alice <#= \case ("", "bob", cmd) -> cmd == UP || cmd == SENT 4; _ -> False
+        alice <#= \case ("", "bob", cmd) -> cmd == UP || cmd == SENT 4; _ -> False
+        bob <# ("", "alice", UP)
+        bob <#= \case ("", "alice", Msg "hello") -> True; _ -> False
+        bob #: ("2", "alice", "ACK 4") #> ("2", "alice", OK)
+        alice #: ("1", "bob", "SEND 11\nhello again") #> ("1", "bob", MID 5)
+        alice <# ("", "bob", SENT 5)
+        bob <#= \case ("", "alice", Msg "hello again") -> True; _ -> False
+
+  removeFile testStoreLogFile
+  removeFile testDB
+  removeFile testDB2
+  where
+    withServer test' = withSmpServerStoreLogOn (ATransport t) testPort2 (const test') `shouldReturn` ()
+    withAgent1 = withAgent agentTestPort testDB
+    withAgent2 = withAgent agentTestPort2 testDB2
+    withAgent :: String -> String -> (c -> IO a) -> IO a
+    withAgent agentPort agentDB = withSmpAgentThreadOn_ (ATransport t) (agentPort, testPort2, agentDB) (pure ()) . const . testSMPAgentClientOn agentPort
+
 testMsgDeliveryAgentRestart :: Transport c => TProxy c -> c -> IO ()
 testMsgDeliveryAgentRestart t bob = do
   withAgent $ \alice -> do