From 056f8e6e5f56da44ff9375043faf687e844f877e Mon Sep 17 00:00:00 2001 From: Evgeny Poberezkin Date: Wed, 16 Jul 2025 22:31:51 +0100 Subject: [PATCH] efficient initial delivery of messages to subscribed service --- src/Simplex/Messaging/Server.hs | 42 +++++++++---------- .../Messaging/Server/MsgStore/Journal.hs | 23 ++++++++-- src/Simplex/Messaging/Server/MsgStore/STM.hs | 6 ++- .../Messaging/Server/MsgStore/Types.hs | 2 +- .../Messaging/Server/QueueStore/Postgres.hs | 18 ++++---- .../Messaging/Server/QueueStore/STM.hs | 20 ++++----- .../Messaging/Server/QueueStore/Types.hs | 1 - 7 files changed, 63 insertions(+), 49 deletions(-) diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index 4cc08cf6c..c163e18b6 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/src/Simplex/Messaging/Server.hs @@ -1489,8 +1489,8 @@ client rcvId <- randId ntf <- forM ntfKeys_ $ \(notifierKey, rcvNtfDhSecret, rcvPubDhKey) -> do notifierId <- randId - let ntfCreds = NtfCreds {notifierId, notifierKey, rcvNtfDhSecret, ntfServiceId = Nothing} - pure (ntfCreds, ServerNtfCreds notifierId rcvPubDhKey) + let ntfCreds' = NtfCreds {notifierId, notifierKey, rcvNtfDhSecret, ntfServiceId = Nothing} + pure (ntfCreds', ServerNtfCreds notifierId rcvPubDhKey) let queueMode = queueReqMode <$> queueReqData qr = QueueRec @@ -1603,19 +1603,19 @@ client msg_ <- tryPeekMsg ms q msg' <- forM msg_ $ \msg -> liftIO $ do ts <- getSystemSeconds - sub <- maybe (atomically $ getSub entId) pure sub_ + sub <- maybe (atomically getSub) pure sub_ atomically $ setDelivered sub msg ts unless hasSub $ incStat $ qSub stats pure (NoCorrId, entId, MSG (encryptMsg qr msg)) pure ((corrId, entId, SOK clntServiceId), msg') - getSub :: RecipientId -> STM Sub - getSub rId = - TM.lookup rId (subscriptions clnt) >>= \case + getSub :: STM Sub + getSub = + TM.lookup entId (subscriptions clnt) >>= \case Just sub -> pure sub Nothing -> do sub <- newSubscription NoSub - TM.insert rId sub $ subscriptions clnt + TM.insert entId sub $ subscriptions clnt pure sub subscribeNewQueue :: RecipientId -> QueueRec -> M s () @@ -1736,12 +1736,11 @@ client incServiceQueueSubs = modifyTVar' (clientServiceSubs clnt) (+ 1) -- service count Nothing -> case queueServiceId of Just _ -> runExceptT $ do - -- getSubscription should never be Just in this branch, because queue was associated with service. - -- So unless storage and session states diverge, this check is redundant. ExceptT $ setQueueService (queueStore ms) q party Nothing liftIO $ incSrvStat srvAssocRemoved + -- getSub may be Just for receiving service, where clientSubs also hold active deliveries for service subscriptions. + -- For notification service it can only be Just if storage and session states diverge. r <- atomically $ getSubscription >>= newSub - atomically writeSub pure r Nothing -> do @@ -1766,26 +1765,27 @@ client pure $ SOKS count where deliverServiceMessages expectedCnt = do - (qCnt, _msgCnt, _dupCnt, _errCnt) <- foldRcvServiceQueues (queueStore ms) (mkQueue ms True) serviceId deliverQueueMsg (0, 0, 0, 0) + (qCnt, _msgCnt, _dupCnt, _errCnt) <- foldRcvServiceMessages ms serviceId deliverQueueMsg (0, 0, 0, 0) -- TODO [cert rcv] compare with expected logNote $ "Service subscriptions for " <> tshow serviceId <> " (" <> tshow qCnt <> " queues)" - deliverQueueMsg :: (Int, Int, Int, Int) -> (StoreQueue s, QueueRec) -> IO (Int, Int, Int, Int) - deliverQueueMsg acc@(!qCnt, !msgCnt, !dupCnt, !errCnt) (q, qr) = - runExceptT (tryPeekMsg ms q) >>= \case - Left e -> pure (qCnt + 1, msgCnt, dupCnt, errCnt + 1) -- TODO deliver subscription error - Right Nothing -> pure acc - Right (Just msg) -> do - let rId = recipientId q - atomically (getNewSub rId) >>= \case + deliverQueueMsg :: (Int, Int, Int, Int) -> RecipientId -> Either ErrorType (Maybe (QueueRec, Message)) -> IO (Int, Int, Int, Int) + deliverQueueMsg (!qCnt, !msgCnt, !dupCnt, !errCnt) rId = \case + Left e -> pure (qCnt + 1, msgCnt, dupCnt, errCnt + 1) -- TODO deliver subscription error + Right qMsg_ -> case qMsg_ of + Nothing -> pure (qCnt + 1, msgCnt, dupCnt, errCnt) + Just (qr, msg) -> + atomically (getSubscription rId) >>= \case Nothing -> pure (qCnt + 1, msgCnt, dupCnt + 1, errCnt) Just sub -> do ts <- getSystemSeconds atomically $ setDelivered sub msg ts atomically $ writeTBQueue msgQ [(NoCorrId, rId, MSG (encryptMsg qr msg))] pure (qCnt + 1, msgCnt + 1, dupCnt, errCnt) - getNewSub rId = + getSubscription rId = TM.lookup rId (subscriptions clnt) >>= \case - Just _sub -> pure Nothing -- if delivery subscription already exists, then there is no need to deliver message + -- If delivery subscription already exists, then there is no need to deliver message. + -- It may have been created when the message is sent after service subscription is created. + Just _sub -> pure Nothing Nothing -> do sub <- newSubscription NoSub TM.insert rId sub $ subscriptions clnt diff --git a/src/Simplex/Messaging/Server/MsgStore/Journal.hs b/src/Simplex/Messaging/Server/MsgStore/Journal.hs index 4fc656552..8561e8f3e 100644 --- a/src/Simplex/Messaging/Server/MsgStore/Journal.hs +++ b/src/Simplex/Messaging/Server/MsgStore/Journal.hs @@ -358,8 +358,6 @@ instance QueueStoreClass (JournalQueue s) (QStore s) where {-# INLINE getQueueNtfServices #-} getServiceQueueCount = withQS (getServiceQueueCount @(JournalQueue s)) {-# INLINE getServiceQueueCount #-} - foldRcvServiceQueues = withQS foldRcvServiceQueues - {-# INLINE foldRcvServiceQueues #-} makeQueue_ :: JournalMsgStore s -> RecipientId -> QueueRec -> Lock -> IO (JournalQueue s) makeQueue_ JournalMsgStore {sharedLock} rId qr queueLock = do @@ -446,8 +444,25 @@ instance MsgStoreClass (JournalMsgStore s) where getLoadedQueue :: JournalQueue s -> IO (JournalQueue s) getLoadedQueue q = fromMaybe q <$> TM.lookupIO (recipientId q) (loadedQueues $ queueStore_ ms) - foldRcvServiceMessages :: JournalMsgStore s -> ServiceId -> (a -> RecipientId -> Message -> IO a) -> IO a - foldRcvServiceMessages = undefined + foldRcvServiceMessages :: JournalMsgStore s -> ServiceId -> (a -> RecipientId -> Either ErrorType (Maybe (QueueRec, Message)) -> IO a) -> a -> IO a + foldRcvServiceMessages ms serviceId f acc = case queueStore_ ms of + MQStore st -> foldRcvServiceQueues st serviceId f' acc + where + f' a (q, qr) = runExceptT (tryPeekMsg ms q) >>= f a (recipientId q) . ((qr,) <$$>) +#if defined(dbServerPostgres) + PQStore st -> foldRcvServiceQueueRecs st serviceId f' acc + where + JournalMsgStore {queueLocks, sharedLock} = ms + f' a (rId, qr) = do + q <- mkQueue ms False rId qr + qMsg_ <- + withSharedWaitLock rId queueLocks sharedLock $ runExceptT $ tryStore' "foldRcvServiceMessages" rId $ + (qr,) . snd <$$> (getLoadedQueue q >>= unStoreIO . getPeekMsgQueue ms) + f a rId qMsg_ + -- Use cached queue if available. + -- Also see the comment in loadQueue in PostgresQueueStore + getLoadedQueue q = fromMaybe q <$> TM.lookupIO (recipientId q) (loadedQueues $ queueStore_ ms) +#endif logQueueStates :: JournalMsgStore s -> IO () logQueueStates ms = withActiveMsgQueues ms $ unStoreIO . logQueueState diff --git a/src/Simplex/Messaging/Server/MsgStore/STM.hs b/src/Simplex/Messaging/Server/MsgStore/STM.hs index 36221542b..72eab6756 100644 --- a/src/Simplex/Messaging/Server/MsgStore/STM.hs +++ b/src/Simplex/Messaging/Server/MsgStore/STM.hs @@ -89,8 +89,10 @@ instance MsgStoreClass STMMsgStore where expireOldMessages _tty ms now ttl = withLoadedQueues (queueStore_ ms) $ atomically . expireQueueMsgs ms now (now - ttl) - foldRcvServiceMessages :: STMMsgStore -> ServiceId -> (a -> RecipientId -> Message -> IO a) -> IO a - foldRcvServiceMessages = undefined + foldRcvServiceMessages :: STMMsgStore -> ServiceId -> (a -> RecipientId -> Either ErrorType (Maybe (QueueRec, Message)) -> IO a) -> a -> IO a + foldRcvServiceMessages ms serviceId f= + foldRcvServiceQueues (queueStore_ ms) serviceId $ \a (q, qr) -> + runExceptT (tryPeekMsg ms q) >>= f a (recipientId q) . ((qr,) <$$>) logQueueStates _ = pure () {-# INLINE logQueueStates #-} diff --git a/src/Simplex/Messaging/Server/MsgStore/Types.hs b/src/Simplex/Messaging/Server/MsgStore/Types.hs index 7f63c1ddf..0729d71b6 100644 --- a/src/Simplex/Messaging/Server/MsgStore/Types.hs +++ b/src/Simplex/Messaging/Server/MsgStore/Types.hs @@ -44,7 +44,7 @@ class (Monad (StoreMonad s), QueueStoreClass (StoreQueue s) (QueueStore s)) => M unsafeWithAllMsgQueues :: Monoid a => Bool -> Bool -> s -> (StoreQueue s -> IO a) -> IO a -- tty, store, now, ttl expireOldMessages :: Bool -> s -> Int64 -> Int64 -> IO MessageStats - foldRcvServiceMessages :: s -> ServiceId -> (a -> RecipientId -> Message -> IO a) -> IO a + foldRcvServiceMessages :: s -> ServiceId -> (a -> RecipientId -> Either ErrorType (Maybe (QueueRec, Message)) -> IO a) -> a -> IO a logQueueStates :: s -> IO () logQueueState :: StoreQueue s -> StoreMonad s () queueStore :: s -> QueueStore s diff --git a/src/Simplex/Messaging/Server/QueueStore/Postgres.hs b/src/Simplex/Messaging/Server/QueueStore/Postgres.hs index 57b1b1493..759dbb501 100644 --- a/src/Simplex/Messaging/Server/QueueStore/Postgres.hs +++ b/src/Simplex/Messaging/Server/QueueStore/Postgres.hs @@ -24,6 +24,7 @@ module Simplex.Messaging.Server.QueueStore.Postgres batchInsertServices, batchInsertQueues, foldServiceRecs, + foldRcvServiceQueueRecs, foldQueueRecs, handleDuplicate, withLog_, @@ -494,14 +495,6 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where SRecipientService -> "SELECT count(1) FROM msg_queues WHERE rcv_service_id = ? AND deleted_at IS NULL" SNotifierService -> "SELECT count(1) FROM msg_queues WHERE ntf_service_id = ? AND deleted_at IS NULL" - foldRcvServiceQueues :: PostgresQueueStore q -> (RecipientId -> QueueRec -> IO q) -> ServiceId -> (a -> (q, QueueRec) -> IO a) -> a -> IO a - foldRcvServiceQueues st mkQ serviceId f acc = - withConnection (dbStore st) $ \db -> - DB.fold db (queueRecQuery <> " WHERE rcv_service_id = ? AND deleted_at IS NULL") (Only serviceId) acc $ \acc' r -> do - let (rId, qr) = rowToQueueRec r - sq <- mkQ rId qr -- TODO [rcv certs] check map in case queue exists - f acc' (sq, qr) - batchInsertServices :: [STMService] -> PostgresQueueStore q -> IO Int64 batchInsertServices services' toStore = withConnection (dbStore toStore) $ \db -> @@ -547,13 +540,18 @@ insertServiceQuery = VALUES (?,?,?,?,?) |] -foldServiceRecs :: forall a q. Monoid a => PostgresQueueStore q -> (ServiceRec -> IO a) -> IO a +foldServiceRecs :: Monoid a => PostgresQueueStore q -> (ServiceRec -> IO a) -> IO a foldServiceRecs st f = withConnection (dbStore st) $ \db -> DB.fold_ db "SELECT service_id, service_role, service_cert, service_cert_hash, created_at FROM services" mempty $ \ !acc -> fmap (acc <>) . f . rowToServiceRec -foldQueueRecs :: forall a q. Monoid a => Bool -> Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a +foldRcvServiceQueueRecs :: PostgresQueueStore q -> ServiceId -> (a -> (RecipientId, QueueRec) -> IO a) -> a -> IO a +foldRcvServiceQueueRecs st serviceId f acc = + withConnection (dbStore st) $ \db -> + DB.fold db (queueRecQuery <> " WHERE rcv_service_id = ? AND deleted_at IS NULL") (Only serviceId) acc $ \a -> f a . rowToQueueRec + +foldQueueRecs :: Monoid a => Bool -> Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a foldQueueRecs tty withData st skipOld_ f = do (n, r) <- withConnection (dbStore st) $ \db -> foldRecs db (0 :: Int, mempty) $ \(i, acc) qr -> do diff --git a/src/Simplex/Messaging/Server/QueueStore/STM.hs b/src/Simplex/Messaging/Server/QueueStore/STM.hs index 7a70db8d5..c444f33f9 100644 --- a/src/Simplex/Messaging/Server/QueueStore/STM.hs +++ b/src/Simplex/Messaging/Server/QueueStore/STM.hs @@ -17,6 +17,7 @@ module Simplex.Messaging.Server.QueueStore.STM ( STMQueueStore (..), STMService (..), + foldRcvServiceQueues, setStoreLog, withLog', readQueueRecIO, @@ -356,16 +357,15 @@ instance StoreQueueClass q => QueueStoreClass q (STMQueueStore q) where SRecipientService -> serviceRcvQueues SNotifierService -> serviceNtfQueues - -- TODO [rcv certs] this should create queue with mkQ - foldRcvServiceQueues :: STMQueueStore q -> (RecipientId -> QueueRec -> IO q) -> ServiceId -> (a -> (q, QueueRec) -> IO a) -> a -> IO a - foldRcvServiceQueues st mkQ serviceId f acc = - TM.lookupIO serviceId (services st) >>= \case - Nothing -> pure acc - Just s -> - readTVarIO (serviceRcvQueues s) - >>= foldM (\a -> get >=> maybe (pure a) (f a)) acc - where - get rId = TM.lookupIO rId (queues st) $>>= \q -> (q,) <$$> readTVarIO (queueRec q) +foldRcvServiceQueues :: StoreQueueClass q => STMQueueStore q -> ServiceId -> (a -> (q, QueueRec) -> IO a) -> a -> IO a +foldRcvServiceQueues st serviceId f acc = + TM.lookupIO serviceId (services st) >>= \case + Nothing -> pure acc + Just s -> + readTVarIO (serviceRcvQueues s) + >>= foldM (\a -> get >=> maybe (pure a) (f a)) acc + where + get rId = TM.lookupIO rId (queues st) $>>= \q -> (q,) <$$> readTVarIO (queueRec q) withQueueRec :: TVar (Maybe QueueRec) -> (QueueRec -> STM a) -> IO (Either ErrorType a) withQueueRec qr a = atomically $ readQueueRec qr >>= mapM a diff --git a/src/Simplex/Messaging/Server/QueueStore/Types.hs b/src/Simplex/Messaging/Server/QueueStore/Types.hs index 3d9808639..55be4d21d 100644 --- a/src/Simplex/Messaging/Server/QueueStore/Types.hs +++ b/src/Simplex/Messaging/Server/QueueStore/Types.hs @@ -49,7 +49,6 @@ class StoreQueueClass q => QueueStoreClass q s where setQueueService :: (PartyI p, ServiceParty p) => s -> q -> SParty p -> Maybe ServiceId -> IO (Either ErrorType ()) getQueueNtfServices :: s -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)])) getServiceQueueCount :: (PartyI p, ServiceParty p) => s -> SParty p -> ServiceId -> IO (Either ErrorType Int64) - foldRcvServiceQueues :: s -> (RecipientId -> QueueRec -> IO q) -> ServiceId -> (a -> (q, QueueRec) -> IO a) -> a -> IO a data EntityCounts = EntityCounts { queueCount :: Int,