smp: batch queue association updates on subscriptions (#1760)

* smp: batch queue association updates on subscriptions

* refactor to fused batching

* simpler

* batch assoc functions

* clean up

* fix

---------

Co-authored-by: Evgeny @ SimpleX Chat <259188159+evgeny-simplex@users.noreply.github.com>

src/Simplex/Messaging/Server/MsgStore/Journal.hs

@@ -353,6 +353,8 @@ instance QueueStoreClass (JournalQueue s) (QStore s) where
   {-# INLINE getCreateService #-}
   setQueueService = withQS setQueueService
   {-# INLINE setQueueService #-}
+  setQueueServices = withQS setQueueServices
+  {-# INLINE setQueueServices #-}
   getQueueNtfServices = withQS (getQueueNtfServices @(JournalQueue s))
   {-# INLINE getQueueNtfServices #-}
   getServiceQueueCountHash = withQS (getServiceQueueCountHash @(JournalQueue s))

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

@@ -91,7 +91,7 @@ import Simplex.Messaging.SystemTime
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPServiceRole (..))
-import Simplex.Messaging.Util (eitherToMaybe, firstRow, ifM, maybeFirstRow, maybeFirstRow', tshow, (<$$>))
+import Simplex.Messaging.Util (eitherToMaybe, firstRow, ifM, maybeFirstRow, maybeFirstRow', tshow, (<$$>), ($>>=))
 import System.Exit (exitFailure)
 import System.IO (IOMode (..), hFlush, stdout)
 import UnliftIO.STM
@@ -504,6 +504,32 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
         atomically $ writeTVar (queueRec sq) $ Just q'
         withLog "setQueueService" st $ \sl -> logQueueService sl rId party serviceId
 
+  setQueueServices :: (PartyI p, ServiceParty p) => PostgresQueueStore q -> SParty p -> Maybe ServiceId -> [q] -> IO (Either ErrorType (M.Map RecipientId (Either ErrorType ())))
+  setQueueServices _ _ _ [] = pure $ Right M.empty
+  setQueueServices st party serviceId qs = E.uninterruptibleMask_ $ runExceptT $ do
+    updated <- S.fromList <$> withDB' "setQueueServices" st (\db ->
+      map fromOnly <$> DB.query db updateQuery (serviceId, In (map recipientId qs)))
+    results <- liftIO $ forM qs $ \sq -> do
+      let rId = recipientId sq
+      (rId,) <$> if S.member rId updated
+        then readQueueRecIO (queueRec sq) $>>= \q -> do
+          atomically $ writeTVar (queueRec sq) $ Just $ updateRec q
+          withLog "setQueueServices" st $ \sl -> logQueueService sl rId party serviceId
+          pure $ Right ()
+        else pure $ Left AUTH
+    pure $ M.fromList results
+    where
+      updateQuery = case party of
+        SRecipientService ->
+          "UPDATE msg_queues SET rcv_service_id = ? WHERE recipient_id IN ? AND deleted_at IS NULL RETURNING recipient_id"
+        SNotifierService ->
+          "UPDATE msg_queues SET ntf_service_id = ? WHERE recipient_id IN ? AND notifier_id IS NOT NULL AND deleted_at IS NULL RETURNING recipient_id"
+      updateRec q = case party of
+        SRecipientService -> q {rcvServiceId = serviceId}
+        SNotifierService -> case notifier q of
+          Just nc -> q {notifier = Just nc {ntfServiceId = serviceId}}
+          Nothing -> q
+
   getQueueNtfServices :: PostgresQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = E.uninterruptibleMask_ $ runExceptT $ do
     snIds <-

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

@@ -337,6 +337,10 @@ instance StoreQueueClass q => QueueStoreClass q (STMQueueStore q) where
         mapM_ (removeServiceQueue st serviceSel qId) prevSrvId
         mapM_ (addServiceQueue st serviceSel qId) serviceId
 
+  setQueueServices st party serviceId qs = Right . M.fromList <$> mapM setQueue qs
+    where
+      setQueue sq = (recipientId sq,) <$> setQueueService st sq party serviceId
+
   getQueueNtfServices :: STMQueueStore q -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getQueueNtfServices st ntfs = do
     ss <- readTVarIO (services st)

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

@@ -16,6 +16,7 @@ import Control.Concurrent.STM
 import Control.Monad
 import Data.Int (Int64)
 import Data.List.NonEmpty (NonEmpty)
+import Data.Map.Strict (Map)
 import Data.Text (Text)
 import Simplex.Messaging.Protocol
 import Simplex.Messaging.Server.QueueStore
@@ -51,6 +52,7 @@ class StoreQueueClass q => QueueStoreClass q s where
   deleteStoreQueue :: s -> q -> IO (Either ErrorType QueueRec)
   getCreateService :: s -> ServiceRec -> IO (Either ErrorType ServiceId)
   setQueueService :: (PartyI p, ServiceParty p) => s -> q -> SParty p -> Maybe ServiceId -> IO (Either ErrorType ())
+  setQueueServices :: (PartyI p, ServiceParty p) => s -> SParty p -> Maybe ServiceId -> [q] -> IO (Either ErrorType (Map RecipientId (Either ErrorType ())))
   getQueueNtfServices :: s -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
   getServiceQueueCountHash :: (PartyI p, ServiceParty p) => s -> SParty p -> ServiceId -> IO (Either ErrorType (Int64, IdsHash))