smp-server: group pending notifications by service in one pass

src/Simplex/Messaging/Server/QueueStore/Postgres.hs

@@ -50,7 +50,7 @@ import Data.Bitraversable (bimapM)
 import Data.Either (fromRight)
 import Data.Functor (($>))
 import Data.Int (Int64)
-import Data.List (foldl', intersperse, partition)
+import Data.List (intersperse)
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.Map.Strict as M
 import Data.Maybe (catMaybes, fromMaybe)
@@ -405,23 +405,15 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
 
   getQueueNtfServices :: PostgresQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = E.uninterruptibleMask_ $ runExceptT $ do
-    snIds <-
+    nsIds <-
       withDB' "getQueueNtfServices" st $ \db ->
-        DB.query db "SELECT ntf_service_id, notifier_id FROM msg_queues WHERE notifier_id IN ? AND deleted_at IS NULL" (Only (In (map fst ntfs)))
-    pure $
-      if null snIds
-        then ([], ntfs)
-        else
-          let snIds' = foldl' (\m (sId, nId) -> M.alter (Just . maybe (S.singleton nId) (S.insert nId)) sId m) M.empty snIds
-           in foldr addService ([], ntfs) (M.assocs snIds')
+        DB.query db "SELECT notifier_id, ntf_service_id FROM msg_queues WHERE notifier_id IN ? AND deleted_at IS NULL" (Only (In (map fst ntfs)))
+    let (ssNtfs, deleteNtfs) = foldr (addNtf $ M.fromList nsIds) (M.empty, []) ntfs
+    pure (M.assocs ssNtfs, deleteNtfs)
     where
-      addService ::
-        (Maybe ServiceId, S.Set NotifierId) ->
-        ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]) ->
-        ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)])
-      addService (serviceId, snIds) (ssNtfs, ntfs') =
-        let (sNtfs, restNtfs) = partition (\(nId, _) -> S.member nId snIds) ntfs'
-         in ((serviceId, sNtfs) : ssNtfs, restNtfs)
+      addNtf nsIds ntf@(nId, _) (ssNtfs, deleteNtfs) = case M.lookup nId nsIds of
+        Just serviceId -> (M.alter (Just . maybe [ntf] (ntf :)) serviceId ssNtfs, deleteNtfs)
+        Nothing -> (ssNtfs, ntf : deleteNtfs)
 
   getServiceQueueCountHash :: (PartyI p, ServiceParty p) => PostgresQueueStore q -> SParty p -> ServiceId -> IO (Either ErrorType (Int64, IdsHash))
   getServiceQueueCountHash st party serviceId =

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

@@ -31,9 +31,9 @@ import Control.Logger.Simple
 import Control.Monad
 import Data.Bifunctor (first)
 import Data.Bitraversable (bimapM)
+import Data.Foldable (foldrM)
 import Data.Functor (($>))
 import Data.Int (Int64)
-import Data.List (partition)
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.Map.Strict as M
 import Data.Set (Set)
@@ -348,17 +348,16 @@ instance StoreQueueClass q => QueueStoreClass q (STMQueueStore q) where
 
   getQueueNtfServices :: STMQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = do
-    ss <- readTVarIO (services st)
-    (ssNtfs, noServiceNtfs) <- if M.null ss then pure ([], ntfs) else foldM addService ([], ntfs) (M.assocs ss)
     ns <- readTVarIO (notifiers st)
-    let (ntfs', deleteNtfs) = partition (\(nId, _) -> M.member nId ns) noServiceNtfs
-        ssNtfs' = (Nothing, ntfs') : ssNtfs
-    pure $ Right (ssNtfs', deleteNtfs)
+    qs <- readTVarIO (queues st)
+    Right . first M.assocs <$> foldrM (addNtf ns qs) (M.empty, []) ntfs
     where
-      addService (ssNtfs, ntfs') (serviceId, s) = do
-        (snIds, _) <- readTVarIO $ serviceNtfQueues s
-        let (sNtfs, restNtfs) = partition (\(nId, _) -> S.member nId snIds) ntfs'
-        pure ((Just serviceId, sNtfs) : ssNtfs, restNtfs)
+      addNtf ns qs ntf@(nId, _) (ssNtfs, deleteNtfs) = case M.lookup nId ns >>= (`M.lookup` qs) of
+        Just q -> do
+          qr <- readTVarIO (queueRec q)
+          let serviceId = qr >>= notifier >>= ntfServiceId
+          pure (M.alter (Just . maybe [ntf] (ntf :)) serviceId ssNtfs, deleteNtfs)
+        Nothing -> pure (ssNtfs, ntf : deleteNtfs)
 
   getServiceQueueCountHash :: (PartyI p, ServiceParty p) => STMQueueStore q -> SParty p -> ServiceId -> IO (Either ErrorType (Int64, IdsHash))
   getServiceQueueCountHash st party serviceId =

tests/CoreTests/MsgStoreTests.hs

@@ -25,6 +25,8 @@ import Data.ByteString.Char8 (ByteString)
 import Data.Int (Int64)
 import qualified Data.Map.Strict as M
 import Data.Time.Clock.System (SystemTime (..), getSystemTime)
+import qualified Data.X509 as X
+import qualified Data.X509.Validation as XV
 import Simplex.Messaging.Crypto (pattern MaxLenBS)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Protocol (EncDataBytes (..), EntityId (..), ErrorType (..), LinkId, Message (..), NotifierId, QueueLinkData, RecipientId, SParty (..), noMsgFlags, randomMsgId)
@@ -35,8 +37,11 @@ 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.SystemTime (getSystemDate)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
+import Simplex.Messaging.Transport (SMPServiceRole (..))
+import Simplex.Messaging.Transport.Credentials (genCredentials)
 import Test.Hspec hiding (fit, it)
 import Util
 
@@ -73,6 +78,7 @@ msgStoreTests = do
       it "should get queue and store/read messages" $ withMsgStore cfg testGetQueue
       it "should write/ack messages" $ withMsgStore cfg testWriteAckMessages
       it "should not add link data to secured messaging queue" $ withMsgStore cfg testLinkDataSecuredQueue
+      it "should group pending notifiers by service" $ withMsgStore cfg testGetQueueNtfServices
 
 -- TODO constrain to STM stores?
 withMsgStore :: MsgStoreClass s => MsgStoreConfig s -> (s -> IO ()) -> IO ()
@@ -147,6 +153,20 @@ testNtfCreds g = do
         ntfServiceId = Nothing
       }
 
+newTestServiceRec :: TVar ChaChaDRG -> IO ServiceRec
+newTestServiceRec g = do
+  serviceId <- atomically $ EntityId <$> C.randomBytes 24 g
+  (_, cert) <- genCredentials g Nothing (0, 2400) "ntf.example.com"
+  serviceCreatedAt <- getSystemDate
+  pure
+    ServiceRec
+      { serviceId,
+        serviceRole = SRNotifier,
+        serviceCert  = X.CertificateChain [cert],
+        serviceCertHash = XV.getFingerprint cert X.HashSHA256,
+        serviceCreatedAt
+      }
+
 testGetQueue :: MsgStoreClass s => s -> IO ()
 testGetQueue ms = do
   g <- C.newRandom
@@ -238,6 +258,30 @@ testLinkDataSecuredQueue ms = do
     void $ ExceptT $ deleteQueue ms q
     void $ ExceptT $ deleteQueue ms cq
 
+testGetQueueNtfServices :: forall s. MsgStoreClass s => s -> IO ()
+testGetQueueNtfServices ms = do
+  g <- C.newRandom
+  ntfCreds <- testNtfCreds g
+  sr <- newTestServiceRec g
+  sr' <- newTestServiceRec g
+  let st = queueStore ms
+      ntf1 = (EntityId "notifier 1", ())
+      ntf2 = (EntityId "notifier 2", ())
+      ntf3 = (EntityId "notifier 3", ())
+      ntf4 = (EntityId "notifier 4", ())
+      addNtfQueue (notifierId, _) ntfServiceId = do
+        (rId, qr) <- liftIO $ testNewQueueRec g QMMessaging
+        void $ ExceptT $ addQueue ms rId qr {notifier = Just ntfCreds {notifierId, ntfServiceId}}
+  runRight_ $ do
+    serviceId <- ExceptT $ getCreateService @(StoreQueue s) st sr
+    void $ ExceptT $ getCreateService @(StoreQueue s) st sr'
+    addNtfQueue ntf1 Nothing
+    addNtfQueue ntf2 $ Just serviceId
+    addNtfQueue ntf3 $ Just serviceId
+    liftIO $
+      getQueueNtfServices @(StoreQueue s) st [ntf1, ntf2, ntf3, ntf4]
+        `shouldReturn` Right ([(Nothing, [ntf1]), (Just serviceId, [ntf2, ntf3])], [ntf4])
+
 -- sizes of queues, senders, links and notifiers maps
 type QueueMapSizes = (Int, Int, Int, Int)
 

tests/CoreTests/StoreLogTests.hs

@@ -12,13 +12,10 @@ module CoreTests.StoreLogTests where
 import Control.Concurrent.STM
 import Control.Monad
 import CoreTests.MsgStoreTests
-import Crypto.Random (ChaChaDRG)
 import qualified Data.ByteString.Char8 as B
 import Data.Either (partitionEithers)
 import qualified Data.List.NonEmpty as L
 import qualified Data.Map.Strict as M
-import qualified Data.X509 as X
-import qualified Data.X509.Validation as XV
 import SMPClient
 import Simplex.FileTransfer.Protocol (FileInfo (..))
 import Simplex.FileTransfer.Server.Store (FileRec (..), FileRecipient (..), FileStoreClass (..), RoundedFileTime, STMFileStore (..))
@@ -36,7 +33,6 @@ import Simplex.Messaging.Server.QueueStore.Types
 import Simplex.Messaging.Server.StoreLog
 import Simplex.Messaging.SystemTime
 import Simplex.Messaging.Transport (SMPServiceRole (..))
-import Simplex.Messaging.Transport.Credentials (genCredentials)
 import System.FilePath ((</>))
 import Test.Hspec hiding (fit, it)
 import Util
@@ -151,20 +147,6 @@ storeLogTests =
           }
       ]
 
-newTestServiceRec :: TVar ChaChaDRG -> IO ServiceRec
-newTestServiceRec g = do
-  serviceId <- atomically $ EntityId <$> C.randomBytes 24 g
-  (_, cert) <- genCredentials g Nothing (0, 2400) "ntf.example.com"
-  serviceCreatedAt <- getSystemDate
-  pure
-    ServiceRec
-      { serviceId,
-        serviceRole = SRNotifier,
-        serviceCert  = X.CertificateChain [cert],
-        serviceCertHash = XV.getFingerprint cert X.HashSHA256,
-        serviceCreatedAt
-      }
-
 testSMPStoreLog :: String -> [SMPStoreLogTestCase] -> Spec
 testSMPStoreLog testSuite tests =
   describe testSuite $ forM_ tests $ \t@SLTC {name, saved} -> it name $ do