core: fix reconnection bugs (#979)

* core: fix reconnection bugs

* untangle newProtocolClient

* refactor

* report busy clientVar error

* log error

* comments

---------

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

src/Simplex/Messaging/Agent/Client.hs

@@ -155,6 +155,7 @@ import Data.Text.Encoding
 import Data.Time (UTCTime, defaultTimeLocale, formatTime, getCurrentTime)
 import Data.Time.Clock.System (getSystemTime)
 import Data.Word (Word16)
+
 -- import GHC.Conc (unsafeIOToSTM)
 import Network.Socket (HostName)
 import Simplex.FileTransfer.Client (XFTPChunkSpec (..), XFTPClient, XFTPClientConfig (..), XFTPClientError)
@@ -228,7 +229,7 @@ data SessionVar a = SessionVar
     sessionVarId :: Int
   }
 
-type ClientVar msg = SessionVar (Either AgentErrorType (Client msg)) 
+type ClientVar msg = SessionVar (Either AgentErrorType (Client msg))
 
 type SMPClientVar = ClientVar SMP.BrokerMsg
 
@@ -668,32 +669,44 @@ newProtocolClient ::
   (AgentClient -> TransportSession msg -> m ()) ->
   ClientVar msg ->
   m (Client msg)
-newProtocolClient c tSess@(userId, srv, entityId_) clients connectClient clientConnected v = tryConnectClient pure tryConnectAsync
+newProtocolClient c tSess@(userId, srv, entityId_) clients connectClient clientConnectedAsync v =
+  -- attempt sync connect first
+  tryAgentError (connectClient v) >>= \case
+    Right client -> putClient client $> client
+    Left e -> do
+      handleErr e $ newAsyncAction asyncConnectLoop (asyncClients c) -- initiate reconnect loop for temporary errors
+      throwError e -- signal error to caller
   where
-    tryConnectClient :: (Client msg -> m a) -> m () -> m a
-    tryConnectClient successAction retryAction =
-      tryAgentError (connectClient v) >>= \case
-        Right client -> do
-          logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv <> " (user " <> bshow userId <> maybe "" (" for entity " <>) entityId_ <> ")"
-          atomically $ putTMVar (sessionVar v) (Right client)
-          liftIO $ incClientStat c userId client "CLIENT" "OK"
-          atomically $ writeTBQueue (subQ c) ("", "", APC SAENone $ hostEvent CONNECT client)
-          successAction client
-        Left e -> do
-          liftIO $ incServerStat c userId srv "CLIENT" $ strEncode e
-          if temporaryAgentError e
-            then retryAction
-            else atomically $ do
-              putTMVar (sessionVar v) (Left e)
-              removeTSessVar v tSess clients
-          throwError e
-    tryConnectAsync :: m ()
-    tryConnectAsync = newAsyncAction connectAsync $ asyncClients c
-    connectAsync :: Int -> m ()
-    connectAsync aId = do
+    putClient :: Client msg -> m ()
+    putClient client = do
+      logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv <> " (user " <> bshow userId <> maybe "" (" for entity " <>) entityId_ <> ")"
+      -- tryPutTMVar is a precaution, it always succeeds here
+      r <- atomically $ tryPutTMVar (sessionVar v) (Right client)
+      unless r $ logError "newProtocolClient: cannot put connected client"
+      liftIO $ incClientStat c userId client "CLIENT" "OK"
+      atomically $ writeTBQueue (subQ c) ("", "", APC SAENone $ hostEvent CONNECT client)
+    handleErr :: AgentErrorType -> m () -> m ()
+    handleErr e handleTmp = do
+      liftIO $ incServerStat c userId srv "CLIENT" $ strEncode e
+      if temporaryAgentError e
+        then handleTmp
+        else do
+          r <- atomically $ do
+            removeTSessVar v tSess clients
+            -- tryPutTMVar is a precaution, it always succeeds here
+            tryPutTMVar (sessionVar v) (Left e)
+          unless r $ logError "newProtocolClient: cannot put client error"
+    asyncConnectLoop :: Int -> m ()
+    asyncConnectLoop aId = do
       ri <- asks $ reconnectInterval . config
-      withRetryInterval ri $ \_ loop -> void $ tryConnectClient (const $ clientConnected c tSess) loop
-      atomically . removeAsyncAction aId $ asyncClients c
+      withRetryInterval ri (const retryConnectClient)
+        `E.finally` atomically (removeAsyncAction aId $ asyncClients c)
+      where
+        -- does not return anything, restarts instead of throwing errors
+        retryConnectClient loop =
+          tryAgentError (connectClient v) >>= \case
+            Right client -> putClient client >> clientConnectedAsync c tSess
+            Left e -> handleErr e loop
 
 hostEvent :: forall err msg. (ProtocolTypeI (ProtoType msg), ProtocolServerClient err msg) => (AProtocolType -> TransportHost -> ACommand 'Agent 'AENone) -> Client msg -> ACommand 'Agent 'AENone
 hostEvent event = event (AProtocolType $ protocolTypeI @(ProtoType msg)) . clientTransportHost
@@ -724,7 +737,7 @@ closeAgentClient c = liftIO $ do
     clearWorkers workers = atomically $ swapTVar (workers c) mempty
     clear :: Monoid m => (AgentClient -> TVar m) -> IO ()
     clear sel = atomically $ writeTVar (sel c) mempty
-    cancelReconnect :: SessionVar (Async ())  -> IO ()
+    cancelReconnect :: SessionVar (Async ()) -> IO ()
     cancelReconnect v = void . forkIO $ atomically (readTMVar $ sessionVar v) >>= uninterruptibleCancel
 
 cancelWorker :: Worker -> IO ()

src/Simplex/Messaging/Client.hs

@@ -346,12 +346,12 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
       action <-
         async $
           runTransportClient tcConfig (Just username) useHost port' (Just $ keyHash srv) (client t c cVar)
-            `finally` atomically (putTMVar cVar $ Left PCENetworkError)
+            `finally` atomically (tryPutTMVar cVar $ Left PCENetworkError)
       c_ <- tcpConnectTimeout `timeout` atomically (takeTMVar cVar)
-      pure $ case c_ of
-        Just (Right c') -> Right c' {action = Just action}
-        Just (Left e) -> Left e
-        Nothing -> Left PCENetworkError
+      case c_ of
+        Just (Right c') -> pure $ Right c' {action = Just action}
+        Just (Left e) -> pure $ Left e
+        Nothing -> cancel action $> Left PCENetworkError
 
     useTransport :: (ServiceName, ATransport)
     useTransport = case port srv of

src/Simplex/Messaging/Transport/Client.hs

@@ -21,6 +21,7 @@ module Simplex.Messaging.Transport.Client
 where
 
 import Control.Applicative (optional)
+import Control.Logger.Simple (logError)
 import Control.Monad.IO.Unlift
 import Data.Aeson (FromJSON (..), ToJSON (..))
 import qualified Data.Attoparsec.ByteString.Char8 as A
@@ -48,7 +49,7 @@ import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers (parseAll, parseString)
 import Simplex.Messaging.Transport
 import Simplex.Messaging.Transport.KeepAlive
-import Simplex.Messaging.Util (bshow, (<$?>))
+import Simplex.Messaging.Util (bshow, (<$?>), catchAll, tshow)
 import System.IO.Error
 import Text.Read (readMaybe)
 import UnliftIO.Exception (IOException)
@@ -135,7 +136,7 @@ runTLSTransportClient tlsParams caStore_ cfg@TransportClientConfig {socksProxy,
         _ -> connectTCPClient hostName
   c <- liftIO $ do
     sock <- connectTCP port
-    mapM_ (setSocketKeepAlive sock) tcpKeepAlive
+    mapM_ (setSocketKeepAlive sock) tcpKeepAlive `catchAll` \e -> logError ("Error setting TCP keep-alive" <> tshow e)
     let tCfg = clientTransportConfig cfg
     connectTLS (Just hostName) tCfg clientParams sock >>= getClientConnection tCfg
   client c `E.finally` liftIO (closeConnection c)