suspend/activate agent (#432)

* suspend/activate agent

* deliver pending messages before agent is suspended

* refactor

src/Simplex/Messaging/Agent.hs

@@ -60,7 +60,8 @@ module Simplex.Messaging.Agent
     deleteNtfToken,
     getNtfToken,
     deleteNtfSub,
-    setAgentPhase,
+    activateAgent,
+    suspendAgent,
     logConnection,
   )
 where
@@ -84,6 +85,7 @@ import Data.Time.Clock
 import Data.Time.Clock.System (systemToUTCTime)
 import Data.Word (Word16)
 import qualified Database.SQLite.Simple as DB
+-- import GHC.Conc (unsafeIOToSTM)
 import Simplex.Messaging.Agent.Client
 import Simplex.Messaging.Agent.Env.SQLite
 import Simplex.Messaging.Agent.NtfSubSupervisor
@@ -107,7 +109,7 @@ import Simplex.Messaging.Util (bshow, eitherToMaybe, liftE, liftError, tryError,
 import Simplex.Messaging.Version
 import System.Random (randomR)
 import UnliftIO.Async (async, race_)
-import UnliftIO.Concurrent (forkFinally)
+import UnliftIO.Concurrent (forkFinally, forkIO, threadDelay)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -214,8 +216,13 @@ getNtfToken c = withAgentEnv c $ getNtfToken' c
 deleteNtfSub :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
 deleteNtfSub c = withAgentEnv c . deleteNtfSub' c
 
-setAgentPhase :: AgentErrorMonad m => AgentClient -> AgentPhase -> m ()
-setAgentPhase c = withAgentEnv c . setAgentPhase' c
+-- | Activate operations
+activateAgent :: AgentErrorMonad m => AgentClient -> m ()
+activateAgent c = withAgentEnv c $ activateAgent' c
+
+-- | Suspend operations with max delay to deliver pending messages
+suspendAgent :: AgentErrorMonad m => AgentClient -> Int -> m ()
+suspendAgent c = withAgentEnv c . suspendAgent' c
 
 withAgentEnv :: AgentClient -> ReaderT Env m a -> m a
 withAgentEnv c = (`runReaderT` agentEnv c)
@@ -479,6 +486,9 @@ resumeMsgDelivery c cData@ConnData {connId} sq@SndQueue {server, sndId} = do
 
 queuePendingMsgs :: AgentMonad m => AgentClient -> ConnId -> SndQueue -> [InternalId] -> m ()
 queuePendingMsgs c connId sq msgIds = atomically $ do
+  modifyTVar' (msgDeliveryOp c) $ \s -> s {opsInProgress = opsInProgress s + length msgIds}
+  -- s <- readTVar (msgDeliveryOp c)
+  -- unsafeIOToSTM $ putStrLn $ "msgDeliveryOp: " <> show (opsInProgress s)
   q <- getPendingMsgQ c connId sq
   mapM_ (writeTQueue q) msgIds
 
@@ -497,9 +507,11 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} cData@ConnData {connId, duplexHandsh
   mq <- atomically $ getPendingMsgQ c connId sq
   ri <- asks $ reconnectInterval . config
   forever $ do
-    atomically $ endAgentOperation c AONetwork
+    atomically $ endAgentOperation c AOSndNetwork
     msgId <- atomically $ readTQueue mq
-    atomically $ beginAgentOperation c AONetwork
+    atomically $ do
+      beginAgentOperation c AOSndNetwork
+      endAgentOperation c AOMsgDelivery
     let mId = unId msgId
     E.try (withStore c $ \db -> getPendingMsgData db connId msgId) >>= \case
       Left (e :: E.SomeException) ->
@@ -583,9 +595,9 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} cData@ConnData {connId, duplexHandsh
     notifyDel msgId cmd = notify cmd >> delMsg msgId
     connError msgId = notifyDel msgId . ERR . CONN
     retrySending loop = do
-      atomically $ do
-        endAgentOperation c AONetwork
-        beginAgentOperation c AONetwork
+      -- end... is in a separate atomically because if begin... blocks, SUSPENDED won't be sent
+      atomically $ endAgentOperation c AOSndNetwork
+      atomically $ beginAgentOperation c AOSndNetwork
       loop
 
 ackMessage' :: forall m. AgentMonad m => AgentClient -> ConnId -> AgentMsgId -> m ()
@@ -779,12 +791,31 @@ initializeNtfSubQ c tkn = do
 setNtfServers' :: AgentMonad m => AgentClient -> [NtfServer] -> m ()
 setNtfServers' c = atomically . writeTVar (ntfServers c)
 
-setAgentPhase' :: AgentMonad m => AgentClient -> AgentPhase -> m ()
-setAgentPhase' c p = do
-  aPhase <- asks agentPhase
-  atomically $ do
-    writeTVar aPhase (p, False)
-    notifyAgentPhaseChanged c
+activateAgent' :: AgentMonad m => AgentClient -> m ()
+activateAgent' c = atomically $ do
+  writeTVar (agentState c) ASActive
+  activate databaseOp
+  activate sndNetworkOp
+  activate msgDeliveryOp
+  activate rcvNetworkOp
+  where
+    activate opSel = modifyTVar' (opSel c) $ \s -> s {opSuspended = False}
+
+suspendAgent' :: AgentMonad m => AgentClient -> Int -> m ()
+suspendAgent' c@AgentClient {agentState = as} maxDelay = do
+  state <-
+    atomically $ do
+      writeTVar as ASSuspending
+      suspendOperation c AORcvNetwork $
+        suspendOperation c AOMsgDelivery $
+          suspendSendingAndDatabase c
+      readTVar as
+  when (state == ASSuspending) . void . forkIO $ do
+    threadDelay maxDelay
+    -- liftIO $ putStrLn "suspendAgent after timeout"
+    atomically . whenSuspending c $ do
+      -- unsafeIOToSTM $ putStrLn $ "in timeout: suspendSendingAndDatabase"
+      suspendSendingAndDatabase c
 
 getSMPServer :: AgentMonad m => AgentClient -> m SMPServer
 getSMPServer c = do
@@ -798,9 +829,9 @@ getSMPServer c = do
 
 subscriber :: (MonadUnliftIO m, MonadReader Env m) => AgentClient -> m ()
 subscriber c@AgentClient {msgQ} = forever $ do
-  atomically $ endAgentOperation c AONetwork
+  atomically $ endAgentOperation c AORcvNetwork
   t <- atomically $ readTBQueue msgQ
-  atomically $ beginAgentOperation c AONetwork
+  atomically $ beginAgentOperation c AORcvNetwork
   withAgentLock c (runExceptT $ processSMPTransmission c t) >>= \case
     Left e -> liftIO $ print e
     Right _ -> return ()

src/Simplex/Messaging/Agent/Client.hs

@@ -49,9 +49,15 @@ module Simplex.Messaging.Agent.Client
     removeSubscription,
     hasActiveSubscription,
     agentDbPath,
+    AgentOperation (..),
+    AgentOpState (..),
+    AgentState (..),
     beginAgentOperation,
     endAgentOperation,
-    notifyAgentPhaseChanged,
+    suspendSendingAndDatabase,
+    suspendOperation,
+    notifySuspended,
+    whenSuspending,
     withStore,
     withStore',
   )
@@ -77,6 +83,7 @@ import Data.Text.Encoding
 import Data.Word (Word16)
 import Database.SQLite.Simple (SQLError)
 import qualified Database.SQLite.Simple as DB
+-- import GHC.Conc (unsafeIOToSTM)
 import Simplex.Messaging.Agent.Env.SQLite
 import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.RetryInterval
@@ -120,6 +127,11 @@ data AgentClient = AgentClient
     connMsgsQueued :: TMap ConnId Bool,
     smpQueueMsgQueues :: TMap (ConnId, SMPServer, SMP.SenderId) (TQueue InternalId),
     smpQueueMsgDeliveries :: TMap (ConnId, SMPServer, SMP.SenderId) (Async ()),
+    rcvNetworkOp :: TVar AgentOpState,
+    msgDeliveryOp :: TVar AgentOpState,
+    sndNetworkOp :: TVar AgentOpState,
+    databaseOp :: TVar AgentOpState,
+    agentState :: TVar AgentState,
     getMsgLocks :: TMap (SMPServer, SMP.RecipientId) (TMVar ()),
     reconnections :: TVar [Async ()],
     asyncClients :: TVar [Async ()],
@@ -128,6 +140,21 @@ data AgentClient = AgentClient
     lock :: TMVar ()
   }
 
+data AgentOperation = AORcvNetwork | AOMsgDelivery | AOSndNetwork | AODatabase
+  deriving (Eq, Show)
+
+agentOpSel :: AgentOperation -> (AgentClient -> TVar AgentOpState)
+agentOpSel = \case
+  AORcvNetwork -> rcvNetworkOp
+  AOMsgDelivery -> msgDeliveryOp
+  AOSndNetwork -> sndNetworkOp
+  AODatabase -> databaseOp
+
+data AgentOpState = AgentOpState {opSuspended :: Bool, opsInProgress :: Int}
+
+data AgentState = ASActive | ASSuspending | ASSuspended
+  deriving (Eq, Show)
+
 newAgentClient :: InitialAgentServers -> Env -> STM AgentClient
 newAgentClient InitialAgentServers {smp, ntf} agentEnv = do
   let qSize = tbqSize $ config agentEnv
@@ -145,12 +172,17 @@ newAgentClient InitialAgentServers {smp, ntf} agentEnv = do
   connMsgsQueued <- TM.empty
   smpQueueMsgQueues <- TM.empty
   smpQueueMsgDeliveries <- TM.empty
+  rcvNetworkOp <- newTVar $ AgentOpState False 0
+  msgDeliveryOp <- newTVar $ AgentOpState False 0
+  sndNetworkOp <- newTVar $ AgentOpState False 0
+  databaseOp <- newTVar $ AgentOpState False 0
+  agentState <- newTVar ASActive
   getMsgLocks <- TM.empty
   reconnections <- newTVar []
   asyncClients <- newTVar []
   clientId <- stateTVar (clientCounter agentEnv) $ \i -> (i + 1, i + 1)
   lock <- newTMVar ()
-  return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, subscrSrvrs, pendingSubscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, getMsgLocks, reconnections, asyncClients, clientId, agentEnv, lock}
+  return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, subscrSrvrs, pendingSubscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, rcvNetworkOp, msgDeliveryOp, sndNetworkOp, databaseOp, agentState, getMsgLocks, reconnections, asyncClients, clientId, agentEnv, lock}
 
 agentDbPath :: AgentClient -> FilePath
 agentDbPath AgentClient {agentEnv = Env {store = SQLiteStore {dbFilePath}}} = dbFilePath
@@ -643,25 +675,55 @@ cryptoError = \case
   e -> INTERNAL $ show e
 
 endAgentOperation :: AgentClient -> AgentOperation -> STM ()
-endAgentOperation c@AgentClient {agentEnv = Env {agentOperations}} op = do
-  TM.alter (Just . maybe 0 (\n -> max 0 $ n - 1)) op agentOperations
-  notifyAgentPhaseChanged c
+endAgentOperation c op = endOperation c op $ case op of
+  AORcvNetwork ->
+    suspendOperation c AOMsgDelivery $
+      suspendSendingAndDatabase c
+  AOMsgDelivery ->
+    suspendSendingAndDatabase c
+  AOSndNetwork ->
+    suspendOperation c AODatabase $
+      notifySuspended c
+  AODatabase ->
+    notifySuspended c
+
+suspendSendingAndDatabase :: AgentClient -> STM ()
+suspendSendingAndDatabase c =
+  suspendOperation c AOSndNetwork $
+    suspendOperation c AODatabase $
+      notifySuspended c
+
+suspendOperation :: AgentClient -> AgentOperation -> STM () -> STM ()
+suspendOperation c op endedAction = do
+  n <- stateTVar (agentOpSel op c) $ \s -> (opsInProgress s, s {opSuspended = True})
+  -- unsafeIOToSTM $ putStrLn $ "suspendOperation_ " <> show op <> " " <> show n
+  when (n == 0) $ whenSuspending c endedAction
+
+notifySuspended :: AgentClient -> STM ()
+notifySuspended c = do
+  -- unsafeIOToSTM $ putStrLn "notifySuspended"
+  writeTBQueue (subQ c) ("", "", SUSPENDED)
+  writeTVar (agentState c) ASSuspended
+
+endOperation :: AgentClient -> AgentOperation -> STM () -> STM ()
+endOperation c op endedAction = do
+  (suspended, n) <- stateTVar (agentOpSel op c) $ \s ->
+    let n = max 0 (opsInProgress s - 1)
+     in ((opSuspended s, n), s {opsInProgress = n})
+  -- unsafeIOToSTM $ putStrLn $ "endOperation: " <> show op <> " " <> show suspended <> " " <> show n
+  when (suspended && n == 0) $ whenSuspending c endedAction
+
+whenSuspending :: AgentClient -> STM () -> STM ()
+whenSuspending c = whenM ((== ASSuspending) <$> readTVar (agentState c))
 
 beginAgentOperation :: AgentClient -> AgentOperation -> STM ()
-beginAgentOperation AgentClient {agentEnv = Env {agentPhase, agentOperations}} op = do
-  (p, _) <- readTVar agentPhase
-  when (op `elem` disallowedOperations p) retry
-  TM.alter (Just . maybe 1 (+ 1)) op agentOperations
-
-notifyAgentPhaseChanged :: AgentClient -> STM ()
-notifyAgentPhaseChanged AgentClient {subQ, agentEnv = Env {agentPhase, agentOperations}} = do
-  (p, notified) <- readTVar agentPhase
-  unless notified $ do
-    ops <- readTVar agentOperations
-    let opsPaused = all (maybe True (== 0) . (`M.lookup` ops)) $ disallowedOperations p
-    when opsPaused $ do
-      writeTBQueue subQ ("", "", PHASE p)
-      writeTVar agentPhase (p, True)
+beginAgentOperation c op = do
+  let opVar = agentOpSel op c
+  s <- readTVar opVar
+  -- unsafeIOToSTM $ putStrLn $ "beginOperation? " <> show op <> " " <> show (opsInProgress s)
+  when (opSuspended s) retry
+  -- unsafeIOToSTM $ putStrLn $ "beginOperation! " <> show op <> " " <> show (opsInProgress s + 1)
+  writeTVar opVar $! s {opsInProgress = opsInProgress s + 1}
 
 withStore' :: AgentMonad m => AgentClient -> (DB.Connection -> IO a) -> m a
 withStore' c action = withStore c $ fmap Right . action

src/Simplex/Messaging/Agent/Env/SQLite.hs

@@ -17,8 +17,6 @@ module Simplex.Messaging.Agent.Env.SQLite
     defaultAgentConfig,
     defaultReconnectInterval,
     Env (..),
-    AgentOperation (..),
-    disallowedOperations,
     newSMPAgentEnv,
     NtfSupervisor (..),
     NtfSupervisorCommand (..),
@@ -121,30 +119,17 @@ data Env = Env
     idsDrg :: TVar ChaChaDRG,
     clientCounter :: TVar Int,
     randomServer :: TVar StdGen,
-    agentPhase :: TVar (AgentPhase, Bool),
-    agentOperations :: TMap AgentOperation Int,
     ntfSupervisor :: NtfSupervisor
   }
 
-data AgentOperation = AONetwork | AODatabase
-  deriving (Eq, Ord, Show)
-
-disallowedOperations :: AgentPhase -> [AgentOperation]
-disallowedOperations = \case
-  APActive -> []
-  APPaused -> [AONetwork]
-  APSuspended -> [AONetwork, AODatabase]
-
 newSMPAgentEnv :: (MonadUnliftIO m, MonadRandom m) => AgentConfig -> m Env
 newSMPAgentEnv config@AgentConfig {dbFile, yesToMigrations} = do
   idsDrg <- newTVarIO =<< drgNew
   store <- liftIO $ createSQLiteStore dbFile Migrations.app yesToMigrations
   clientCounter <- newTVarIO 0
   randomServer <- newTVarIO =<< liftIO newStdGen
-  agentPhase <- newTVarIO (APActive, True)
-  agentOperations <- atomically TM.empty
   ntfSupervisor <- atomically . newNtfSubSupervisor $ tbqSize config
-  return Env {config, store, idsDrg, clientCounter, randomServer, agentPhase, agentOperations, ntfSupervisor}
+  return Env {config, store, idsDrg, clientCounter, randomServer, ntfSupervisor}
 
 data NtfSupervisor = NtfSupervisor
   { ntfTkn :: TVar (Maybe NtfToken),

src/Simplex/Messaging/Agent/Protocol.hs

@@ -82,7 +82,6 @@ module Simplex.Messaging.Agent.Protocol
     QueueStatus (..),
     ACorrId,
     AgentMsgId,
-    AgentPhase (..),
     NotificationsMode (..),
     NotificationInfo (..),
 
@@ -231,39 +230,13 @@ data ACommand (p :: AParty) where
   DEL :: ACommand Client
   OK :: ACommand Agent
   ERR :: AgentErrorType -> ACommand Agent
-  PHASE :: AgentPhase -> ACommand Agent
+  SUSPENDED :: ACommand Agent
   NTFMODE :: NtfTknStatus -> NotificationsMode -> ACommand Agent
 
 deriving instance Eq (ACommand p)
 
 deriving instance Show (ACommand p)
 
--- | Agent phase allows to have two agent processes concurrently working with the same database
-data AgentPhase
-  = -- | agent is operating normally
-    APActive
-  | -- | agent is paused - no new send/receive operations will be started - they will STM-retry
-    APPaused
-  | -- | agent is suspended - no new send/receive/database operations will be started - they will STM-retry
-    APSuspended
-  deriving (Eq, Show)
-
-instance StrEncoding AgentPhase where
-  strEncode = \case
-    APActive -> "ACTIVE"
-    APPaused -> "PAUSED"
-    APSuspended -> "SUSPENDED"
-  strP =
-    A.takeTill (== ' ') >>= \case
-      "ACTIVE" -> pure APActive
-      "PAUSED" -> pure APPaused
-      "SUSPENDED" -> pure APSuspended
-      _ -> fail "bad AgentPhase"
-
-instance ToJSON AgentPhase where
-  toEncoding = strToJEncoding
-  toJSON = strToJSON
-
 data NotificationsMode = NMOff | NMPeriodic | NMInstant
   deriving (Eq, Show)
 
@@ -974,7 +947,7 @@ serializeCommand = \case
   CON -> "CON"
   ERR e -> "ERR " <> strEncode e
   OK -> "OK"
-  PHASE p -> "PHASE " <> strEncode p
+  SUSPENDED -> "SUSPENDED"
   NTFMODE t m -> "NTFMODE " <> smpEncode t <> " " <> strEncode m
   where
     showTs :: UTCTime -> ByteString

tests/AgentTests/FunctionalAPITests.hs

@@ -89,9 +89,13 @@ functionalAPITests t = do
       testInactiveClientDisconnected t
     it "should NOT disconnect active clients" $
       testActiveClientNotDisconnected t
-  describe "Agent phases" $ do
-    it "should update clients when agent phase changes" $
-      withSmpServer t testAgentPhaseChanges
+  describe "Suspending agent" $ do
+    it "should update client when agent is suspended" $
+      withSmpServer t testSuspendingAgent
+    it "should complete sending messages when agent is suspended" $
+      testSuspendingAgentCompleteSending t
+    it "should suspend agent on timeout, even if pending messages not sent" $
+      testSuspendingAgentTimeout t
 
 testAgentClient :: IO ()
 testAgentClient = do
@@ -417,8 +421,8 @@ testActiveClientNotDisconnected t = do
           get alice ##> ("", "", DOWN testSMPServer [connId])
     milliseconds ts = systemSeconds ts * 1000 + fromIntegral (systemNanoseconds ts `div` 1000000)
 
-testAgentPhaseChanges :: IO ()
-testAgentPhaseChanges = do
+testSuspendingAgent :: IO ()
+testSuspendingAgent = do
   a <- getSMPAgentClient agentCfg initAgentServers
   b <- getSMPAgentClient agentCfg {dbFile = testDB2} initAgentServers
   Right () <- runExceptT $ do
@@ -427,18 +431,72 @@ testAgentPhaseChanges = do
     get a ##> ("", bId, SENT 4)
     get b =##> \case ("", c, Msg "hello") -> c == aId; _ -> False
     ackMessage b aId 4
-    setAgentPhase b APPaused
-    get b ##> ("", "", PHASE APPaused)
+    suspendAgent b 1000000
+    get b ##> ("", "", SUSPENDED)
     5 <- sendMessage a bId SMP.noMsgFlags "hello 2"
     get a ##> ("", bId, SENT 5)
     Nothing <- 100000 `timeout` get b
-    setAgentPhase b APSuspended
-    get b ##> ("", "", PHASE APSuspended)
-    setAgentPhase b APActive
-    get b ##> ("", "", PHASE APActive)
+    activateAgent b
     get b =##> \case ("", c, Msg "hello 2") -> c == aId; _ -> False
   pure ()
 
+testSuspendingAgentCompleteSending :: ATransport -> IO ()
+testSuspendingAgentCompleteSending t = do
+  a <- getSMPAgentClient agentCfg initAgentServers
+  b <- getSMPAgentClient agentCfg {dbFile = testDB2} initAgentServers
+  Right (aId, bId) <- withSmpServerStoreLogOn t testPort $ \_ -> runExceptT $ do
+    (aId, bId) <- makeConnection a b
+    4 <- sendMessage a bId SMP.noMsgFlags "hello"
+    get a ##> ("", bId, SENT 4)
+    get b =##> \case ("", c, Msg "hello") -> c == aId; _ -> False
+    ackMessage b aId 4
+    pure (aId, bId)
+
+  Right () <- runExceptT $ do
+    ("", "", DOWN {}) <- get a
+    ("", "", DOWN {}) <- get b
+    5 <- sendMessage b aId SMP.noMsgFlags "hello too"
+    6 <- sendMessage b aId SMP.noMsgFlags "how are you?"
+    liftIO $ threadDelay 100000
+    suspendAgent b 5000000
+
+  Right () <- withSmpServerStoreLogOn t testPort $ \_ -> runExceptT $ do
+    get b =##> \case ("", c, SENT 5) -> c == aId; ("", "", UP {}) -> True; _ -> False
+    get b =##> \case ("", c, SENT 5) -> c == aId; ("", "", UP {}) -> True; _ -> False
+    get b =##> \case ("", c, SENT 6) -> c == aId; _ -> False
+    ("", "", SUSPENDED) <- get b
+
+    ("", "", UP {}) <- get a
+    get a =##> \case ("", c, Msg "hello too") -> c == bId; _ -> False
+    ackMessage a bId 5
+    get a =##> \case ("", c, Msg "how are you?") -> c == bId; _ -> False
+    ackMessage a bId 6
+
+  pure ()
+
+testSuspendingAgentTimeout :: ATransport -> IO ()
+testSuspendingAgentTimeout t = do
+  a <- getSMPAgentClient agentCfg initAgentServers
+  b <- getSMPAgentClient agentCfg {dbFile = testDB2} initAgentServers
+  Right (aId, _) <- withSmpServer t . runExceptT $ do
+    (aId, bId) <- makeConnection a b
+    4 <- sendMessage a bId SMP.noMsgFlags "hello"
+    get a ##> ("", bId, SENT 4)
+    get b =##> \case ("", c, Msg "hello") -> c == aId; _ -> False
+    ackMessage b aId 4
+    pure (aId, bId)
+
+  Right () <- runExceptT $ do
+    ("", "", DOWN {}) <- get a
+    ("", "", DOWN {}) <- get b
+    5 <- sendMessage b aId SMP.noMsgFlags "hello too"
+    6 <- sendMessage b aId SMP.noMsgFlags "how are you?"
+    suspendAgent b 100000
+    ("", "", SUSPENDED) <- get b
+    pure ()
+
+  pure ()
+
 exchangeGreetings :: AgentClient -> ConnId -> AgentClient -> ConnId -> ExceptT AgentErrorType IO ()
 exchangeGreetings alice bobId bob aliceId = do
   4 <- sendMessage alice bobId SMP.noMsgFlags "hello"