agent: avoid race condition between worker and supervisor when getting work (#928)

src/Simplex/Messaging/Agent/Client.hs

@@ -85,6 +85,8 @@ module Simplex.Messaging.Agent.Client
     AgentState (..),
     AgentLocks (..),
     AgentStatsKey (..),
+    waitForWork,
+    withWork,
     agentOperations,
     agentOperationBracket,
     waitUntilActive,
@@ -133,6 +135,7 @@ import Data.Functor (($>))
 import Data.List (deleteFirstsBy, foldl', partition, (\\))
 import Data.List.NonEmpty (NonEmpty (..), (<|))
 import qualified Data.List.NonEmpty as L
+import Data.Maybe (isNothing)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
 import Data.Maybe (isJust, listToMaybe)
@@ -1224,6 +1227,19 @@ cryptoError = \case
   where
     c = AGENT . A_CRYPTO
 
+waitForWork :: AgentMonad m => TMVar () -> m ()
+waitForWork = void . atomically . readTMVar
+
+withWork :: AgentMonad m => AgentClient -> TMVar () -> (DB.Connection -> IO (Maybe a)) -> (a -> m ()) -> m ()
+withWork c doWork getWork action = do
+  r <- withStore' c $ \db -> do
+    r' <- getWork db
+    when (isNothing r') noWorkToDo
+    pure r'
+  mapM_ action r
+  where
+    noWorkToDo = void . atomically $ tryTakeTMVar doWork
+
 endAgentOperation :: AgentClient -> AgentOperation -> STM ()
 endAgentOperation c op = endOperation c op $ case op of
   AONtfNetwork -> pure ()

src/Simplex/Messaging/Agent/NtfSubSupervisor.hs

@@ -159,24 +159,21 @@ runNtfWorker :: forall m. AgentMonad m => AgentClient -> NtfServer -> TMVar () -
 runNtfWorker c srv doWork = do
   delay <- asks $ ntfWorkerDelay . config
   forever $ do
-    void . atomically $ readTMVar doWork
+    waitForWork doWork
     agentOperationBracket c AONtfNetwork throwWhenInactive runNtfOperation
     threadDelay delay
   where
     runNtfOperation :: m ()
-    runNtfOperation = do
-      nextSub_ <- withStore' c (`getNextNtfSubNTFAction` srv)
-      logInfo $ "runNtfWorker, nextSub_ " <> tshow nextSub_
-      case nextSub_ of
-        Nothing -> noWorkToDo
-        Just a@(NtfSubscription {connId}, _, _) -> do
+    runNtfOperation =
+      withWork c doWork (`getNextNtfSubNTFAction` srv) $
+        \nextSub@(NtfSubscription {connId}, _, _) -> do
+          logInfo $ "runNtfWorker, nextSub " <> tshow nextSub
           ri <- asks $ reconnectInterval . config
           withRetryInterval ri $ \_ loop ->
-            processAction a
+            processSub nextSub
               `catchAgentError` retryOnError c "NtfWorker" loop (workerInternalError c connId . show)
-    noWorkToDo = void . atomically $ tryTakeTMVar doWork
-    processAction :: (NtfSubscription, NtfSubNTFAction, NtfActionTs) -> m ()
-    processAction (sub@NtfSubscription {connId, smpServer, ntfSubId}, action, actionTs) = do
+    processSub :: (NtfSubscription, NtfSubNTFAction, NtfActionTs) -> m ()
+    processSub (sub@NtfSubscription {connId, smpServer, ntfSubId}, action, actionTs) = do
       ts <- liftIO getCurrentTime
       unlessM (rescheduleAction doWork ts actionTs) $
         case action of
@@ -244,23 +241,20 @@ runNtfSMPWorker :: forall m. AgentMonad m => AgentClient -> SMPServer -> TMVar (
 runNtfSMPWorker c srv doWork = do
   delay <- asks $ ntfSMPWorkerDelay . config
   forever $ do
-    void . atomically $ readTMVar doWork
+    waitForWork doWork
     agentOperationBracket c AONtfNetwork throwWhenInactive runNtfSMPOperation
     threadDelay delay
   where
-    runNtfSMPOperation = do
-      nextSub_ <- withStore' c (`getNextNtfSubSMPAction` srv)
-      logInfo $ "runNtfSMPWorker, nextSub_ " <> tshow nextSub_
-      case nextSub_ of
-        Nothing -> noWorkToDo
-        Just a@(NtfSubscription {connId}, _, _) -> do
+    runNtfSMPOperation =
+      withWork c doWork (`getNextNtfSubSMPAction` srv) $
+        \nextSub@(NtfSubscription {connId}, _, _) -> do
+          logInfo $ "runNtfSMPWorker, nextSub " <> tshow nextSub
           ri <- asks $ reconnectInterval . config
           withRetryInterval ri $ \_ loop ->
-            processAction a
+            processSub nextSub
               `catchAgentError` retryOnError c "NtfSMPWorker" loop (workerInternalError c connId . show)
-    noWorkToDo = void . atomically $ tryTakeTMVar doWork
-    processAction :: (NtfSubscription, NtfSubSMPAction, NtfActionTs) -> m ()
-    processAction (sub@NtfSubscription {connId, ntfServer}, smpAction, actionTs) = do
+    processSub :: (NtfSubscription, NtfSubSMPAction, NtfActionTs) -> m ()
+    processSub (sub@NtfSubscription {connId, ntfServer}, smpAction, actionTs) = do
       ts <- liftIO getCurrentTime
       unlessM (rescheduleAction doWork ts actionTs) $
         case smpAction of