try async commands without servers on different servers (#516)

* refactor

* retry commands with different servers

* refactor

* remove comment

Co-authored-by: JRoberts <8711996+jr-simplex@users.noreply.github.com>

Co-authored-by: JRoberts <8711996+jr-simplex@users.noreply.github.com>

src/Simplex/Messaging/Agent.hs

@@ -85,6 +85,7 @@ import Data.Bifunctor (bimap, first, second)
 import Data.ByteString.Char8 (ByteString)
 import Data.Composition ((.:), (.:.))
 import Data.Functor (($>))
+import Data.List (deleteFirstsBy)
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import Data.Map.Strict (Map)
@@ -359,8 +360,11 @@ ackMessageAsync' c connId msgId =
       enqueueCommand c connId (Just server) $ ACK msgId
 
 newConn :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> SConnectionMode c -> m (ConnId, ConnectionRequestUri c)
-newConn c connId asyncMode enableNtfs cMode = do
-  srv <- getSMPServer c
+newConn c connId asyncMode enableNtfs cMode =
+  getSMPServer c >>= newConnSrv c connId asyncMode enableNtfs cMode
+
+newConnSrv :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> SConnectionMode c -> SMPServer -> m (ConnId, ConnectionRequestUri c)
+newConnSrv c connId asyncMode enableNtfs cMode srv = do
   clientVRange <- asks $ smpClientVRange . config
   (rq, qUri) <- newRcvQueue c srv clientVRange
   connId' <- setUpConn asyncMode rq
@@ -387,7 +391,11 @@ newConn c connId asyncMode enableNtfs cMode = do
       withStore c $ \db -> createRcvConn db g cData rq cMode
 
 joinConn :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> ConnectionRequestUri c -> ConnInfo -> m ConnId
-joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentVRange (qUri :| _)) e2eRcvParamsUri) cInfo = do
+joinConn c connId asyncMode enableNtfs connReq cInfo =
+  getSMPServer c >>= joinConnSrv c connId asyncMode enableNtfs connReq cInfo
+
+joinConnSrv :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> ConnectionRequestUri c -> ConnInfo -> SMPServer -> m ConnId
+joinConnSrv c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentVRange (qUri :| _)) e2eRcvParamsUri) cInfo srv = do
   aVRange <- asks $ smpAgentVRange . config
   clientVRange <- asks $ smpClientVRange . config
   case ( qUri `compatibleVersion` clientVRange,
@@ -403,7 +411,7 @@ joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentV
           cData = ConnData {connId, connAgentVersion, enableNtfs, duplexHandshake = Just duplexHS}
       connId' <- setUpConn asyncMode cData sq rc
       let cData' = (cData :: ConnData) {connId = connId'}
-      tryError (confirmQueue aVersion c cData' sq cInfo $ Just e2eSndParams) >>= \case
+      tryError (confirmQueue aVersion c cData' sq srv cInfo $ Just e2eSndParams) >>= \case
         Right _ -> do
           unless duplexHS . void $ enqueueMessage c cData' sq SMP.noMsgFlags HELLO
           pure connId'
@@ -424,23 +432,22 @@ joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentV
             liftIO $ createRatchet db connId' rc
             pure connId'
     _ -> throwError $ AGENT A_VERSION
-joinConn c connId False enableNtfs (CRContactUri (ConnReqUriData _ agentVRange (qUri :| _))) cInfo = do
+joinConnSrv c connId False enableNtfs (CRContactUri (ConnReqUriData _ agentVRange (qUri :| _))) cInfo srv = do
   aVRange <- asks $ smpAgentVRange . config
   clientVRange <- asks $ smpClientVRange . config
   case ( qUri `compatibleVersion` clientVRange,
          agentVRange `compatibleVersion` aVRange
        ) of
     (Just qInfo, Just vrsn) -> do
-      (connId', cReq) <- newConn c connId False enableNtfs SCMInvitation
+      (connId', cReq) <- newConnSrv c connId False enableNtfs SCMInvitation srv
       sendInvitation c qInfo vrsn cReq cInfo
       pure connId'
     _ -> throwError $ AGENT A_VERSION
-joinConn _c _connId True _enableNtfs (CRContactUri _) _cInfo = do
+joinConnSrv _c _connId True _enableNtfs (CRContactUri _) _cInfo _srv = do
   throwError $ CMD PROHIBITED
 
-createReplyQueue :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> m SMPQueueInfo
-createReplyQueue c ConnData {connId, enableNtfs} SndQueue {smpClientVersion} = do
-  srv <- getSMPServer c
+createReplyQueue :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> SMPServer -> m SMPQueueInfo
+createReplyQueue c ConnData {connId, enableNtfs} SndQueue {smpClientVersion} srv = do
   (rq, qUri) <- newRcvQueue c srv $ versionToRange smpClientVersion
   let qInfo = toVersionT qUri smpClientVersion
   addSubscription c rq connId
@@ -490,21 +497,21 @@ processConfirmation c rq@RcvQueue {e2ePrivKey, smpClientVersion = v} SMPConfirma
 -- | Subscribe to receive connection messages (SUB command) in Reader monad
 subscribeConnection' :: forall m. AgentMonad m => AgentClient -> ConnId -> m ()
 subscribeConnection' c connId =
-  withStore c (`getConn` connId) >>= \case
-    SomeConn _ (DuplexConnection cData rq sq) -> do
-      resumeMsgDelivery c cData sq
-      subscribe rq
-      resumeConnCmds c connId
-    SomeConn _ (SndConnection cData sq) -> do
-      resumeMsgDelivery c cData sq
-      case status (sq :: SndQueue) of
-        Confirmed -> pure ()
-        Active -> throwError $ CONN SIMPLEX
-        _ -> throwError $ INTERNAL "unexpected queue status"
-      resumeConnCmds c connId
-    SomeConn _ (RcvConnection _ rq) -> subscribe rq >> resumeConnCmds c connId
-    SomeConn _ (ContactConnection _ rq) -> subscribe rq >> resumeConnCmds c connId
-    SomeConn _ (NewConnection _) -> resumeConnCmds c connId
+  withStore c (`getConn` connId) >>= \conn -> do
+    resumeConnCmds c connId
+    case conn of
+      SomeConn _ (DuplexConnection cData rq sq) -> do
+        resumeMsgDelivery c cData sq
+        subscribe rq
+      SomeConn _ (SndConnection cData sq) -> do
+        resumeMsgDelivery c cData sq
+        case status (sq :: SndQueue) of
+          Confirmed -> pure ()
+          Active -> throwError $ CONN SIMPLEX
+          _ -> throwError $ INTERNAL "unexpected queue status"
+      SomeConn _ (RcvConnection _ rq) -> subscribe rq
+      SomeConn _ (ContactConnection _ rq) -> subscribe rq
+      SomeConn _ (NewConnection _) -> pure ()
   where
     subscribe :: RcvQueue -> m ()
     subscribe rq = do
@@ -521,7 +528,7 @@ subscribeConnections' c connIds = do
       (subRs, rcvQs) = M.mapEither rcvQueueOrResult cs
       srvRcvQs :: Map SMPServer (Map ConnId (RcvQueue, ConnData)) = M.foldlWithKey' addRcvQueue M.empty rcvQs
   mapM_ (mapM_ (uncurry $ resumeMsgDelivery c) . sndQueue) cs
-  forM_ (M.keys cs) $ resumeConnCmds c
+  mapM_ (resumeConnCmds c) $ M.keys cs
   rcvRs <- mapConcurrently subscribe (M.assocs srvRcvQs)
   ns <- asks ntfSupervisor
   tkn <- readTVarIO (ntfTkn ns)
@@ -679,26 +686,35 @@ runCommandProcessing c@AgentClient {subQ} server = do
     E.try (withStore c $ \db -> getPendingCommand db cmdId) >>= \case
       Left (e :: E.SomeException) ->
         notify "" $ ERR (INTERNAL $ show e)
-      Right (connId, ACmd _ cmd) ->
+      Right (connId, ACmd _ cmd) -> do
+        usedSrvs <- newTVarIO ([] :: [SMPServer])
         withRetryInterval ri $ \loop -> do
           resp <- tryError $ case cmd of
-            NEW enableNtfs (ACM cMode) -> do
-              (_, cReq) <- newConn c connId True enableNtfs cMode
-              notify connId $ INV (ACR cMode cReq)
-            JOIN enableNtfs (ACR _ cReq) connInfo -> void $ joinConn c connId True enableNtfs cReq connInfo
+            NEW enableNtfs (ACM cMode) ->
+              withNextSrv usedSrvs $ \srv -> do
+                (_, cReq) <- newConnSrv c connId True enableNtfs cMode srv
+                notify connId $ INV (ACR cMode cReq)
+            JOIN enableNtfs (ACR _ cReq) connInfo ->
+              withNextSrv usedSrvs $ \srv ->
+                void $ joinConnSrv c connId True enableNtfs cReq connInfo srv
             LET confId ownCInfo -> allowConnection' c connId confId ownCInfo
             ACK msgId -> ackMessage' c connId msgId
-            _ -> notify "" $ ERR (INTERNAL "")
+            _ -> notify connId $ ERR $ INTERNAL $ "unsupported async command " <> show cmd
           case resp of
-            Left _ ->
-              -- TODO retry NEW and JOIN on different server
-              -- TODO depending on command, some errors shouldn't be retried
-              retryCommand loop
-            Right () -> do
-              delCmd cmdId
+            Left e
+              | temporaryAgentError e || e == BROKER HOST -> retryCommand loop
+              | otherwise -> notify connId $ ERR e
+            Right () -> withStore' c (`deleteCommand` cmdId)
   where
-    delCmd :: AsyncCmdId -> m ()
-    delCmd cmdId = withStore' c $ \db -> deleteCommand db cmdId
+    withNextSrv :: TVar [SMPServer] -> (SMPServer -> m ()) -> m ()
+    withNextSrv usedSrvs action = do
+      used <- readTVarIO usedSrvs
+      srv <- getNextSMPServer c used
+      atomically $ do
+        srvs <- readTVar $ smpServers c
+        let used' = if length used + 1 >= L.length srvs then [] else srv : used
+        writeTVar usedSrvs used'
+      action srv
     notify :: ConnId -> ACommand 'Agent -> m ()
     notify connId cmd = atomically $ writeTBQueue subQ ("", connId, cmd)
     retryCommand loop = do
@@ -844,7 +860,8 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} cData@ConnData {connId, duplexHandsh
                     -- and this branch should never be reached as receive is created before the confirmation,
                     -- so the condition is not necessary here, strictly speaking.
                     _ -> unless (duplexHandshake == Just True) $ do
-                      qInfo <- createReplyQueue c cData sq
+                      srv <- getSMPServer c
+                      qInfo <- createReplyQueue c cData sq srv
                       void . enqueueMessage c cData sq SMP.noMsgFlags $ REPLY [qInfo]
                 AM_A_MSG_ -> notify $ SENT mId
                 _ -> pure ()
@@ -1155,14 +1172,23 @@ suspendAgent' c@AgentClient {agentState = as} maxDelay = do
       suspendSendingAndDatabase c
 
 getSMPServer :: AgentMonad m => AgentClient -> m SMPServer
-getSMPServer c = do
-  smpServers <- readTVarIO $ smpServers c
-  case smpServers of
-    srv :| [] -> pure srv
-    servers -> do
-      gen <- asks randomServer
-      atomically . stateTVar gen $
-        first (servers L.!!) . randomR (0, L.length servers - 1)
+getSMPServer c = readTVarIO (smpServers c) >>= pickServer
+
+pickServer :: AgentMonad m => NonEmpty SMPServer -> m SMPServer
+pickServer = \case
+  srv :| [] -> pure srv
+  servers -> do
+    gen <- asks randomServer
+    atomically $ (servers L.!!) <$> stateTVar gen (randomR (0, L.length servers - 1))
+
+getNextSMPServer :: AgentMonad m => AgentClient -> [SMPServer] -> m SMPServer
+getNextSMPServer c usedSrvs = do
+  srvs <- readTVarIO $ smpServers c
+  case L.nonEmpty $ deleteFirstsBy different (L.toList srvs) usedSrvs of
+    Just srvs' -> pickServer srvs'
+    _ -> pickServer srvs
+  where
+    different (SMPServer host port _) (SMPServer host' port' _) = host /= host' || port /= port'
 
 subscriber :: (MonadUnliftIO m, MonadReader Env m) => AgentClient -> m ()
 subscriber c@AgentClient {msgQ} = forever $ do
@@ -1400,8 +1426,8 @@ connectReplyQueues c cData@ConnData {connId} ownConnInfo (qInfo :| _) = do
       withStore c $ \db -> upgradeRcvConnToDuplex db connId sq
       enqueueConfirmation c cData sq ownConnInfo Nothing
 
-confirmQueue :: forall m. AgentMonad m => Compatible Version -> AgentClient -> ConnData -> SndQueue -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m ()
-confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq connInfo e2eEncryption = do
+confirmQueue :: forall m. AgentMonad m => Compatible Version -> AgentClient -> ConnData -> SndQueue -> SMPServer -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m ()
+confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq srv connInfo e2eEncryption = do
   aMessage <- mkAgentMessage agentVersion
   msg <- mkConfirmation aMessage
   sendConfirmation c sq msg
@@ -1415,7 +1441,7 @@ confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq connInfo e2e
     mkAgentMessage :: Version -> m AgentMessage
     mkAgentMessage 1 = pure $ AgentConnInfo connInfo
     mkAgentMessage _ = do
-      qInfo <- createReplyQueue c cData sq
+      qInfo <- createReplyQueue c cData sq srv
       pure $ AgentConnInfoReply (qInfo :| []) connInfo
 
 enqueueConfirmation :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m ()