agent: read queues in batches for subscriptions (#1758)

* agent: read queues in batches for subscriptions

* resubscribe in batches too

---------

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

Author: epoberezkin

src/Simplex/Messaging/Agent.hs

@@ -1564,12 +1564,11 @@ subscribeAllConnections' :: AgentClient -> Bool -> Maybe UserId -> AM ()
 subscribeAllConnections' c onlyNeeded activeUserId_ = handleErr $ do
   userSrvs <- withStore' c (`getSubscriptionServers` onlyNeeded)
   unless (null userSrvs) $ do
-    maxPending <- asks $ maxPendingSubscriptions . config
-    currPending <- newTVarIO 0
+    batchSize <- asks $ subsBatchSize . config
     let userSrvs' = case activeUserId_ of
           Just activeUserId -> sortOn (\(uId, _) -> if uId == activeUserId then 0 else 1 :: Int) userSrvs
           Nothing -> userSrvs
-    rs <- lift $ mapConcurrently (subscribeUserServer maxPending currPending) userSrvs'
+    rs <- lift $ mapConcurrently (subscribeUserServer batchSize) userSrvs'
     let (errs, oks) = partitionEithers rs
     logInfo $ "subscribed " <> tshow (sum oks) <> " queues"
     forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map ("",)
@@ -1578,18 +1577,16 @@ subscribeAllConnections' c onlyNeeded activeUserId_ = handleErr $ do
   resumeAllCommands c
   where
     handleErr = (`catchAllErrors` \e -> notifySub' c "" (ERR e) >> throwE e)
-    subscribeUserServer :: Int -> TVar Int -> (UserId, SMPServer) -> AM' (Either AgentErrorType Int)
-    subscribeUserServer maxPending currPending (userId, srv) = do
-      atomically $ whenM ((maxPending <=) <$> readTVar currPending) retry
-      tryAllErrors' $ do
-        qs <- withStore' c $ \db -> do
-          qs <- getUserServerRcvQueueSubs db userId srv onlyNeeded
-          unless (null qs) $ atomically $ modifyTVar' currPending (+ length qs) -- update before leaving transaction
-          pure qs
-        let n = length qs
-        unless (null qs) $ lift $ subscribe qs `E.finally` atomically (modifyTVar' currPending $ subtract n)
-        pure n
+    subscribeUserServer :: Int -> (UserId, SMPServer) -> AM' (Either AgentErrorType Int)
+    subscribeUserServer batchSize (userId, srv) = tryAllErrors' $ loop 0 Nothing
       where
+        loop !n cursor_ = do
+          qs <- withStore' c $ \db -> getUserServerRcvQueueSubs db userId srv onlyNeeded batchSize cursor_
+          if null qs then pure n else do
+            lift $ subscribe qs
+            let n' = n + length qs
+                lastRcvId = Just $ queueId $ last qs
+            if length qs < batchSize then pure n' else loop n' lastRcvId
         subscribe qs = do
           rs <- subscribeUserServerQueues c userId srv qs
           -- TODO [certs rcv] storeClientServiceAssocs store associations of queues with client service ID

src/Simplex/Messaging/Agent/Client.hs

@@ -1598,9 +1598,15 @@ checkQueues c = fmap partitionEithers . mapM checkQueue
 -- and that they are already added to pending subscriptions.
 resubscribeSessQueues :: AgentClient -> SMPTransportSession -> [RcvQueueSub] -> AM' ()
 resubscribeSessQueues c tSess qs = do
+  batchSize <- asks $ subsBatchSize . config
   (errs, qs_) <- checkQueues c qs
-  forM_ (L.nonEmpty qs_) $ \qs' -> void $ subscribeSessQueues_ c True (tSess, qs')
+  subscribeChunks $ toChunks batchSize qs_
   forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
+  where
+    subscribeChunks [] = pure ()
+    subscribeChunks (qs' : rest) = do
+      (_, active) <- subscribeSessQueues_ c True (tSess, qs')
+      when active $ subscribeChunks rest
 
 subscribeSessQueues_ :: AgentClient -> Bool -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId), Bool)
 subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c NRMBackground qs

src/Simplex/Messaging/Agent/Env/SQLite.hs

@@ -168,7 +168,7 @@ data AgentConfig = AgentConfig
     ntfBatchSize :: Int,
     ntfSubFirstCheckInterval :: NominalDiffTime,
     ntfSubCheckInterval :: NominalDiffTime,
-    maxPendingSubscriptions :: Int,
+    subsBatchSize :: Int,
     caCertificateFile :: FilePath,
     privateKeyFile :: FilePath,
     certificateFile :: FilePath,
@@ -241,7 +241,7 @@ defaultAgentConfig =
       ntfBatchSize = 150,
       ntfSubFirstCheckInterval = nominalDay,
       ntfSubCheckInterval = 3 * nominalDay,
-      maxPendingSubscriptions = 35000,
+      subsBatchSize = 1350,
       -- CA certificate private key is not needed for initialization
       -- ! we do not generate these
       caCertificateFile = "/etc/opt/simplex-agent/ca.crt",

src/Simplex/Messaging/Agent/Store/AgentStore.hs

@@ -2211,14 +2211,14 @@ getSubscriptionServers db onlyNeeded =
     toUserServer :: (UserId, NonEmpty TransportHost, ServiceName, C.KeyHash) -> (UserId, SMPServer)
     toUserServer (userId, host, port, keyHash) = (userId, SMPServer host port keyHash)
 
-getUserServerRcvQueueSubs :: DB.Connection -> UserId -> SMPServer -> Bool -> IO [RcvQueueSub]
-getUserServerRcvQueueSubs db userId (SMPServer h p kh) onlyNeeded =
-  map toRcvQueueSub
-    <$> DB.query
-      db
-      (rcvQueueSubQuery <> toSubscribe <> " c.deleted = 0 AND q.deleted = 0 AND c.user_id = ? AND q.host = ? AND q.port = ? AND COALESCE(q.server_key_hash, s.key_hash) = ?")
-      (userId, h, p, kh)
+getUserServerRcvQueueSubs :: DB.Connection -> UserId -> SMPServer -> Bool -> Int -> Maybe SMP.RecipientId -> IO [RcvQueueSub]
+getUserServerRcvQueueSubs db userId (SMPServer h p kh) onlyNeeded limit cursor_ =
+  map toRcvQueueSub <$> case cursor_ of
+    Nothing -> DB.query db (q <> orderLimit) (userId, h, p, kh, limit)
+    Just cursor -> DB.query db (q <> " AND q.rcv_id > ? " <> orderLimit) (userId, h, p, kh, cursor, limit)
   where
+    q = rcvQueueSubQuery <> toSubscribe <> " c.deleted = 0 AND q.deleted = 0 AND c.user_id = ? AND q.host = ? AND q.port = ? AND COALESCE(q.server_key_hash, s.key_hash) = ?"
+    orderLimit = " ORDER BY q.rcv_id LIMIT ?"
     toSubscribe
       | onlyNeeded = " WHERE q.to_subscribe = 1 AND "
       | otherwise = " WHERE "