smp-server: fix notification store key leak

src/Simplex/Messaging/Server.hs

@@ -366,7 +366,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
         runDeliverNtfs ms ns' stats
       where
         runDeliverNtfs :: s -> NtfStore -> ServerStats -> IO ()
-        runDeliverNtfs ms (NtfStore ns) stats = do
+        runDeliverNtfs ms ns'@(NtfStore ns) stats = do
           ntfs <- M.assocs <$> readTVarIO ns
           unless (null ntfs) $
             getQueueNtfServices @(StoreQueue s) (queueStore ms) ntfs >>= \case
@@ -378,6 +378,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                     cIds <- IS.toList <$> readTVarIO subClients
                     forM_ cIds $ \cId -> getServerClient cId srv >>= mapM_ (deliverQueueNtfs ntfs')
                 atomically $ modifyTVar' ns (`M.withoutKeys` S.fromList (map fst deleted))
+                deleteEmptyNtfs ns' ntfs
           where
             deliverQueueNtfs ntfs' c@Client {ntfSubscriptions} =
               whenM (currentClient readTVarIO c) $ do

src/Simplex/Messaging/Server/NtfStore.hs

@@ -10,6 +10,7 @@ module Simplex.Messaging.Server.NtfStore
     storeNtf,
     deleteNtfs,
     deleteExpiredNtfs,
+    deleteEmptyNtfs,
   ) where
 
 import Control.Concurrent.STM
@@ -22,6 +23,7 @@ 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)
 
 newtype NtfStore = NtfStore (TMap NotifierId (TVar [MsgNtf]))
 
@@ -33,13 +35,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 +47,25 @@ 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
+
+deleteEmptyNtfs :: NtfStore -> [(NotifierId, TVar [MsgNtf])] -> IO ()
+deleteEmptyNtfs (NtfStore ns) = mapM_ deleteEmpty
+  where
+    deleteEmpty (nId, v) =
+      whenM (null <$> readTVarIO v) $
+        atomically $
+          TM.lookup nId ns >>= mapM_ (\v' -> whenM (null <$> readTVar v') $ TM.delete nId ns)
 
 data NtfLogRecord = NLRv1 NotifierId MsgNtf
 

tests/CoreTests/MsgStoreTests.hs

@@ -24,24 +24,25 @@ import Crypto.Random (ChaChaDRG)
 import Data.ByteString.Char8 (ByteString)
 import Data.Int (Int64)
 import qualified Data.Map.Strict as M
-import Data.Time.Clock.System (getSystemTime)
+import Data.Time.Clock.System (SystemTime (..), getSystemTime)
 import Simplex.Messaging.Crypto (pattern MaxLenBS)
 import qualified Simplex.Messaging.Crypto as C
-import Simplex.Messaging.Protocol (EncDataBytes (..), EntityId (..), ErrorType (..), LinkId, Message (..), QueueLinkData, RecipientId, SParty (..), noMsgFlags, randomMsgId)
+import Simplex.Messaging.Protocol (EncDataBytes (..), EntityId (..), ErrorType (..), LinkId, Message (..), NotifierId, QueueLinkData, RecipientId, SParty (..), noMsgFlags, randomMsgId)
 import Simplex.Messaging.Server.MsgStore.STM
 import Simplex.Messaging.Server.MsgStore.Types
+import Simplex.Messaging.Server.NtfStore
 import Simplex.Messaging.Server.QueueStore
 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 Test.Hspec hiding (fit, it)
 import Util
 
 #if defined(dbServerPostgres)
 import AgentTests.FunctionalAPITests (runRight)
 import Control.Concurrent (threadDelay)
-import Data.Time.Clock.System (SystemTime (..))
 import Database.PostgreSQL.Simple (Only (..))
 import qualified Database.PostgreSQL.Simple as DB
 import Simplex.Messaging.Agent.Store.Postgres.Common
@@ -63,6 +64,9 @@ msgStoreTests = do
     it "tryDelPeekMsg (ACK not from NSE) should reset message counts when queue is empty" $ withMsgStore testPostgresStoreConfig testResetMessageCounts
     it "should expire messages across commit batches" $ withMsgStore testPostgresStoreConfig testExpireMessagesInBatches
 #endif
+  describe "Notification store" $ do
+    it "should remove keys of expired notifications" testExpireNtfs
+    it "should remove keys of delivered notifications" testDeleteEmptyNtfs
   where
     someMsgStoreTests :: MsgStoreClass s => (HasTestEnv => MsgStoreConfig s) -> Spec
     someMsgStoreTests cfg = do
@@ -281,6 +285,54 @@ testDeleteQueueMaps ms = do
   where
     expectAuth = either (`shouldBe` AUTH) (\_ -> expectationFailure "deleted queue is still found")
 
+mkMsgNtf :: TVar ChaChaDRG -> Int64 -> IO MsgNtf
+mkMsgNtf g ts = do
+  ntfMsgId <- atomically $ randomMsgId 24 g
+  ntfNonce <- atomically $ C.randomCbNonce g
+  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)
+
+testExpireNtfs :: IO ()
+testExpireNtfs = do
+  g <- C.newRandom
+  st@(NtfStore ns) <- NtfStore <$> TM.emptyIO
+  now <- systemSeconds <$> getSystemTime
+  let old = now - 100
+      nId1 = EntityId "notifier 1"
+      nId2 = EntityId "notifier 2"
+      nId3 = EntityId "notifier 3"
+  storeNtf st nId1 =<< mkMsgNtf g old
+  storeNtf st nId2 =<< mkMsgNtf g old
+  storeNtf st nId2 =<< mkMsgNtf g now
+  atomically $ TM.insertM nId3 (newTVar []) ns
+  ntfCounts st `shouldReturn` [(nId1, 1), (nId2, 2), (nId3, 0)]
+  deleteExpiredNtfs st (now - 50) `shouldReturn` 2
+  ntfCounts st `shouldReturn` [(nId2, 1)]
+  storeNtf st nId1 =<< mkMsgNtf g now
+  deleteExpiredNtfs st (now - 50) `shouldReturn` 0
+  ntfCounts st `shouldReturn` [(nId1, 1), (nId2, 1)]
+
+testDeleteEmptyNtfs :: IO ()
+testDeleteEmptyNtfs = do
+  g <- C.newRandom
+  st@(NtfStore ns) <- NtfStore <$> TM.emptyIO
+  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
+  let flush nId = forM_ (lookup nId ntfs) $ \v -> atomically $ writeTVar v []
+  mapM_ flush ([nId1, nId2, nId4] :: [NotifierId])
+  storeNtf st nId2 =<< mkMsgNtf g now
+  deleteNtfs st nId4 `shouldReturn` 0
+  storeNtf st nId4 =<< mkMsgNtf g now
+  deleteEmptyNtfs st ntfs
+  ntfCounts st `shouldReturn` [(nId2, 1), (nId3, 1), (nId4, 1)]
+
 testLinkData :: QueueLinkData
 testLinkData = (EncDataBytes "fixed data", EncDataBytes "user data")