smp server: optimize concurrency and memory usage, refactor (#1544)

* smp server: optimize concurrency and memory usage, refactor

* hide clients IntMap

* reduce STM contention

* comment

* version

* correct stats for subscriptions

* version

* comment

* remove subscribed clients from map

* version

* optimze, refactor

* version

* debug test

* enable all tests

* remove test logs

* retry failed tests with debug logging

* increase test timeout

* sync between tests

Author: epoberezkin

simplexmq.cabal

@@ -519,6 +519,7 @@ test-suite simplexmq-test
     , generic-random ==1.5.*
     , hashable
     , hspec ==2.11.*
+    , hspec-core ==2.11.*
     , http-client
     , http-types
     , http2

src/Simplex/Messaging/Notifications/Server/Env.hs

@@ -53,7 +53,6 @@ data NtfServerConfig = NtfServerConfig
     subIdBytes :: Int,
     regCodeBytes :: Int,
     clientQSize :: Natural,
-    subQSize :: Natural,
     pushQSize :: Natural,
     smpAgentCfg :: SMPClientAgentConfig,
     apnsConfig :: APNSPushClientConfig,
@@ -94,11 +93,11 @@ data NtfEnv = NtfEnv
   }
 
 newNtfServerEnv :: NtfServerConfig -> IO NtfEnv
-newNtfServerEnv config@NtfServerConfig {subQSize, pushQSize, smpAgentCfg, apnsConfig, dbStoreConfig, ntfCredentials, startOptions} = do
+newNtfServerEnv config@NtfServerConfig {pushQSize, smpAgentCfg, apnsConfig, dbStoreConfig, ntfCredentials, startOptions} = do
   when (compactLog startOptions) $ compactDbStoreLog $ dbStoreLogPath dbStoreConfig
   random <- C.newRandom
   store <- newNtfDbStore dbStoreConfig
-  subscriber <- newNtfSubscriber subQSize smpAgentCfg random
+  subscriber <- newNtfSubscriber smpAgentCfg random
   pushServer <- newNtfPushServer pushQSize apnsConfig
   tlsServerCreds <- loadServerCredential ntfCredentials
   Fingerprint fp <- loadFingerprint ntfCredentials
@@ -121,8 +120,8 @@ data NtfSubscriber = NtfSubscriber
 
 type SMPSubscriberVar = SessionVar SMPSubscriber
 
-newNtfSubscriber :: Natural -> SMPClientAgentConfig -> TVar ChaChaDRG -> IO NtfSubscriber
-newNtfSubscriber qSize smpAgentCfg random = do
+newNtfSubscriber :: SMPClientAgentConfig -> TVar ChaChaDRG -> IO NtfSubscriber
+newNtfSubscriber smpAgentCfg random = do
   smpSubscribers <- TM.emptyIO
   subscriberSeq <- newTVarIO 0
   smpAgent <- newSMPClientAgent smpAgentCfg random

src/Simplex/Messaging/Notifications/Server/Main.hs

@@ -233,7 +233,6 @@ ntfServerCLI cfgPath logPath =
               subIdBytes = 24,
               regCodeBytes = 32,
               clientQSize = 64,
-              subQSize = 2048,
               pushQSize = 32768,
               smpAgentCfg =
                 defaultSMPClientAgentConfig

src/Simplex/Messaging/Server.hs

@@ -18,7 +18,59 @@
 #endif
 {-# OPTIONS_GHC -fno-warn-ambiguous-fields #-}
 
-module Simplex.Messaging.Server.Env.STM where
+module Simplex.Messaging.Server.Env.STM
+  ( ServerConfig (..),
+    ServerStoreCfg (..),
+    AServerStoreCfg (..),
+    StorePaths (..),
+    StartOptions (..),
+    Env (..),
+    Server (..),
+    ServerSubscribers (..),
+    SubscribedClients,
+    ProxyAgent (..),
+    Client (..),
+    AClient (..),
+    ClientId,
+    Subscribed,
+    Sub (..),
+    ServerSub (..),
+    SubscriptionThread (..),
+    MsgStore,
+    AMsgStore (..),
+    AStoreType (..),
+    newEnv,
+    mkJournalStoreConfig,
+    newClient,
+    getServerClients,
+    getServerClient,
+    insertServerClient,
+    deleteServerClient,
+    getSubscribedClients,
+    getSubscribedClient,
+    upsertSubscribedClient,
+    lookupDeleteSubscribedClient,
+    deleteSubcribedClient,
+    sameClientId,
+    sameClient,
+    clientId',
+    newSubscription,
+    newProhibitedSub,
+    defaultMsgQueueQuota,
+    defMsgExpirationDays,
+    defNtfExpirationHours,
+    defaultMessageExpiration,
+    defaultNtfExpiration,
+    defaultInactiveClientExpiration,
+    defaultProxyClientConcurrency,
+    defaultMaxJournalMsgCount,
+    defaultMaxJournalStateLines,
+    defaultIdleQueueInterval,
+    journalMsgStoreDepth,
+    readWriteQueueStore,
+    noPostgresExit,
+  )
+where
 
 import Control.Concurrent (ThreadId)
 import Control.Logger.Simple
@@ -29,9 +81,12 @@ import Data.ByteString.Char8 (ByteString)
 import Data.Int (Int64)
 import Data.IntMap.Strict (IntMap)
 import qualified Data.IntMap.Strict as IM
+import Data.IntSet (IntSet)
+import qualified Data.IntSet as IS
 import Data.Kind (Constraint)
 import Data.List (intercalate)
 import Data.List.NonEmpty (NonEmpty)
+import Data.Map.Strict (Map)
 import Data.Maybe (isJust)
 import qualified Data.Text as T
 import Data.Time.Clock (getCurrentTime, nominalDay)
@@ -66,6 +121,7 @@ import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (ATransport, VersionRangeSMP, VersionSMP)
 import Simplex.Messaging.Transport.Server
+import Simplex.Messaging.Util (ifM, whenM, ($>>=))
 import System.Directory (doesFileExist)
 import System.Exit (exitFailure)
 import System.IO (IOMode (..))
@@ -203,7 +259,6 @@ data Env = Env
     serverStats :: ServerStats,
     sockets :: TVar [(ServiceName, SocketState)],
     clientSeq :: TVar ClientId,
-    clients :: TVar (IntMap (Maybe AClient)),
     proxyAgent :: ProxyAgent -- senders served on this proxy
   }
 
@@ -236,17 +291,72 @@ data AMsgStore =
 type Subscribed = Bool
 
 data Server = Server
-  { subscribedQ :: TQueue (RecipientId, ClientId, Subscribed),
-    subscribers :: TMap RecipientId (TVar AClient),
-    ntfSubscribedQ :: TQueue (NotifierId, ClientId, Subscribed),
-    notifiers :: TMap NotifierId (TVar AClient),
-    subClients :: TVar (IntMap AClient), -- clients with SMP subscriptions
-    ntfSubClients :: TVar (IntMap AClient), -- clients with Ntf subscriptions
-    pendingSubEvents :: TVar (IntMap (NonEmpty (RecipientId, Subscribed))),
-    pendingNtfSubEvents :: TVar (IntMap (NonEmpty (NotifierId, Subscribed))),
+  { clients :: ServerClients,
+    subscribers :: ServerSubscribers,
+    ntfSubscribers :: ServerSubscribers,
     savingLock :: Lock
   }
 
+-- not exported, to prevent concurrent IntMap lookups inside STM transactions.
+newtype ServerClients = ServerClients {serverClients :: TVar (IntMap AClient)}
+
+data ServerSubscribers = ServerSubscribers
+  { subQ :: TQueue (QueueId, ClientId, Subscribed),
+    queueSubscribers :: SubscribedClients,
+    subClients :: TVar IntSet,
+    pendingEvents :: TVar (IntMap (NonEmpty (EntityId, BrokerMsg)))
+  }
+
+-- not exported, to prevent accidental concurrent Map lookups inside STM transactions.
+-- Map stores TVars with pointers to the clients rather than client ID to allow reading the same TVar
+-- inside transactions to ensure that transaction is re-evaluated in case subscriber changes.
+-- Storing Maybe allows to have continuity of subscription when the same user client disconnects and re-connects -
+-- any STM transaction that reads subscribed client will re-evaluate in this case.
+-- The subscriptions that were made at any point are not removed -
+-- this is a better trade-off with intermittently connected mobile clients.
+data SubscribedClients = SubscribedClients (TMap EntityId (TVar (Maybe AClient)))
+
+getSubscribedClients :: SubscribedClients -> IO (Map EntityId (TVar (Maybe AClient)))
+getSubscribedClients (SubscribedClients cs) = readTVarIO cs
+
+getSubscribedClient :: EntityId -> SubscribedClients -> IO (Maybe (TVar (Maybe AClient)))
+getSubscribedClient entId (SubscribedClients cs) = TM.lookupIO entId cs
+{-# INLINE getSubscribedClient #-}
+
+-- insert subscribed and current client, return previously subscribed client if it is different
+upsertSubscribedClient :: EntityId -> AClient -> SubscribedClients -> STM (Maybe AClient)
+upsertSubscribedClient entId ac@(AClient _ _ c) (SubscribedClients cs) =
+  TM.lookup entId cs >>= \case
+    Nothing -> Nothing <$ TM.insertM entId (newTVar (Just ac)) cs
+    Just cv ->
+      readTVar cv >>= \case
+        Just c' | sameClientId c c' -> pure Nothing
+        c_ -> c_ <$ writeTVar cv (Just ac)
+
+-- lookup and delete currently subscribed client
+lookupDeleteSubscribedClient :: EntityId -> SubscribedClients -> STM (Maybe AClient)
+lookupDeleteSubscribedClient entId (SubscribedClients cs) =
+  TM.lookupDelete entId cs $>>= (`swapTVar` Nothing)
+
+deleteSubcribedClient :: EntityId -> Client s -> SubscribedClients -> IO ()
+deleteSubcribedClient entId c (SubscribedClients cs) =
+  -- lookup of the subscribed client TVar can be in separate transaction,
+  -- as long as the client is read in the same transaction -
+  -- it prevents removing the next subscribed client and also avoids STM contention for the Map.
+  TM.lookupIO entId cs >>= mapM_ (\cv -> atomically $ whenM (sameClient c cv) $ delete cv)
+  where
+    delete cv = do
+      writeTVar cv Nothing
+      TM.delete entId cs
+
+sameClientId :: Client s -> AClient -> Bool
+sameClientId Client {clientId} ac = clientId == clientId' ac
+{-# INLINE sameClientId #-}
+
+sameClient :: Client s -> TVar (Maybe AClient) -> STM Bool
+sameClient c cv = maybe False (sameClientId c) <$> readTVar cv
+{-# INLINE sameClient #-}
+
 newtype ProxyAgent = ProxyAgent
   { smpAgent :: SMPClientAgent
   }
@@ -288,16 +398,40 @@ data Sub = Sub
 
 newServer :: IO Server
 newServer = do
-  subscribedQ <- newTQueueIO
-  subscribers <- TM.emptyIO
-  ntfSubscribedQ <- newTQueueIO
-  notifiers <- TM.emptyIO
-  subClients <- newTVarIO IM.empty
-  ntfSubClients <- newTVarIO IM.empty
-  pendingSubEvents <- newTVarIO IM.empty
-  pendingNtfSubEvents <- newTVarIO IM.empty
+  clients <- ServerClients <$> newTVarIO mempty
+  subscribers <- newServerSubscribers
+  ntfSubscribers <- newServerSubscribers
   savingLock <- createLockIO
-  return Server {subscribedQ, subscribers, ntfSubscribedQ, notifiers, subClients, ntfSubClients, pendingSubEvents, pendingNtfSubEvents, savingLock}
+  return Server {clients, subscribers, ntfSubscribers, savingLock}
+
+getServerClients :: Server -> IO (IntMap AClient)
+getServerClients = readTVarIO . serverClients . clients
+{-# INLINE getServerClients #-}
+
+getServerClient :: ClientId -> Server -> IO (Maybe AClient)
+getServerClient cId s = IM.lookup cId <$> getServerClients s
+{-# INLINE getServerClient #-}
+
+insertServerClient :: AClient -> Server -> IO Bool
+insertServerClient ac@(AClient _ _ Client {clientId, connected}) Server {clients} =
+  atomically $
+    ifM
+      (readTVar connected)
+      (True <$ modifyTVar' (serverClients clients) (IM.insert clientId ac))
+      (pure False)
+{-# INLINE insertServerClient #-}
+
+deleteServerClient :: ClientId -> Server -> IO ()
+deleteServerClient cId Server {clients} = atomically $ modifyTVar' (serverClients clients) $ IM.delete cId
+{-# INLINE deleteServerClient #-}
+
+newServerSubscribers :: IO ServerSubscribers
+newServerSubscribers = do
+  subQ <- newTQueueIO
+  queueSubscribers <- SubscribedClients <$> TM.emptyIO
+  subClients <- newTVarIO IS.empty
+  pendingEvents <- newTVarIO IM.empty
+  pure ServerSubscribers {subQ, queueSubscribers, subClients, pendingEvents}
 
 newClient :: SQSType qs -> SMSType ms -> ClientId -> Natural -> VersionSMP -> ByteString -> SystemTime -> IO (Client (MsgStore qs ms))
 newClient _ _ clientId qSize thVersion sessionId createdAt = do
@@ -312,7 +446,24 @@ newClient _ _ clientId qSize thVersion sessionId createdAt = do
   connected <- newTVarIO True
   rcvActiveAt <- newTVarIO createdAt
   sndActiveAt <- newTVarIO createdAt
-  return Client {clientId, subscriptions, ntfSubscriptions, rcvQ, sndQ, msgQ, procThreads, endThreads, endThreadSeq, thVersion, sessionId, connected, createdAt, rcvActiveAt, sndActiveAt}
+  return
+    Client
+      { clientId,
+        subscriptions,
+        ntfSubscriptions,
+        rcvQ,
+        sndQ,
+        msgQ,
+        procThreads,
+        endThreads,
+        endThreadSeq,
+        thVersion,
+        sessionId,
+        connected,
+        createdAt,
+        rcvActiveAt,
+        sndActiveAt
+      }
 
 newSubscription :: SubscriptionThread -> STM Sub
 newSubscription st = do
@@ -362,9 +513,24 @@ newEnv config@ServerConfig {smpCredentials, httpCredentials, serverStoreCfg, smp
   serverStats <- newServerStats =<< getCurrentTime
   sockets <- newTVarIO []
   clientSeq <- newTVarIO 0
-  clients <- newTVarIO mempty
   proxyAgent <- newSMPProxyAgent smpAgentCfg random
-  pure Env {serverActive, config, serverInfo, server, serverIdentity, msgStore, ntfStore, random, tlsServerCreds, httpServerCreds, serverStats, sockets, clientSeq, clients, proxyAgent}
+  pure
+    Env
+      { serverActive,
+        config,
+        serverInfo,
+        server,
+        serverIdentity,
+        msgStore,
+        ntfStore,
+        random,
+        tlsServerCreds,
+        httpServerCreds,
+        serverStats,
+        sockets,
+        clientSeq,
+        proxyAgent
+      }
   where
     loadStoreLog :: StoreQueueClass q => (RecipientId -> QueueRec -> IO q) -> FilePath -> STMQueueStore q -> IO ()
     loadStoreLog mkQ f st = do

src/Simplex/Messaging/Server/Env/STM.hs

@@ -33,13 +33,16 @@ data RealTimeMetrics = RealTimeMetrics
   { socketStats :: [(ServiceName, SocketStats)],
     threadsCount :: Int,
     clientsCount :: Int,
-    smpSubsCount :: Int,
-    smpSubClientsCount :: Int,
-    ntfSubsCount :: Int,
-    ntfSubClientsCount :: Int,
+    smpSubs :: RTSubscriberMetrics,
+    ntfSubs :: RTSubscriberMetrics,
     loadedCounts :: LoadedQueueCounts
   }
 
+data RTSubscriberMetrics = RTSubscriberMetrics
+  { subsCount :: Int,
+    subClientsCount :: Int
+  }
+
 {-# FOURMOLU_DISABLE\n#-}
 prometheusMetrics :: ServerMetrics -> RealTimeMetrics -> UTCTime -> Text
 prometheusMetrics sm rtm ts =
@@ -50,10 +53,8 @@ prometheusMetrics sm rtm ts =
       { socketStats,
         threadsCount,
         clientsCount,
-        smpSubsCount,
-        smpSubClientsCount,
-        ntfSubsCount,
-        ntfSubClientsCount,
+        smpSubs,
+        ntfSubs,
         loadedCounts
       } = rtm
     ServerStatsData
@@ -367,21 +368,21 @@ prometheusMetrics sm rtm ts =
       \# TYPE simplex_smp_clients_total gauge\n\
       \simplex_smp_clients_total " <> mshow clientsCount <> "\n\
       \\n\
-      \# HELP simplex_smp_subscribtion_total Total subscriptions\n\
+      \# HELP simplex_smp_subscribtion_total Total SMP subscriptions\n\
       \# TYPE simplex_smp_subscribtion_total gauge\n\
-      \simplex_smp_subscribtion_total " <> mshow smpSubsCount <> "\n# smpSubs\n\
+      \simplex_smp_subscribtion_total " <> mshow (subsCount smpSubs) <> "\n# smp.subsCount\n\
       \\n\
-      \# HELP simplex_smp_subscribtion_clients_total Subscribed clients, first counting method\n\
+      \# HELP simplex_smp_subscribtion_clients_total Subscribed clients\n\
       \# TYPE simplex_smp_subscribtion_clients_total gauge\n\
-      \simplex_smp_subscribtion_clients_total " <> mshow smpSubClientsCount <> "\n# smpSubClients\n\
+      \simplex_smp_subscribtion_clients_total " <> mshow (subClientsCount smpSubs) <> "\n# smp.subClientsCount\n\
       \\n\
       \# HELP simplex_smp_subscription_ntf_total Total notification subscripbtions (from ntf server)\n\
       \# TYPE simplex_smp_subscription_ntf_total gauge\n\
-      \simplex_smp_subscription_ntf_total " <> mshow ntfSubsCount <> "\n# ntfSubs\n\
+      \simplex_smp_subscription_ntf_total " <> mshow (subsCount ntfSubs) <> "\n# ntf.subsCount\n\
       \\n\
-      \# HELP simplex_smp_subscription_ntf_clients_total Total subscribed NTF servers, first counting method\n\
+      \# HELP simplex_smp_subscription_ntf_clients_total Total subscribed NTF servers\n\
       \# TYPE simplex_smp_subscription_ntf_clients_total gauge\n\
-      \simplex_smp_subscription_ntf_clients_total " <> mshow ntfSubClientsCount <> "\n# ntfSubClients\n\
+      \simplex_smp_subscription_ntf_clients_total " <> mshow (subClientsCount ntfSubs) <> "\n# ntf.subClientsCount\n\
       \\n\
       \# HELP simplex_smp_loaded_queues_queue_count Total loaded queues count (all queues for memory/journal storage)\n\
       \# TYPE simplex_smp_loaded_queues_queue_count gauge\n\