smp server: use stm-containers for the notification store

simplexmq.cabal

@@ -375,6 +375,7 @@ library
     , data-default ==0.7.*
     , directory ==1.3.*
     , filepath ==1.4.*
+    , hashable ==1.4.*
     , hourglass ==0.2.*
     , http-types ==0.12.*
     , http2 >=4.2.2 && <4.3
@@ -405,13 +406,14 @@ library
   if !flag(client_library)
     build-depends:
         case-insensitive ==1.2.*
-      , hashable ==1.4.*
       , http-client >=0.7 && <0.8
       , http-client-tls >=0.3 && <0.4
       , ini ==0.4.1
+      , list-t ==1.0.*
       , network-uri >=2.6 && <2.7
       , optparse-applicative >=0.15 && <0.17
       , process ==1.6.*
+      , stm-containers ==1.2.*
       , temporary ==1.3.*
       , wai >=3.2 && <3.3
       , wai-app-static >=3.1 && <3.2
@@ -638,6 +640,7 @@ test-suite simplexmq-test
     , simple-logger
     , simplexmq
     , stm
+    , stm-containers
     , text
     , time
     , tls >=1.9.0 && <1.10

src/Simplex/Messaging/Server.hs

@@ -75,7 +75,6 @@ import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
 import Data.Maybe (fromMaybe, isJust, isNothing)
 import Data.Semigroup (Sum (..))
-import qualified Data.Set as S
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeLatin1)
@@ -366,8 +365,8 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
         runDeliverNtfs ms ns' stats
       where
         runDeliverNtfs :: s -> NtfStore -> ServerStats -> IO ()
-        runDeliverNtfs ms ns'@(NtfStore ns) stats = do
-          ntfs <- M.assocs <$> readTVarIO ns
+        runDeliverNtfs ms ns' stats = do
+          ntfs <- getNtfs ns'
           unless (null ntfs) $
             getQueueNtfServices @(StoreQueue s) (queueStore ms) ntfs >>= \case
               Left e -> logError $ "NOTIFICATIONS: getQueueNtfServices error " <> tshow e
@@ -377,7 +376,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                   Nothing -> do -- legacy code that does almost the same as before for non-service subscribers
                     cIds <- IS.toList <$> readTVarIO subClients
                     forM_ cIds $ \cId -> getServerClient cId srv >>= mapM_ (deliverQueueNtfs ntfs')
-                atomically $ modifyTVar' ns (`M.withoutKeys` S.fromList (map fst deleted))
+                mapM_ (deleteNtfs ns' . fst) deleted
                 deleteEmptyNtfs ns' ntfs
           where
             deliverQueueNtfs ntfs' c@Client {ntfSubscriptions} =
@@ -854,8 +853,8 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                 putStat "ntfCount" ntfCount
                 readTVarIO role >>= \case
                   CPRAdmin -> do
-                    NtfStore ns <- unliftIO u $ asks ntfStore
-                    ntfCount2 <- liftIO . foldM (\(!n) q -> (n +) . length <$> readTVarIO q) 0 =<< readTVarIO ns
+                    ns <- unliftIO u $ asks ntfStore
+                    ntfCount2 <- liftIO . foldM (\(!n) q -> (n +) . length <$> readTVarIO q) 0 . map snd =<< getNtfs ns
                     hPutStrLn h $ "ntfCount 2: " <> show ntfCount2
                   _ -> pure ()
                 putProxyStat "pRelays" pRelays
@@ -2470,9 +2469,9 @@ saveServerNtfs = asks (storeNtfsFile . config) >>= mapM_ saveNtfs
   where
     saveNtfs f = do
       logNote $ "saving notifications to file " <> T.pack f
-      NtfStore ns <- asks ntfStore
+      ns <- asks ntfStore
       liftIO . withFile f WriteMode $ \h ->
-        readTVarIO ns >>= mapM_ (saveQueueNtfs h) . M.assocs
+        getNtfs ns >>= mapM_ (saveQueueNtfs h)
       logNote "notifications saved"
       where
         -- reverse on save, to save notifications in order, will become reversed again when restoring.
@@ -2496,8 +2495,7 @@ restoreServerNtfs =
             liftIO exitFailure
           Right (lineCount, storedMsgsCount, expiredMsgsCount) -> do
             renameFile f $ f <> ".bak"
-            let NtfStore ns' = ns
-            storedQueues <- M.size <$> readTVarIO ns'
+            storedQueues <- length <$> getNtfs ns
             logNote $ "notifications restored, " <> tshow lineCount <> " lines processed"
             pure MessageStats {storedMsgsCount, expiredMsgsCount, storedQueues}
       where

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

@@ -127,6 +127,7 @@ import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (ASrvTransport, SMPVersion, THandleParams, TransportPeer (..), VersionRangeSMP)
 import Simplex.Messaging.Transport.Server
 import Simplex.Messaging.Util (ifM, tshow, whenM, ($>>=))
+import qualified StmContainers.Map as SM
 import System.Directory (doesFileExist)
 import System.Exit (exitFailure)
 import System.IO (IOMode (..))
@@ -564,7 +565,7 @@ newEnv config@ServerConfig {smpCredentials, httpCredentials, serverStoreCfg, smp
       when compactLog $ compactDbStoreLog $ dbStoreLogPath storeCfg
       StoreDatabase <$> newMsgStore cfg
 #endif
-  ntfStore <- NtfStore <$> TM.emptyIO
+  ntfStore <- NtfStore <$> SM.newIO
   random <- C.newRandom
   tlsServerCreds <- getCredentials "SMP" smpCredentials
   httpServerCreds <- mapM (getCredentials "HTTPS") httpCredentials

src/Simplex/Messaging/Server/NtfStore.hs

@@ -11,22 +11,22 @@ module Simplex.Messaging.Server.NtfStore
     deleteNtfs,
     deleteExpiredNtfs,
     deleteEmptyNtfs,
+    getNtfs,
   ) where
 
 import Control.Concurrent.STM
 import Control.Monad (foldM)
 import Data.Functor (($>))
 import Data.Int (Int64)
-import qualified Data.Map.Strict as M
 import Data.Time.Clock.System (SystemTime (..))
+import qualified ListT
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (EncNMsgMeta, MsgId, NotifierId)
-import Simplex.Messaging.TMap (TMap)
-import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Util (whenM)
+import qualified StmContainers.Map as SM
 
-newtype NtfStore = NtfStore (TMap NotifierId (TVar [MsgNtf]))
+newtype NtfStore = NtfStore (SM.Map NotifierId (TVar [MsgNtf]))
 
 data MsgNtf = MsgNtf
   { ntfMsgId :: MsgId,
@@ -37,25 +37,25 @@ data MsgNtf = MsgNtf
 
 storeNtf :: NtfStore -> NotifierId -> MsgNtf -> IO ()
 storeNtf (NtfStore ns) nId ntf =
-  atomically $ TM.lookup nId ns >>= maybe (TM.insertM nId (newTVar [ntf]) ns) (`modifyTVar'` (ntf :))
+  atomically $ SM.lookup nId ns >>= maybe (newTVar [ntf] >>= \v -> SM.insert v nId ns) (`modifyTVar'` (ntf :))
   -- TODO [ntfdb] coalesce messages here once the client is updated to process multiple messages
   -- for single notification.
   -- when (isJust prevNtf) $ incStat $ msgNtfReplaced stats
 
 deleteNtfs :: NtfStore -> NotifierId -> IO Int
-deleteNtfs (NtfStore ns) nId = atomically (TM.lookupDelete nId ns) >>= maybe (pure 0) (fmap length . readTVarIO)
+deleteNtfs (NtfStore ns) nId = atomically (SM.lookup nId ns <* SM.delete nId ns) >>= maybe (pure 0) (fmap length . readTVarIO)
 
 deleteExpiredNtfs :: NtfStore -> Int64 -> IO Int
-deleteExpiredNtfs (NtfStore ns) old =
-  foldM (\expired -> fmap (expired +) . expireQueue) 0 . M.keys =<< readTVarIO ns
+deleteExpiredNtfs st@(NtfStore ns) old =
+  foldM (\expired -> fmap (expired +) . expireQueue) 0 . map fst =<< getNtfs st
   where
-    expireQueue nId = atomically $ TM.lookup nId ns >>= maybe (pure 0) (expire nId)
+    expireQueue nId = atomically $ SM.lookup nId ns >>= maybe (pure 0) (expire nId)
     expire nId v = readTVar v >>= \case
-      [] -> TM.delete nId ns $> 0
+      [] -> SM.delete nId ns $> 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
-        if null ntfs' then TM.delete nId ns else writeTVar v ntfs'
+        if null ntfs' then SM.delete nId ns else writeTVar v ntfs'
         pure $! length ntfs - length ntfs'
       _ -> pure 0
 
@@ -65,7 +65,12 @@ deleteEmptyNtfs (NtfStore ns) = mapM_ deleteEmpty
     deleteEmpty (nId, v) =
       whenM (null <$> readTVarIO v) $
         atomically $
-          TM.lookup nId ns >>= mapM_ (\v' -> whenM (null <$> readTVar v') $ TM.delete nId ns)
+          SM.lookup nId ns >>= mapM_ (\v' -> whenM (null <$> readTVar v') $ SM.delete nId ns)
+
+getNtfs :: NtfStore -> IO [(NotifierId, TVar [MsgNtf])]
+getNtfs (NtfStore ns) = go [] $ SM.listT ns
+  where
+    go ntfs l = atomically (ListT.uncons l) >>= maybe (pure ntfs) (\(ntf, l') -> go (ntf : ntfs) l')
 
 data NtfLogRecord = NLRv1 NotifierId MsgNtf
 

src/Simplex/Messaging/Transport.hs

@@ -122,6 +122,7 @@ import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Default (def)
 import Data.Functor (($>))
+import Data.Hashable (Hashable)
 import Data.Int (Int64)
 import Data.Kind (Type)
 import Data.Tuple (swap)
@@ -528,7 +529,7 @@ type ServiceId = EntityId
 -- this type is used for server entities only
 newtype EntityId = EntityId {unEntityId :: ByteString}
   deriving (Eq, Ord, Show)
-  deriving newtype (Encoding, StrEncoding)
+  deriving newtype (Encoding, Hashable, StrEncoding)
 
 pattern NoEntity :: EntityId
 pattern NoEntity = EntityId ""

tests/CoreTests/MsgStoreTests.hs

@@ -36,7 +36,7 @@ import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.Server.QueueStore.STM (STMQueueStore (..))
 import Simplex.Messaging.Server.QueueStore.Types
 import Simplex.Messaging.TMap (TMap)
-import qualified Simplex.Messaging.TMap as TM
+import qualified StmContainers.Map as SM
 import Test.Hspec hiding (fit, it)
 import Util
 
@@ -292,12 +292,12 @@ mkMsgNtf g ts = do
   pure MsgNtf {ntfMsgId, ntfTs = MkSystemTime ts 0, ntfNonce, ntfEncMeta = "meta"}
 
 ntfCounts :: NtfStore -> IO [(NotifierId, Int)]
-ntfCounts (NtfStore ns) = M.assocs <$> (mapM (fmap length . readTVarIO) =<< readTVarIO ns)
+ntfCounts st = M.assocs <$> (mapM (fmap length . readTVarIO) . M.fromList =<< getNtfs st)
 
 testExpireNtfs :: IO ()
 testExpireNtfs = do
   g <- C.newRandom
-  st@(NtfStore ns) <- NtfStore <$> TM.emptyIO
+  st@(NtfStore ns) <- NtfStore <$> SM.newIO
   now <- systemSeconds <$> getSystemTime
   let old = now - 100
       nId1 = EntityId "notifier 1"
@@ -306,7 +306,7 @@ testExpireNtfs = do
   storeNtf st nId1 =<< mkMsgNtf g old
   storeNtf st nId2 =<< mkMsgNtf g old
   storeNtf st nId2 =<< mkMsgNtf g now
-  atomically $ TM.insertM nId3 (newTVar []) ns
+  atomically $ newTVar [] >>= \v -> SM.insert v nId3 ns
   ntfCounts st `shouldReturn` [(nId1, 1), (nId2, 2), (nId3, 0)]
   deleteExpiredNtfs st (now - 50) `shouldReturn` 2
   ntfCounts st `shouldReturn` [(nId2, 1)]
@@ -317,14 +317,14 @@ testExpireNtfs = do
 testDeleteEmptyNtfs :: IO ()
 testDeleteEmptyNtfs = do
   g <- C.newRandom
-  st@(NtfStore ns) <- NtfStore <$> TM.emptyIO
+  st <- NtfStore <$> SM.newIO
   now <- systemSeconds <$> getSystemTime
   let nId1 = EntityId "notifier 1"
       nId2 = EntityId "notifier 2"
       nId3 = EntityId "notifier 3"
       nId4 = EntityId "notifier 4"
   forM_ ([nId1, nId2, nId3, nId4] :: [NotifierId]) $ \nId -> storeNtf st nId =<< mkMsgNtf g now
-  ntfs <- M.assocs <$> readTVarIO ns
+  ntfs <- getNtfs st
   let flush nId = forM_ (lookup nId ntfs) $ \v -> atomically $ writeTVar v []
   mapM_ flush ([nId1, nId2, nId4] :: [NotifierId])
   storeNtf st nId2 =<< mkMsgNtf g now