smp-server: fix several memory leaks

simplexmq.cabal

@@ -652,3 +652,47 @@ test-suite simplexmq-test
   if flag(client_postgres) || flag(server_postgres)
     build-depends:
         postgresql-simple ==0.7.*
+
+executable smp-mem-bench
+  if flag(client_library)
+      buildable: False
+  main-is: MemBench.hs
+  other-modules:
+      SMPClient
+      Util
+  hs-source-dirs:
+      bench
+      tests
+  default-extensions:
+      StrictData
+  ghc-options: -Wall -Wno-unused-imports -Wno-unused-top-binds -Wno-name-shadowing -threaded -rtsopts "-with-rtsopts=-T"
+  build-depends:
+      base
+    , async
+    , bytestring
+    , containers
+    , crypton
+    , crypton-x509
+    , crypton-x509-store
+    , crypton-x509-validation
+    , directory
+    , hspec ==2.11.*
+    , hspec-core ==2.11.*
+    , mtl
+    , network
+    , process
+    , simple-logger
+    , simplexmq
+    , sqlcipher-simple
+    , stm
+    , text
+    , time
+    , tls >=1.9.0 && <1.10
+    , transformers
+    , unliftio
+    , unliftio-core
+  if flag(server_postgres)
+    cpp-options: -DdbServerPostgres
+    build-depends:
+        postgresql-simple ==0.7.*
+  default-language: Haskell2010

src/Simplex/Messaging/Client.hs

@@ -35,6 +35,7 @@ module Simplex.Messaging.Client
     ProxiedRelay (..),
     getProtocolClient,
     closeProtocolClient,
+    pClientSentCommandsCount,
     protocolClientServer,
     protocolClientServer',
     transportHost',
@@ -167,6 +168,7 @@ import Simplex.Messaging.Protocol
 import Simplex.Messaging.Protocol.Types
 import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.SimplexName (SimplexDomain)
+import qualified Data.Map.Strict as M
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport
@@ -737,6 +739,10 @@ useWebPort cfg presetDomains ProtocolServer {host = h :| _} = case smpWebPortSer
   SWPPreset -> isPresetDomain presetDomains h
   SWPOff -> False
 
+-- | Count of in-flight (awaiting-response) commands on a client - for leak diagnostics.
+pClientSentCommandsCount :: ProtocolClient v err msg -> IO Int
+pClientSentCommandsCount ProtocolClient {client_ = PClient {sentCommands}} = M.size <$> readTVarIO sentCommands
+
 isPresetDomain :: [HostName] -> TransportHost -> Bool
 isPresetDomain presetDomains = \case
   THDomainName h -> any (`isSuffixOf` h) presetDomains

src/Simplex/Messaging/Client/Agent.hs

@@ -31,6 +31,8 @@ module Simplex.Messaging.Client.Agent
     removeActiveSubs,
     removePendingSub,
     removePendingSubs,
+    AgentLeakStats (..),
+    getAgentLeakStats,
   )
 where
 
@@ -158,6 +160,33 @@ data SMPClientAgent p = SMPClientAgent
 
 type OwnServer = Bool
 
+-- | Sizes of every per-server/per-session map in the client agent, plus total in-flight
+-- forwarded commands - for leak diagnostics on the proxy path.
+data AgentLeakStats = AgentLeakStats
+  { alSmpClients :: Int,
+    alSmpSessions :: Int,
+    alActiveServiceSubs :: Int,
+    alActiveQueueSubs :: Int,
+    alPendingServiceSubs :: Int,
+    alPendingQueueSubs :: Int,
+    alSmpSubWorkers :: Int,
+    alSentCommands :: Int
+  }
+
+getAgentLeakStats :: SMPClientAgent p -> IO AgentLeakStats
+getAgentLeakStats SMPClientAgent {smpClients, smpSessions, activeServiceSubs, activeQueueSubs, pendingServiceSubs, pendingQueueSubs, smpSubWorkers} = do
+  alSmpClients <- msize smpClients
+  sess <- readTVarIO smpSessions
+  alActiveServiceSubs <- msize activeServiceSubs
+  alActiveQueueSubs <- msize activeQueueSubs
+  alPendingServiceSubs <- msize pendingServiceSubs
+  alPendingQueueSubs <- msize pendingQueueSubs
+  alSmpSubWorkers <- msize smpSubWorkers
+  alSentCommands <- foldM (\ !a (_, c) -> (a +) <$> pClientSentCommandsCount c) 0 (M.elems sess)
+  pure AgentLeakStats {alSmpClients, alSmpSessions = M.size sess, alActiveServiceSubs, alActiveQueueSubs, alPendingServiceSubs, alPendingQueueSubs, alSmpSubWorkers, alSentCommands}
+  where
+    msize m = M.size <$> readTVarIO m
+
 newSMPClientAgent :: SParty p -> SMPClientAgentConfig -> Maybe DBService -> TVar ChaChaDRG -> IO (SMPClientAgent p)
 newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize} dbService randomDrg = do
   active <- newTVarIO True

src/Simplex/Messaging/Protocol.hs

@@ -1546,11 +1546,14 @@ queueIdHash = IdsHash . C.md5Hash . unEntityId
 {-# INLINE queueIdHash #-}
 
 addServiceSubs :: (Int64, IdsHash) -> (Int64, IdsHash) -> (Int64, IdsHash)
-addServiceSubs (n', idsHash') (n, idsHash) = (n + n', idsHash <> idsHash')
+addServiceSubs (n', idsHash') (n, idsHash) =
+  let !n'' = n + n'
+      !h = idsHash <> idsHash'
+   in (n'', h)
 
 subtractServiceSubs :: (Int64, IdsHash) -> (Int64, IdsHash) -> (Int64, IdsHash)
 subtractServiceSubs (n', idsHash') (n, idsHash)
-  | n > n' = (n - n', idsHash <> idsHash') -- concat is a reversible xor: (x `xor` y) `xor` y == x
+  | n > n' = let !n'' = n - n'; !h = idsHash <> idsHash' in (n'', h) -- concat is a reversible xor: (x `xor` y) `xor` y == x
   | otherwise = (0, mempty)
 
 data ProtocolErrorType = PECmdSyntax | PECmdUnknown | PESession | PEBlock

src/Simplex/Messaging/Server.hs

@@ -98,7 +98,7 @@ import qualified Network.TLS as TLS
 import Numeric.Natural (Natural)
 import Simplex.Messaging.Agent.Lock
 import Simplex.Messaging.Client (ProtocolClient (thParams), ProtocolClientError (..), SMPClient, SMPClientError, clientHandlers, forwardSMPTransmission, smpProxyError, temporaryClientError)
-import Simplex.Messaging.Client.Agent (OwnServer, SMPClientAgent (..), SMPClientAgentEvent (..), closeSMPClientAgent, getSMPServerClient'', isOwnServer, lookupSMPServerClient, getConnectedSMPServerClient)
+import Simplex.Messaging.Client.Agent (AgentLeakStats (..), OwnServer, SMPClientAgent (..), SMPClientAgentEvent (..), closeSMPClientAgent, getAgentLeakStats, getSMPServerClient'', isOwnServer, lookupSMPServerClient, getConnectedSMPServerClient)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
@@ -128,6 +128,7 @@ import Simplex.Messaging.Transport.Server
 import Simplex.Messaging.Util
 import Simplex.Messaging.Version
 import System.Environment (lookupEnv)
+import Text.Read (readMaybe)
 import System.Exit (exitFailure, exitSuccess)
 import System.IO (hPrint, hPutStrLn, hSetNewlineMode, universalNewlineMode)
 import System.Mem.Weak (deRefWeak)
@@ -174,6 +175,23 @@ data ClientSubAction
 
 type PrevClientSub s = (Client s, ClientSubAction, (EntityId, BrokerMsg))
 
+-- accumulator for per-client leak diagnostics (summed across all connected clients)
+data ClientAgg = ClientAgg
+  { aggEndThreads :: !Int, -- Weak ThreadId registrations in endThreads (forkClient leak)
+    aggEndThreadSeq :: !Int, -- total forkClient forks ever (fork rate)
+    aggProcThreads :: !Int,
+    aggSubs :: !Int, -- entries in per-client subscriptions map
+    aggNoSub :: !Int,
+    aggPending :: !Int, -- SubPending: forked delivery threads not yet completed (blocked-thread candidates)
+    aggThread :: !Int, -- SubThread: live delivery threads
+    aggProhibit :: !Int,
+    aggNtfSubs :: !Int,
+    aggSvcSubsCount :: !Int,
+    aggRcvQ :: !Int,
+    aggSndQ :: !Int, -- full sndQ => forkDeliver blocks (leak trigger)
+    aggMsgQ :: !Int
+  }
+
 smpServer :: forall s. MsgStoreClass s => TMVar Bool -> ServerConfig s -> Maybe AttachHTTP -> M s ()
 smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOptions} attachHTTP_ = do
   s <- asks server
@@ -191,6 +209,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
     ( serverThread "server subscribers" s subscribers subscriptions serviceSubsCount (Just cancelSub)
         : serverThread "server ntfSubscribers" s ntfSubscribers ntfSubscriptions ntfServiceSubsCount Nothing
         : deliverNtfsThread s
+        : leakDiagnosticsThread s
         : sendPendingEvtsThread s
         : receiveFromProxyAgent pa
         : expireNtfsThread cfg
@@ -496,6 +515,98 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
               printMessageStats "STORE: messages" msgStats
             Left e -> logError $ "STORE: expireOldMessages, error expiring messages, " <> tshow e
 
+    -- Periodic comprehensive leak diagnostics: sizes of every growable structure in the
+    -- server, summed across clients, plus the proxy agent. Interval seconds via SMP_LEAKDIAG_SEC
+    -- (default 60). A single greppable "LEAKDIAG ..." line per interval; whichever counter grows
+    -- monotonically over time is the leak.
+    leakDiagnosticsThread :: Server s -> M s ()
+    leakDiagnosticsThread srv = do
+      secStr <- liftIO $ lookupEnv "SMP_LEAKDIAG_SEC"
+      let sec = max 5 $ fromMaybe 60 (secStr >>= readMaybe)
+      ms <- asks msgStore
+      ns <- asks ntfStore
+      ProxyAgent {smpAgent} <- asks proxyAgent
+      labelMyThread "leakDiagnosticsThread"
+      -- never let a diagnostics error crash the server (this thread is in raceAny_)
+      liftIO $ forever $ do
+        threadDelay $ sec * 1000000
+        tryAny (logLeakStats srv ms ns smpAgent) >>= either (logError . ("LEAKDIAG error: " <>) . tshow) (const $ pure ())
+
+    logLeakStats :: Server s -> s -> NtfStore -> SMPClientAgent 'Sender -> IO ()
+    logLeakStats srv ms (NtfStore nsv) smpAgent = do
+#if MIN_VERSION_base(4,18,0)
+      nThreads <- length <$> listThreads
+#else
+      let nThreads = 0 :: Int
+#endif
+      cls <- IM.elems <$> getServerClients srv
+      cc <- foldM accClient emptyAgg cls
+      (smpQ, smpS, smpC, smpT, smpP) <- subAgg (subscribers srv)
+      (ntfQ, ntfS, ntfC, ntfT, ntfP) <- subAgg (ntfSubscribers srv)
+      ntfMap <- readTVarIO nsv
+      ntfMsgs <- foldM (\ !a v -> (a +) . length <$> readTVarIO v) (0 :: Int) (M.elems ntfMap)
+      EntityCounts {queueCount, notifierCount, rcvServiceCount, ntfServiceCount, rcvServiceQueuesCount, ntfServiceQueuesCount} <- getEntityCounts @(StoreQueue s) (queueStore ms)
+      LoadedQueueCounts {loadedQueueCount, loadedNotifierCount, openJournalCount, queueLockCount, notifierLockCount} <- loadedQueueCounts ms
+      AgentLeakStats {alSmpClients, alSmpSessions, alActiveServiceSubs, alActiveQueueSubs, alPendingServiceSubs, alPendingQueueSubs, alSmpSubWorkers, alSentCommands} <- getAgentLeakStats smpAgent
+      logNote $
+        T.concat
+          [ "LEAKDIAG",
+            f "threads" nThreads, f "clients" (length cls),
+            f "endThreads" (aggEndThreads cc), f "endThreadSeq" (aggEndThreadSeq cc), f "procThreads" (aggProcThreads cc),
+            f "subs" (aggSubs cc), f "subs_nosub" (aggNoSub cc), f "subs_pending" (aggPending cc), f "subs_thread" (aggThread cc), f "subs_prohibit" (aggProhibit cc),
+            f "ntfSubsClient" (aggNtfSubs cc), f "svcSubsCount" (aggSvcSubsCount cc),
+            f "rcvQ" (aggRcvQ cc), f "sndQ" (aggSndQ cc), f "msgQ" (aggMsgQ cc),
+            f "smp_qSubscribers" smpQ, f "smp_svcSubscribers" smpS, f "smp_subClients" smpC, f "smp_totalSvcSubs" smpT, f "smp_pendingEvents" smpP,
+            f "ntf_qSubscribers" ntfQ, f "ntf_svcSubscribers" ntfS, f "ntf_subClients" ntfC, f "ntf_totalSvcSubs" ntfT, f "ntf_pendingEvents" ntfP,
+            f "ntfStore_keys" (M.size ntfMap), f "ntfStore_msgs" ntfMsgs,
+            f "store_queues" queueCount, f "store_notifiers" notifierCount, f "store_rcvServices" rcvServiceCount, f "store_ntfServices" ntfServiceCount, f "store_rcvSvcQueues" rcvServiceQueuesCount, f "store_ntfSvcQueues" ntfServiceQueuesCount,
+            f "loaded_queues" loadedQueueCount, f "loaded_notifiers" loadedNotifierCount, f "open_journals" openJournalCount, f "queue_locks" queueLockCount, f "notifier_locks" notifierLockCount,
+            f "proxy_smpClients" alSmpClients, f "proxy_smpSessions" alSmpSessions, f "proxy_activeSvcSubs" alActiveServiceSubs, f "proxy_activeQSubs" alActiveQueueSubs, f "proxy_pendingSvcSubs" alPendingServiceSubs, f "proxy_pendingQSubs" alPendingQueueSubs, f "proxy_subWorkers" alSmpSubWorkers, f "proxy_sentCommands" alSentCommands
+          ]
+      where
+        f :: Show a => Text -> a -> Text
+        f k v = " " <> k <> "=" <> tshow v
+        emptyAgg = ClientAgg 0 0 0 0 0 0 0 0 0 0 0 0 0
+        accClient !agg Client {subscriptions, ntfSubscriptions, serviceSubsCount, procThreads, endThreads, endThreadSeq, rcvQ, sndQ, msgQ} = do
+          et <- IM.size <$> readTVarIO endThreads
+          es <- readTVarIO endThreadSeq
+          pt <- readTVarIO procThreads
+          subs <- readTVarIO subscriptions
+          (no, pe, th, pr) <- foldM accSub (0, 0, 0, 0) (M.elems subs)
+          nt <- M.size <$> readTVarIO ntfSubscriptions
+          sc <- fst <$> readTVarIO serviceSubsCount
+          (rl, sl, ml) <- atomically $ (,,) <$> lengthTBQueue rcvQ <*> lengthTBQueue sndQ <*> lengthTBQueue msgQ
+          pure
+            agg
+              { aggEndThreads = aggEndThreads agg + et,
+                aggEndThreadSeq = aggEndThreadSeq agg + es,
+                aggProcThreads = aggProcThreads agg + pt,
+                aggSubs = aggSubs agg + M.size subs,
+                aggNoSub = aggNoSub agg + no,
+                aggPending = aggPending agg + pe,
+                aggThread = aggThread agg + th,
+                aggProhibit = aggProhibit agg + pr,
+                aggNtfSubs = aggNtfSubs agg + nt,
+                aggSvcSubsCount = aggSvcSubsCount agg + fromIntegral sc,
+                aggRcvQ = aggRcvQ agg + fromIntegral rl,
+                aggSndQ = aggSndQ agg + fromIntegral sl,
+                aggMsgQ = aggMsgQ agg + fromIntegral ml
+              }
+        accSub (no, pe, th, pr) Sub {subThread} = case subThread of
+          ServerSub t ->
+            readTVarIO t >>= \case
+              NoSub -> pure (no + 1, pe, th, pr)
+              SubPending -> pure (no, pe + 1, th, pr)
+              SubThread _ -> pure (no, pe, th + 1, pr)
+          ProhibitSub -> pure (no, pe, th, pr + 1)
+        subAgg ServerSubscribers {queueSubscribers, serviceSubscribers, subClients, totalServiceSubs, pendingEvents} = do
+          q <- M.size <$> getSubscribedClients queueSubscribers
+          s' <- M.size <$> getSubscribedClients serviceSubscribers
+          sc <- IS.size <$> readTVarIO subClients
+          ts <- fst <$> readTVarIO totalServiceSubs
+          pe <- IM.size <$> readTVarIO pendingEvents
+          pure (q, s', sc, ts, pe)
+
     expireNtfsThread :: ServerConfig s -> M s ()
     expireNtfsThread ServerConfig {notificationExpiration = expCfg} = do
       ns <- asks ntfStore

src/Simplex/Messaging/Server/NtfStore.hs

@@ -33,13 +33,10 @@ data MsgNtf = MsgNtf
   }
 
 storeNtf :: NtfStore -> NotifierId -> MsgNtf -> IO ()
-storeNtf (NtfStore ns) nId ntf = do
-  TM.lookupIO nId ns >>= atomically . maybe newNtfs (`modifyTVar'` (ntf :))
+storeNtf (NtfStore ns) nId ntf =
   -- TODO [ntfdb] coalesce messages here once the client is updated to process multiple messages
   -- for single notification.
-  -- when (isJust prevNtf) $ incStat $ msgNtfReplaced stats
-  where
-    newNtfs = TM.lookup nId ns >>= maybe (TM.insertM nId (newTVar [ntf]) ns) (`modifyTVar'` (ntf :))
+  atomically $ TM.lookup nId ns >>= maybe (TM.insertM nId (newTVar [ntf]) ns) (`modifyTVar'` (ntf :))
 
 deleteNtfs :: NtfStore -> NotifierId -> IO Int
 deleteNtfs (NtfStore ns) nId = atomically (TM.lookupDelete nId ns) >>= maybe (pure 0) (fmap length . readTVarIO)
@@ -48,18 +45,17 @@ deleteExpiredNtfs :: NtfStore -> Int64 -> IO Int
 deleteExpiredNtfs (NtfStore ns) old =
   foldM (\expired -> fmap (expired +) . expireQueue) 0 . M.keys =<< readTVarIO ns
   where
-    expireQueue nId = TM.lookupIO nId ns >>= maybe (pure 0) expire
-    expire v = readTVarIO v >>= \case
-      [] -> pure 0
-      _ ->
-        atomically $ readTVar v >>= \case
-          [] -> pure 0
-          -- check the last message first, it is the earliest
-          ntfs | systemSeconds (ntfTs $ last $ ntfs) < old -> do
+    expireQueue nId = atomically $ TM.lookup nId ns >>= maybe (pure 0) (expire nId)
+    expire nId v = readTVar v >>= \case
+      [] -> TM.delete nId ns >> pure 0
+      -- check the last message first, it is the earliest
+      ntfs
+        | systemSeconds (ntfTs $ last ntfs) < old -> do
             let !ntfs' = filter (\MsgNtf {ntfTs = ts} -> systemSeconds ts >= old) ntfs
-            writeTVar v ntfs'
-            pure $! length ntfs - length ntfs'
-          _ -> pure 0
+            if null ntfs'
+              then TM.delete nId ns >> pure (length ntfs)
+              else writeTVar v ntfs' >> pure (length ntfs - length ntfs')
+        | otherwise -> pure 0
 
 data NtfLogRecord = NLRv1 NotifierId MsgNtf