From c380c795600b887fcae1614a52fb5cda691b569d Mon Sep 17 00:00:00 2001 From: Efim Poberezkin <8711996+efim-poberezkin@users.noreply.github.com> Date: Thu, 10 Feb 2022 18:48:11 +0400 Subject: [PATCH] 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 | 128 ++++++++++++++++++-------- tests/AgentTests.hs | 41 ++++++++- 2 files changed, 129 insertions(+), 40 deletions(-) diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs index e596602e0..d01bc30f9 100644 --- a/src/Simplex/Messaging/Agent/Client.hs +++ b/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 = diff --git a/tests/AgentTests.hs b/tests/AgentTests.hs index e214aafc4..422065ac3 100644 --- a/tests/AgentTests.hs +++ b/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