notify about critical error on agent crash (#1062)

* notify about critical error on agent crash

* waitUntilActive

* disposeAgent

* fix

src/Simplex/Messaging/Agent.hs

@@ -41,6 +41,7 @@ module Simplex.Messaging.Agent
     getSMPAgentClient,
     getSMPAgentClient_,
     disconnectAgentClient,
+    disposeAgentClient,
     resumeAgentClient,
     withConnLock,
     withInvLock,
@@ -122,7 +123,7 @@ import Control.Monad
 import Control.Monad.Except
 import Control.Monad.IO.Unlift (MonadUnliftIO)
 import Control.Monad.Reader
-import Crypto.Random (ChaChaDRG, MonadRandom)
+import Crypto.Random (ChaChaDRG)
 import qualified Data.Aeson as J
 import Data.Bifunctor (bimap, first, second)
 import Data.ByteString.Char8 (ByteString)
@@ -179,28 +180,40 @@ import Simplex.Messaging.Version
 import Simplex.RemoteControl.Client
 import Simplex.RemoteControl.Invitation
 import Simplex.RemoteControl.Types
+import System.Mem.Weak (deRefWeak)
 import UnliftIO.Async (race_)
-import UnliftIO.Concurrent (forkFinally, forkIO, threadDelay)
+import UnliftIO.Concurrent (forkFinally, forkIO, killThread, mkWeakThreadId, threadDelay)
+import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
 -- import GHC.Conc (unsafeIOToSTM)
 
 -- | Creates an SMP agent client instance
-getSMPAgentClient :: (MonadRandom m, MonadUnliftIO m) => AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> m AgentClient
+getSMPAgentClient :: MonadIO m => AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> m AgentClient
 getSMPAgentClient = getSMPAgentClient_ 1
 {-# INLINE getSMPAgentClient #-}
 
-getSMPAgentClient_ :: (MonadRandom m, MonadUnliftIO m) => Int -> AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> m AgentClient
+getSMPAgentClient_ :: MonadIO m => Int -> AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> m AgentClient
 getSMPAgentClient_ clientId cfg initServers store backgroundMode =
-  liftIO (newSMPAgentEnv cfg store) >>= runReaderT runAgent
+  liftIO $ newSMPAgentEnv cfg store >>= runReaderT runAgent
   where
     runAgent = do
-      c <- getAgentClient clientId initServers
-      void $ runAgentThreads c `forkFinally` const (disconnectAgentClient c)
+      c@AgentClient {acThread}  <- getAgentClient clientId initServers
+      t <- runAgentThreads c `forkFinally` const (disconnectAgentClient c)
+      atomically . writeTVar acThread . Just =<< mkWeakThreadId t
       pure c
     runAgentThreads c
-      | backgroundMode = subscriber c
-      | otherwise = raceAny_ [subscriber c, runNtfSupervisor c, cleanupManager c]
+      | backgroundMode = run c "subscriber" $ subscriber c
+      | otherwise =
+          raceAny_
+            [ run c "subscriber" $ subscriber c,
+              run c "runNtfSupervisor" $ runNtfSupervisor c,
+              run c "cleanupManager" $ cleanupManager c
+            ]
+    run AgentClient {subQ, acThread} name a =
+      a `E.catchAny` \e -> whenM (isJust <$> readTVarIO acThread) $ do
+        logError $ "Agent thread " <> name <> " crashed: " <> tshow e
+        atomically $ writeTBQueue subQ ("", "", APC SAEConn $ ERR $ CRITICAL True $ show e)
 
 disconnectAgentClient :: MonadUnliftIO m => AgentClient -> m ()
 disconnectAgentClient c@AgentClient {agentEnv = Env {ntfSupervisor = ns, xftpAgent = xa}} = do
@@ -209,6 +222,14 @@ disconnectAgentClient c@AgentClient {agentEnv = Env {ntfSupervisor = ns, xftpAge
   closeXFTPAgent xa
   logConnection c False
 
+-- only used in the tests
+disposeAgentClient :: MonadUnliftIO m => AgentClient -> m ()
+disposeAgentClient c@AgentClient {acThread, agentEnv = Env {store}} = do
+  t_ <- atomically (swapTVar acThread Nothing) $>>= (liftIO . deRefWeak)
+  disconnectAgentClient c
+  mapM_ killThread t_
+  liftIO $ closeSQLiteStore store
+
 resumeAgentClient :: MonadIO m => AgentClient -> m ()
 resumeAgentClient c = atomically $ writeTVar (active c) True
 
@@ -1916,9 +1937,11 @@ cleanupManager c@AgentClient {subQ} = do
   where
     run :: forall e. AEntityI e => (AgentErrorType -> ACommand 'Agent e) -> ExceptT AgentErrorType m () -> m ()
     run err a = do
-      void . runExceptT $ a `catchAgentError` (notify "" . err)
+      waitActive . runExceptT $ a `catchAgentError` (notify "" . err)
       step <- asks $ cleanupStepInterval . config
       liftIO $ threadDelay step
+    -- we are catching it to avoid CRITICAL errors in tests when this is the only remaining handle to active
+    waitActive a = liftIO (E.tryAny . atomically $ waitUntilActive c) >>= either (\_ -> pure ()) (\_ -> void a)
     deleteConns =
       withLock (deleteLock c) "cleanupManager" $ do
         void $ withStore' c getDeletedConnIds >>= deleteDeletedConns c

src/Simplex/Messaging/Agent/Client.hs

@@ -129,7 +129,7 @@ module Simplex.Messaging.Agent.Client
 where
 
 import Control.Applicative ((<|>))
-import Control.Concurrent (forkIO, threadDelay)
+import Control.Concurrent (ThreadId, forkIO, threadDelay)
 import Control.Concurrent.Async (Async, uninterruptibleCancel)
 import Control.Concurrent.STM (retry, throwSTM)
 import Control.Exception (AsyncException (..))
@@ -225,6 +225,7 @@ import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport.Client (TransportHost)
 import Simplex.Messaging.Util
 import Simplex.Messaging.Version
+import System.Mem.Weak (Weak)
 import System.Random (randomR)
 import UnliftIO (mapConcurrently, timeout)
 import UnliftIO.Async (async)
@@ -252,7 +253,8 @@ type NtfTransportSession = TransportSession NtfResponse
 type XFTPTransportSession = TransportSession FileResponse
 
 data AgentClient = AgentClient
-  { active :: TVar Bool,
+  { acThread :: TVar (Maybe (Weak ThreadId)),
+    active :: TVar Bool,
     rcvQ :: TBQueue (ATransmission 'Client),
     subQ :: TBQueue (ATransmission 'Agent),
     msgQ :: TBQueue (ServerTransmission SMPVersion BrokerMsg),
@@ -395,6 +397,7 @@ data AgentStatsKey = AgentStatsKey
 newAgentClient :: Int -> InitialAgentServers -> Env -> STM AgentClient
 newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg} agentEnv = do
   let qSize = tbqSize $ config agentEnv
+  acThread <- newTVar Nothing
   active <- newTVar True
   rcvQ <- newTBQueue qSize
   subQ <- newTBQueue qSize
@@ -428,7 +431,8 @@ newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg} agentEnv =
   agentStats <- TM.empty
   return
     AgentClient
-      { active,
+      { acThread,
+        active,
         rcvQ,
         subQ,
         msgQ,