From 1c2604f6a3a0f54a3e058862435be53d46201d00 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Fri, 22 Dec 2023 08:59:07 +0000
Subject: [PATCH] agent: background mode for agent NSE (#924)

---
 src/Simplex/Messaging/Agent.hs         | 21 ++++++++++++++-------
 src/Simplex/Messaging/Agent/Client.hs  | 16 +++++++++++++---
 tests/AgentTests/FunctionalAPITests.hs |  6 +++---
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index f1a56d110..c6d216c1e 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -174,14 +174,17 @@ import UnliftIO.STM
 -- import GHC.Conc (unsafeIOToSTM)
 
 -- | Creates an SMP agent client instance
-getSMPAgentClient :: (MonadRandom m, MonadUnliftIO m) => AgentConfig -> InitialAgentServers -> SQLiteStore -> m AgentClient
-getSMPAgentClient cfg initServers store =
+getSMPAgentClient :: (MonadRandom m, MonadUnliftIO m) => AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> m AgentClient
+getSMPAgentClient cfg initServers store backgroundMode =
   liftIO (newSMPAgentEnv cfg store) >>= runReaderT runAgent
   where
     runAgent = do
       c <- getAgentClient initServers
-      void $ raceAny_ [subscriber c, runNtfSupervisor c, cleanupManager c] `forkFinally` const (disconnectAgentClient c)
+      void $ runAgentThreads c `forkFinally` const (disconnectAgentClient c)
       pure c
+    runAgentThreads c
+      | backgroundMode = subscriber c
+      | otherwise = raceAny_ [subscriber c, runNtfSupervisor c, cleanupManager c]
 
 disconnectAgentClient :: MonadUnliftIO m => AgentClient -> m ()
 disconnectAgentClient c@AgentClient {agentEnv = Env {ntfSupervisor = ns, xftpAgent = xa}} = do
@@ -428,8 +431,9 @@ rcDiscoverCtrl' pairings = do
   liftError RCP $ discoverRCCtrl subs pairings
 
 -- | Activate operations
-foregroundAgent :: MonadUnliftIO m => AgentClient -> m ()
-foregroundAgent c = withAgentEnv c $ foregroundAgent' c
+-- stopPending = True is used in iOS NSE
+foregroundAgent :: MonadUnliftIO m => AgentClient -> Bool -> m ()
+foregroundAgent c stopPending = withAgentEnv c $ foregroundAgent' c stopPending
 
 -- | Suspend operations with max delay to deliver pending messages
 suspendAgent :: MonadUnliftIO m => AgentClient -> Int -> m ()
@@ -1800,8 +1804,11 @@ sendNtfConnCommands c cmd = do
 setNtfServers' :: AgentMonad' m => AgentClient -> [NtfServer] -> m ()
 setNtfServers' c = atomically . writeTVar (ntfServers c)
 
-foregroundAgent' :: AgentMonad' m => AgentClient -> m ()
-foregroundAgent' c = do
+-- stopPending = True is used in iOS NSE
+foregroundAgent' :: AgentMonad' m => AgentClient -> Bool -> m ()
+foregroundAgent' c stopPending = do
+  -- to prevent race conditions with message deliveries and async commands in the app
+  when stopPending $ stopPendingActions c
   atomically $ writeTVar (agentState c) ASForeground
   mapM_ activate $ reverse agentOperations
   where
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index d5cf2db17..77c3a9312 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -29,6 +29,7 @@ module Simplex.Messaging.Agent.Client
     closeAgentClient,
     closeProtocolServerClients,
     closeXFTPServerClient,
+    stopPendingActions,
     runSMPServerTest,
     runXFTPServerTest,
     getXFTPWorkPath,
@@ -604,16 +605,25 @@ closeAgentClient c = liftIO $ do
   closeProtocolServerClients c xftpClients
   cancelActions . actions $ reconnections c
   cancelActions . actions $ asyncClients c
-  cancelActions $ smpQueueMsgDeliveries c
-  cancelActions $ asyncCmdProcesses c
+  stopPendingActions c
   atomically . RQ.clear $ activeSubs c
   atomically . RQ.clear $ pendingSubs c
   clear subscrConns
+  clear getMsgLocks
+  where
+    clear :: Monoid m => (AgentClient -> TVar m) -> IO ()
+    clear sel = atomically $ writeTVar (sel c) mempty
+
+stopPendingActions :: MonadIO m => AgentClient -> m ()
+stopPendingActions c = liftIO $ do
+  -- messages
+  cancelActions $ smpQueueMsgDeliveries c
   clear pendingMsgsQueued
   clear smpQueueMsgQueues
+  -- commands
+  cancelActions $ asyncCmdProcesses c
   clear connCmdsQueued
   clear asyncCmdQueues
-  clear getMsgLocks
   where
     clear :: Monoid m => (AgentClient -> TVar m) -> IO ()
     clear sel = atomically $ writeTVar (sel c) mempty
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
index 3c3abb0f9..49b914fee 100644
--- a/tests/AgentTests/FunctionalAPITests.hs
+++ b/tests/AgentTests/FunctionalAPITests.hs
@@ -956,7 +956,7 @@ testRatchetSyncSuspendForeground t = do
 
   suspendAgent bob2 0
   threadDelay 100000
-  foregroundAgent bob2
+  foregroundAgent bob2 False
 
   withSmpServerStoreMsgLogOn t testPort $ \_ -> do
     runRight_ $ do
@@ -1119,7 +1119,7 @@ testSuspendingAgent =
     5 <- sendMessage a bId SMP.noMsgFlags "hello 2"
     get a ##> ("", bId, SENT 5)
     Nothing <- 100000 `timeout` get b
-    foregroundAgent b
+    foregroundAgent b False
     get b =##> \case ("", c, Msg "hello 2") -> c == aId; _ -> False
 
 testSuspendingAgentCompleteSending :: ATransport -> IO ()
@@ -2065,7 +2065,7 @@ testTwoUsers = withAgentClients2 $ \a b -> do
 getSMPAgentClient' :: AgentConfig -> InitialAgentServers -> FilePath -> IO AgentClient
 getSMPAgentClient' cfg' initServers dbPath = do
   Right st <- liftIO $ createAgentStore dbPath "" False MCError
-  getSMPAgentClient cfg' initServers st
+  getSMPAgentClient cfg' initServers st False
 
 testServerMultipleIdentities :: HasCallStack => IO ()
 testServerMultipleIdentities =