agent: subscribe all connections (#1655)

* agent: subscribe all connections

* query, version

* BoolInt

* add query to errors

* Revert "add query to errors"

This reverts commit 32a1f7f.

* fix optional field

* version

* limit number of in-flight subscriptions to 35000

Author: epoberezkin

src/Simplex/Messaging/Agent.hs

@@ -48,6 +48,7 @@ module Simplex.Messaging.Agent.Client
     newRcvQueue,
     newRcvQueue_,
     subscribeQueues,
+    subscribeUserServerQueues,
     getQueueMessage,
     decryptSMPMessage,
     failSubscription,
@@ -427,7 +428,7 @@ getAgentWorker' toW fromW name hasWork c@AgentClient {agentEnv} key ws work = do
 newWorker :: AgentClient -> STM Worker
 newWorker c = do
   workerId <- stateTVar (workerSeq c) $ \next -> (next, next + 1)
-  doWork <- newTMVar ()
+  doWork <- newTMVar () -- new worker is created  with "some work to do" (indicated by () in TMVar)
   action <- newTMVar Nothing
   restarts <- newTVar $ RestartCount 0 0
   pure Worker {workerId, doWork, action, restarts}
@@ -733,7 +734,7 @@ smpClientDisconnected c@AgentClient {active, smpClients, smpProxiedRelays} tSess
         mode <- getSessionModeIO c
         let resubscribe
               | (mode == TSMEntity) == isJust cId = resubscribeSMPSession c tSess
-              | otherwise = void $ subscribeQueues c qs True
+              | otherwise = void $ subscribeQueues c True qs
         runReaderT resubscribe env
 
 resubscribeSMPSession :: AgentClient -> SMPTransportSession -> AM' ()
@@ -1401,6 +1402,7 @@ newRcvQueue_ c nm userId connId (ProtoServerWithAuth srv auth) vRange cqrd enabl
             shortLink,
             clientService = ClientService DBNewEntity <$> serviceId,
             status = New,
+            enableNtfs,
             dbQueueId = DBNewEntity,
             primary = True,
             dbReplaceQueueId = Nothing,
@@ -1509,30 +1511,51 @@ serverHostError = \case
       _ -> False
 
 -- | Batch by transport session and subscribe queues. The list of results can have a different order.
-subscribeQueues :: AgentClient -> [RcvQueueSub] -> Bool -> AM' [(RcvQueueSub, Either AgentErrorType (Maybe ServiceId))]
-subscribeQueues c qs withEvents = do
+subscribeQueues :: AgentClient -> Bool -> [RcvQueueSub] -> AM' [(RcvQueueSub, Either AgentErrorType (Maybe ServiceId))]
+subscribeQueues c withEvents qs = do
   (errs, qs') <- checkQueues c qs
   atomically $ modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId qs'))
   qss <- batchQueues mkSMPTSession qs' <$> getSessionModeIO c
-  mapM_ addPendingSubs qss
-  rs <- mapConcurrently subscribeQueues_ qss
+  mapM_ (addPendingSubs c) qss
+  rs <- mapConcurrently (subscribeQueues_ c withEvents) qss
   when withEvents $ forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
   pure $ map (second Left) errs <> concatMap L.toList rs
+
+addPendingSubs :: AgentClient -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' ()
+addPendingSubs c (tSess, qs') = atomically $ SS.batchAddPendingSubs tSess (L.toList qs') $ currentSubs c
+
+subscribeQueues_ :: AgentClient -> Bool -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId))
+subscribeQueues_ c withEvents qs'@(tSess@(_, srv, _), _) = do
+  (rs, active) <- subscribeSessQueues_ c withEvents qs'
+  if active
+    then when (hasTempErrors rs) resubscribe $> rs
+    else do
+      logWarn "subcription batch result for replaced SMP client, resubscribing"
+      -- we use BROKER NETWORK error here instead of the original error, so it becomes temporary.
+      resubscribe $> L.map (second $ Left . toNESubscribeError) rs
   where
-    addPendingSubs (tSess, qs') = atomically $ SS.batchAddPendingSubs tSess (L.toList qs') $ currentSubs c
-    subscribeQueues_ qs'@(tSess@(_, srv, _), _) = do
-      (rs, active) <- subscribeSessQueues_ c qs' withEvents
-      if active
-        then when (hasTempErrors rs) resubscribe $> rs
-        else do
-          logWarn "subcription batch result for replaced SMP client, resubscribing"
-          -- we use BROKER NETWORK error here instead of the original error, so it becomes temporary.
-          resubscribe $> L.map (second $ Left . toNESubscribeError) rs
-      where
-        -- treating host errors as temporary here as well
-        hasTempErrors = any (either temporaryOrHostError (const False) . snd)
-        toNESubscribeError = BROKER (B.unpack $ strEncode srv) . NETWORK . NESubscribeError . show
-        resubscribe = resubscribeSMPSession c tSess
+    -- treating host errors as temporary here as well
+    hasTempErrors = any (either temporaryOrHostError (const False) . snd)
+    toNESubscribeError = BROKER (B.unpack $ strEncode srv) . NETWORK . NESubscribeError . show
+    resubscribe = resubscribeSMPSession c tSess
+
+subscribeUserServerQueues :: AgentClient -> UserId -> SMPServer -> [RcvQueueSub] -> AM' [(RcvQueueSub, Either AgentErrorType (Maybe ServiceId))]
+subscribeUserServerQueues c userId srv qs = do
+  mode <- getSessionModeIO c
+  if mode == TSMEntity
+    then subscribeQueues c True qs
+    else do
+      let tSess = (userId, srv, Nothing)
+      (errs, qs_) <- checkQueues c qs
+      forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
+      let errs' = map (second Left) errs
+      case L.nonEmpty qs_ of
+        Just qs' -> do
+          atomically $ modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId $ L.toList qs'))
+          addPendingSubs c (tSess, qs')
+          rs <- subscribeQueues_ c True (tSess, qs')
+          pure $ errs' <> L.toList rs
+        Nothing -> pure errs'
 
 -- only "checked" queues are subscribed
 checkQueues :: AgentClient -> [RcvQueueSub] -> AM' ([(RcvQueueSub, AgentErrorType)], [RcvQueueSub])
@@ -1547,14 +1570,14 @@ checkQueues c = fmap partitionEithers . mapM checkQueue
 resubscribeSessQueues :: AgentClient -> SMPTransportSession -> [RcvQueueSub] -> AM' ()
 resubscribeSessQueues c tSess qs = do
   (errs, qs_) <- checkQueues c qs
-  forM_ (L.nonEmpty qs_) $ \qs' -> void $ subscribeSessQueues_ c (tSess, qs') True
+  forM_ (L.nonEmpty qs_) $ \qs' -> void $ subscribeSessQueues_ c True (tSess, qs')
   forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
 
-subscribeSessQueues_ :: AgentClient -> (SMPTransportSession, NonEmpty RcvQueueSub) -> Bool -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId), Bool)
-subscribeSessQueues_ c qs withEvents = sendClientBatch_ "SUB" False subscribeQueues_ c NRMBackground qs
+subscribeSessQueues_ :: AgentClient -> Bool -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId), Bool)
+subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c NRMBackground qs
   where
-    subscribeQueues_ :: SMPClient -> NonEmpty RcvQueueSub -> IO (BatchResponses RcvQueueSub SMPClientError (Maybe ServiceId), Bool)
-    subscribeQueues_ smp qs' = do
+    subscribe_ :: SMPClient -> NonEmpty RcvQueueSub -> IO (BatchResponses RcvQueueSub SMPClientError (Maybe ServiceId), Bool)
+    subscribe_ smp qs' = do
       let (userId, srv, _) = tSess
       atomically $ incSMPServerStat' c userId srv connSubAttempts $ length qs'
       rs <- sendBatch (\smp' _ -> subscribeSMPQueues smp') smp NRMBackground qs'
@@ -1646,12 +1669,6 @@ getRemovedSubs AgentClient {removedSubs} k = TM.lookup k removedSubs >>= maybe n
       TM.insert k s removedSubs
       pure s
 
-addPendingSubscription :: AgentClient -> RcvQueueSub -> STM ()
-addPendingSubscription c rq = do
-  modifyTVar' (subscrConns c) $ S.insert $ qConnId rq
-  tSess <- mkSMPTransportSession c rq
-  SS.addPendingSub tSess rq $ currentSubs c
-
 addNewQueueSubscription :: AgentClient -> RcvQueue -> SMPTransportSession -> SessionId -> AM' ()
 addNewQueueSubscription c rq' tSess sessId = do
   let rq = rcvQueueSub rq'

src/Simplex/Messaging/Agent/Client.hs

@@ -166,6 +166,7 @@ data AgentConfig = AgentConfig
     ntfBatchSize :: Int,
     ntfSubFirstCheckInterval :: NominalDiffTime,
     ntfSubCheckInterval :: NominalDiffTime,
+    maxPendingSubscriptions :: Int,
     caCertificateFile :: FilePath,
     privateKeyFile :: FilePath,
     certificateFile :: FilePath,
@@ -237,6 +238,7 @@ defaultAgentConfig =
       ntfBatchSize = 150,
       ntfSubFirstCheckInterval = nominalDay,
       ntfSubCheckInterval = 3 * nominalDay,
+      maxPendingSubscriptions = 35000,
       -- 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/Env/SQLite.hs

@@ -88,6 +88,8 @@ data StoredRcvQueue (q :: DBStored) = RcvQueue
     clientService :: Maybe (StoredClientService q),
     -- | queue status
     status :: QueueStatus,
+    -- | to enable notifications for this queue - this field is duplicated from ConnData
+    enableNtfs :: Bool,
     -- | database queue ID (within connection)
     dbQueueId :: DBEntityId' q,
     -- | True for a primary or a next primary queue of the connection (next if dbReplaceQueueId is set)
@@ -110,15 +112,16 @@ data RcvQueueSub = RcvQueueSub
     rcvId :: SMP.RecipientId,
     rcvPrivateKey :: RcvPrivateAuthKey,
     status :: QueueStatus,
+    enableNtfs :: Bool,
     dbQueueId :: Int64,
     primary :: Bool,
     dbReplaceQueueId :: Maybe Int64
   }
   deriving (Show)
 
 rcvQueueSub :: RcvQueue -> RcvQueueSub
-rcvQueueSub RcvQueue {userId, connId, server, rcvId, rcvPrivateKey, status, dbQueueId = DBEntityId dbQueueId, primary, dbReplaceQueueId} =
-  RcvQueueSub {userId, connId, server, rcvId, rcvPrivateKey, status, dbQueueId, primary, dbReplaceQueueId}
+rcvQueueSub RcvQueue {userId, connId, server, rcvId, rcvPrivateKey, status, enableNtfs, dbQueueId = DBEntityId dbQueueId, primary, dbReplaceQueueId} =
+  RcvQueueSub {userId, connId, server, rcvId, rcvPrivateKey, status, enableNtfs, dbQueueId, primary, dbReplaceQueueId}
 
 data ShortLinkCreds = ShortLinkCreds
   { shortLinkId :: SMP.LinkId,

src/Simplex/Messaging/Agent/Store.hs

@@ -39,6 +39,8 @@ module Simplex.Messaging.Agent.Store.AgentStore
     updateNewConnRcv,
     updateNewConnSnd,
     createSndConn,
+    getSubscriptionServers,
+    getUserServerRcvQueueSubs,
     getConn,
     getDeletedConn,
     getConns,
@@ -110,6 +112,7 @@ module Simplex.Messaging.Agent.Store.AgentStore
     updateSndMsgRcpt,
     getPendingQueueMsg,
     getConnectionsForDelivery,
+    getAllSndQueuesForDelivery,
     updatePendingMsgRIState,
     deletePendingMsgs,
     getExpiredSndMessages,
@@ -137,6 +140,7 @@ module Simplex.Messaging.Agent.Store.AgentStore
     -- Async commands
     createCommand,
     getPendingCommandServers,
+    getAllPendingCommandConns,
     getPendingServerCommand,
     updateCommandServer,
     deleteCommand,
@@ -884,8 +888,8 @@ createSndMsg db connId sndMsgData@SndMsgData {internalSndId, internalHash} = do
   insertSndMsgDetails_ db connId sndMsgData
   updateSndMsgHash db connId internalSndId internalHash
 
-createSndMsgDelivery :: DB.Connection -> ConnId -> SndQueue -> InternalId -> IO ()
-createSndMsgDelivery db connId SndQueue {dbQueueId} msgId =
+createSndMsgDelivery :: DB.Connection -> SndQueue -> InternalId -> IO ()
+createSndMsgDelivery db SndQueue {connId, dbQueueId} msgId =
   DB.execute db "INSERT INTO snd_message_deliveries (conn_id, snd_queue_id, internal_id) VALUES (?, ?, ?)" (connId, dbQueueId, msgId)
 
 getSndMsgViaRcpt :: DB.Connection -> ConnId -> InternalSndId -> IO (Either StoreError SndMsg)
@@ -917,6 +921,15 @@ getConnectionsForDelivery :: DB.Connection -> IO [ConnId]
 getConnectionsForDelivery db =
   map fromOnly <$> DB.query_ db "SELECT DISTINCT conn_id FROM snd_message_deliveries WHERE failed = 0"
 
+getAllSndQueuesForDelivery :: DB.Connection -> IO [SndQueue]
+getAllSndQueuesForDelivery db = map toSndQueue <$> DB.query_ db (sndQueueQuery <> " " <> delivery)
+  where
+    delivery = [sql|
+      JOIN (SELECT DISTINCT conn_id, snd_queue_id FROM snd_message_deliveries WHERE failed = 0) d
+      ON d.conn_id = q.conn_id AND d.snd_queue_id = q.snd_queue_id
+      WHERE c.deleted = 0
+    |]
+
 getPendingQueueMsg :: DB.Connection -> ConnId -> SndQueue -> IO (Either StoreError (Maybe (Maybe RcvQueue, PendingMsgData)))
 getPendingQueueMsg db connId SndQueue {dbQueueId} =
   getWorkItem "message" getMsgId getMsgData markMsgFailed
@@ -1319,6 +1332,21 @@ getPendingCommandServers db connIds =
     smpServer (host, port, keyHash) = SMPServer <$> host <*> port <*> keyHash
     conns = S.fromList connIds
 
+getAllPendingCommandConns :: DB.Connection -> IO [(ConnId, Maybe SMPServer)]
+getAllPendingCommandConns db =
+  map toResult
+    <$> DB.query_
+      db
+      [sql|
+        SELECT DISTINCT c.conn_id, c.host, c.port, COALESCE(c.server_key_hash, s.key_hash)
+        FROM commands c
+        JOIN connections cs ON c.conn_id = cs.conn_id
+        LEFT JOIN servers s ON s.host = c.host AND s.port = c.port
+        WHERE cs.deleted = 0
+      |]
+  where
+    toResult (connId, host, port, keyHash) = (connId, SMPServer <$> host <*> port <*> keyHash)
+
 getPendingServerCommand :: DB.Connection -> ConnId -> Maybe SMPServer -> IO (Either StoreError (Maybe PendingCommand))
 getPendingServerCommand db connId srv_ = getWorkItem "command" getCmdId getCommand markCommandFailed
   where
@@ -2023,6 +2051,30 @@ newQueueId_ :: [Only Int64] -> DBEntityId
 newQueueId_ [] = DBEntityId 1
 newQueueId_ (Only maxId : _) = DBEntityId (maxId + 1)
 
+-- * subscribe all connections
+
+getSubscriptionServers :: DB.Connection -> IO [(UserId, SMPServer)]
+getSubscriptionServers db =
+  map toUserServer
+    <$> DB.query_
+      db
+      [sql|
+        SELECT DISTINCT c.user_id, q.host, q.port, COALESCE(q.server_key_hash, s.key_hash)
+        FROM rcv_queues q
+        JOIN servers s ON q.host = s.host AND q.port = s.port
+        JOIN connections c ON q.conn_id = c.conn_id
+        WHERE c.deleted = 0 AND q.deleted = 0
+      |]
+  where
+    toUserServer :: (UserId, NonEmpty TransportHost, ServiceName, C.KeyHash) -> (UserId, SMPServer)
+    toUserServer (userId, host, port, keyHash) = (userId, SMPServer host port keyHash)
+
+getUserServerRcvQueueSubs :: DB.Connection -> UserId -> SMPServer -> IO [RcvQueueSub]
+getUserServerRcvQueueSubs db userId srv =
+  map toRcvQueueSub <$> DB.query db (rcvQueueSubQuery <> condition) (userId, host srv, port srv)
+  where
+    condition = " WHERE c.deleted = 0 AND q.deleted = 0 AND c.user_id = ? AND q.host = ? AND q.port = ?"
+
 -- * getConn helpers
 
 getConn :: DB.Connection -> ConnId -> IO (Either StoreError SomeConn)
@@ -2229,7 +2281,7 @@ rcvQueueQuery :: Query
 rcvQueueQuery =
   [sql|
     SELECT c.user_id, COALESCE(q.server_key_hash, s.key_hash), q.conn_id, q.host, q.port, q.rcv_id, q.rcv_private_key, q.rcv_dh_secret,
-      q.e2e_priv_key, q.e2e_dh_secret, q.snd_id, q.queue_mode, q.status,
+      q.e2e_priv_key, q.e2e_dh_secret, q.snd_id, q.queue_mode, q.status, c.enable_ntfs,
       q.rcv_queue_id, q.rcv_primary, q.replace_rcv_queue_id, q.switch_status, q.smp_client_version, q.delete_errors,
       q.ntf_public_key, q.ntf_private_key, q.ntf_id, q.rcv_ntf_dh_secret,
       q.link_id, q.link_key, q.link_priv_sig_key, q.link_enc_fixed_data
@@ -2240,13 +2292,13 @@ rcvQueueQuery =
 
 toRcvQueue ::
   (UserId, C.KeyHash, ConnId, NonEmpty TransportHost, ServiceName, SMP.RecipientId, SMP.RcvPrivateAuthKey, SMP.RcvDhSecret, C.PrivateKeyX25519, Maybe C.DhSecretX25519, SMP.SenderId, Maybe QueueMode)
-    :. (QueueStatus, DBEntityId, BoolInt, Maybe Int64, Maybe RcvSwitchStatus, Maybe VersionSMPC, Int)
+    :. (QueueStatus, Maybe BoolInt, DBEntityId, BoolInt, Maybe Int64, Maybe RcvSwitchStatus, Maybe VersionSMPC, Int)
     :. (Maybe SMP.NtfPublicAuthKey, Maybe SMP.NtfPrivateAuthKey, Maybe SMP.NotifierId, Maybe RcvNtfDhSecret)
     :. (Maybe SMP.LinkId, Maybe LinkKey, Maybe C.PrivateKeyEd25519, Maybe EncDataBytes) ->
   RcvQueue
 toRcvQueue
   ( (userId, keyHash, connId, host, port, rcvId, rcvPrivateKey, rcvDhSecret, e2ePrivKey, e2eDhSecret, sndId, queueMode)
-      :. (status, dbQueueId, BI primary, dbReplaceQueueId, rcvSwchStatus, smpClientVersion_, deleteErrors)
+      :. (status, enableNtfs_, dbQueueId, BI primary, dbReplaceQueueId, rcvSwchStatus, smpClientVersion_, deleteErrors)
       :. (ntfPublicKey_, ntfPrivateKey_, notifierId_, rcvNtfDhSecret_)
       :. (shortLinkId_, shortLinkKey_, linkPrivSigKey_, linkEncFixedData_)
   ) =
@@ -2258,8 +2310,9 @@ toRcvQueue
       shortLink = case (shortLinkId_, shortLinkKey_, linkPrivSigKey_, linkEncFixedData_) of
         (Just shortLinkId, Just shortLinkKey, Just linkPrivSigKey, Just linkEncFixedData) -> Just ShortLinkCreds {shortLinkId, shortLinkKey, linkPrivSigKey, linkEncFixedData}
         _ -> Nothing
+      enableNtfs = maybe True unBI enableNtfs_
       -- TODO [certs rcv] read client service
-   in RcvQueue {userId, connId, server, rcvId, rcvPrivateKey, rcvDhSecret, e2ePrivKey, e2eDhSecret, sndId, queueMode, shortLink, clientService = Nothing, status, dbQueueId, primary, dbReplaceQueueId, rcvSwchStatus, smpClientVersion, clientNtfCreds, deleteErrors}
+   in RcvQueue {userId, connId, server, rcvId, rcvPrivateKey, rcvDhSecret, e2ePrivKey, e2eDhSecret, sndId, queueMode, shortLink, clientService = Nothing, status, enableNtfs, dbQueueId, primary, dbReplaceQueueId, rcvSwchStatus, smpClientVersion, clientNtfCreds, deleteErrors}
 
 -- | returns all connection queue credentials, the first queue is the primary one
 getRcvQueueSubsByConnId_ :: DB.Connection -> ConnId -> IO (Maybe (NonEmpty RcvQueueSub))
@@ -2270,16 +2323,17 @@ getRcvQueueSubsByConnId_ db connId =
 rcvQueueSubQuery :: Query
 rcvQueueSubQuery =
   [sql|
-    SELECT c.user_id, q.conn_id, q.host, q.port, COALESCE(q.server_key_hash, s.key_hash), q.rcv_id, q.rcv_private_key, q.status,
+    SELECT c.user_id, q.conn_id, q.host, q.port, COALESCE(q.server_key_hash, s.key_hash), q.rcv_id, q.rcv_private_key, q.status, c.enable_ntfs,
       q.rcv_queue_id, q.rcv_primary, q.replace_rcv_queue_id
     FROM rcv_queues q
     JOIN servers s ON q.host = s.host AND q.port = s.port
     JOIN connections c ON q.conn_id = c.conn_id
   |]
 
-toRcvQueueSub :: (UserId, ConnId, NonEmpty TransportHost, ServiceName, C.KeyHash, SMP.RecipientId, SMP.RcvPrivateAuthKey, QueueStatus, Int64, BoolInt, Maybe Int64) -> RcvQueueSub
-toRcvQueueSub (userId, connId, host, port, keyHash, rcvId, rcvPrivateKey, status, dbQueueId, BI primary, dbReplaceQueueId) =
-  RcvQueueSub {userId, connId, server = SMPServer host port keyHash, rcvId, dbQueueId, primary, dbReplaceQueueId, rcvPrivateKey, status}
+toRcvQueueSub :: (UserId, ConnId, NonEmpty TransportHost, ServiceName, C.KeyHash, SMP.RecipientId, SMP.RcvPrivateAuthKey, QueueStatus, Maybe BoolInt, Int64, BoolInt, Maybe Int64) -> RcvQueueSub
+toRcvQueueSub (userId, connId, host, port, keyHash, rcvId, rcvPrivateKey, status, enableNtfs_, dbQueueId, BI primary, dbReplaceQueueId) =
+  let enableNtfs = maybe True unBI enableNtfs_
+   in RcvQueueSub {userId, connId, server = SMPServer host port keyHash, rcvId, rcvPrivateKey, status, enableNtfs, dbQueueId, primary, dbReplaceQueueId}
 
 getRcvQueueById :: DB.Connection -> ConnId -> Int64 -> IO (Either StoreError RcvQueue)
 getRcvQueueById db connId dbRcvId =