try tracking service subs in the agent (WIP, does not compile)

Author: epoberezkin

src/Simplex/Messaging/Agent.hs

@@ -850,7 +850,7 @@ newConn c nm userId enableNtfs cMode userData_ clientData pqInitKeys subMode = d
   srv <- getSMPServer c userId
   connId <- newConnNoQueues c userId enableNtfs cMode (CR.connPQEncryption pqInitKeys)
   (connId,) <$> newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKeys subMode srv
-    `catchE` \e -> withStore' c (`deleteConnRecord` connId) >> throwE e
+    `catchE` \e -> withStore' c (\db -> deleteConnRecord db userId connId) >> throwE e
 
 setConnShortLink' :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserLinkData -> Maybe CRClientData -> AM (ConnShortLink c)
 setConnShortLink' c nm connId cMode userData clientData =
@@ -2160,7 +2160,8 @@ prepareDeleteConnections_ getConnections c waitDelivery connIds = do
   -- ! if it was used to notify about the result, it might be necessary to differentiate
   -- ! between completed deletions of connections, and deletions delayed due to wait for delivery (see deleteConn)
   deliveryTimeout <- if waitDelivery then asks (Just . connDeleteDeliveryTimeout . config) else pure Nothing
-  cIds_ <- lift $ L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) (M.keys delRs))
+  let delConns = map (\(connId, RcvQueue {userId}) -> (userId, connId)) $ M.toList delRs
+  cIds_ <- lift $ L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (uncurry $ deleteConn db deliveryTimeout) delConns)
   forM_ cIds_ $ \cIds -> notify ("", "", AEvt SAEConn $ DEL_CONNS cIds)
   pure (errs' <> delRs, rqs, connIds')
   where
@@ -2179,9 +2180,9 @@ prepareDeleteConnections_ getConnections c waitDelivery connIds = do
 deleteConnQueues :: AgentClient -> NetworkRequestMode -> Bool -> Bool -> [RcvQueue] -> AM' (Map ConnId (Either AgentErrorType ()))
 deleteConnQueues c nm waitDelivery ntf rqs = do
   rs <- connResults <$> (deleteQueueRecs =<< deleteQueues c nm rqs)
-  let connIds = M.keys $ M.filter isRight rs
+  let conns = map (\(connId, RcvQueue {userId}) -> (userId, connId)) $ M.toList $ M.filter isRight rs
   deliveryTimeout <- if waitDelivery then asks (Just . connDeleteDeliveryTimeout . config) else pure Nothing
-  cIds_ <- L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) connIds)
+  cIds_ <- L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (uncurry $ deleteConn db deliveryTimeout) conns)
   forM_ cIds_ $ \cIds -> notify ("", "", AEvt SAEConn $ DEL_CONNS cIds)
   pure rs
   where
@@ -2474,7 +2475,7 @@ sendNtfConnCommands :: AgentClient -> NtfSupervisorCommand -> AM ()
 sendNtfConnCommands c cmd = do
   ns <- asks ntfSupervisor
   connIds <- liftIO $ S.toList <$> getSubscriptions c
-  rs <- lift $ withStoreBatch' c (\db -> map (getConnData db) connIds)
+  rs <- lift $ withStoreBatch' c (\db -> map (getConnData False db) connIds)
   let (connIds', cErrs) = enabledNtfConns (zip connIds rs)
   forM_ (L.nonEmpty connIds') $ \connIds'' ->
     atomically $ writeTBQueue (ntfSubQ ns) (cmd, connIds'')

src/Simplex/Messaging/Agent/Client.hs

@@ -271,7 +271,7 @@ import Simplex.Messaging.Protocol
     RcvNtfPublicDhKey,
     SMPMsgMeta (..),
     SProtocolType (..),
-    ServiceSub,
+    ServiceSub (..),
     SndPublicAuthKey,
     SubscriptionMode (..),
     NewNtfCreds (..),
@@ -599,10 +599,9 @@ getServiceCredentials c userId srv =
   liftIO (TM.lookupIO userId $ useClientServices c)
     $>>= \useService -> if useService then Just <$> getService else pure Nothing
   where
-    getService :: AM (ServiceCredentials, Maybe ServiceId)
     getService = do
       let g = agentDRG c
-      ((C.KeyHash kh, serviceCreds), serviceId_) <-
+      ((C.KeyHash kh, serviceCreds), serviceSub_) <-
         withStore' c $ \db ->
           getClientService db userId srv >>= \case
             Just service -> pure service
@@ -614,7 +613,7 @@ getServiceCredentials c userId srv =
       (_, pk) <- atomically $ C.generateKeyPair g
       let serviceSignKey = C.APrivateSignKey C.SEd25519 pk
           creds = ServiceCredentials {serviceRole = SRMessaging, serviceCreds, serviceCertHash = XV.Fingerprint kh, serviceSignKey}
-      pure (creds, serviceId_)
+      pure (creds, smpServiceId <$> serviceSub_)
 
 class (Encoding err, Show err) => ProtocolServerClient v err msg | msg -> v, msg -> err where
   type Client msg = c | c -> msg

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

@@ -395,19 +395,36 @@ createClientService db userId srv (kh, (cert, pk)) =
     |]
     (userId, host srv, port srv, kh, cert, pk)
 
-getClientService :: DB.Connection -> UserId -> SMPServer -> IO (Maybe ((C.KeyHash, TLS.Credential), Maybe ServiceId))
+getClientService :: DB.Connection -> UserId -> SMPServer -> IO (Maybe ((C.KeyHash, TLS.Credential), Maybe ServiceSub))
 getClientService db userId srv =
   maybeFirstRow toService $
     DB.query
       db
       [sql|
-        SELECT service_cert_hash, service_cert, service_priv_key, rcv_service_id
+        SELECT service_cert_hash, service_cert, service_priv_key, rcv_service_id,
         FROM client_services
         WHERE user_id = ? AND host = ? AND port = ?
       |]
       (userId, host srv, port srv)
   where
-    toService (kh, cert, pk, serviceId_) = ((kh, (cert, pk)), serviceId_)
+    toService (kh, cert, pk, serviceId_, n, idsHash) =
+      let service_ = (\serviceId -> ServiceSub serviceId n idsHash) <$> serviceId_
+       in ((kh, (cert, pk)), service_)
+
+updateServiceAggregates :: DB.Connection -> UserId -> SMPServer -> Int64 -> [RecipientId] -> IO ()
+updateServiceAggregates db userId server change rcvIds =
+  getClientService db userId server >>= \case
+    Just (_, Just (ServiceSub serviceId n idsHash)) ->
+      DB.execute
+        db
+        [sql|
+          UPDATE client_services
+          SET service_queue_count = ?,
+              service_queue_ids_hash = ?
+          WHERE user_id = ? AND rcv_service_id = ?
+        |]
+        (n + change, idsHash <> queueIdsHash rcvIds, userId, serviceId)
+    _ -> pure ()
 
 getClientServiceServers :: DB.Connection -> UserId -> IO [(SMPServer, ServiceSub)]
 getClientServiceServers db userId =
@@ -499,8 +516,13 @@ createConnRecord db connId ConnData {userId, connAgentVersion, enableNtfs, pqSup
     |]
     (userId, connId, cMode, connAgentVersion, BI enableNtfs, pqSupport, BI True)
 
-deleteConnRecord :: DB.Connection -> ConnId -> IO ()
-deleteConnRecord db connId = DB.execute db "DELETE FROM connections WHERE conn_id = ?" (Only connId)
+deleteConnRecord :: DB.Connection -> UserId -> ConnId -> IO ()
+deleteConnRecord db userId connId = do
+  -- TODO [certs rcv] it needs to be grouped per server here
+  rIds :: [RecipientId] <-
+    map fromOnly <$> DB.query db "SELECT rcv_id FROM rcv_queues WHERE conn_id = ? AND deleted = 0 AND rcv_service_assoc = 1" (Only connId)
+  unless (null rIds) $ updateServiceAggregates db userId server (- length rIds') rIds'
+  DB.execute db "DELETE FROM connections WHERE conn_id = ?" (Only connId)
 
 checkConfirmedSndQueueExists_ :: DB.Connection -> NewSndQueue -> IO Bool
 checkConfirmedSndQueueExists_ db SndQueue {server, sndId} = do
@@ -521,8 +543,8 @@ getRcvConn db ProtocolServer {host, port} rcvId = runExceptT $ do
   (rq,) <$> ExceptT (getConn db connId)
 
 -- | Deletes connection, optionally checking for pending snd message deliveries; returns connection id if it was deleted
-deleteConn :: DB.Connection -> Maybe NominalDiffTime -> ConnId -> IO (Maybe ConnId)
-deleteConn db waitDeliveryTimeout_ connId = case waitDeliveryTimeout_ of
+deleteConn :: DB.Connection -> Maybe NominalDiffTime -> UserId -> ConnId -> IO (Maybe ConnId)
+deleteConn db waitDeliveryTimeout_ userId connId = case waitDeliveryTimeout_ of
   Nothing -> delete
   Just timeout ->
     ifM
@@ -534,7 +556,7 @@ deleteConn db waitDeliveryTimeout_ connId = case waitDeliveryTimeout_ of
           (pure Nothing)
       )
   where
-    delete = deleteConnRecord db connId $> Just connId
+    delete = deleteConnRecord db userId connId $> Just connId
     checkNoPendingDeliveries_ = do
       r :: (Maybe Int64) <-
         maybeFirstRow fromOnly $
@@ -609,15 +631,13 @@ setRcvSwitchStatus db rq@RcvQueue {rcvId, server = ProtocolServer {host, port}}
   pure rq {rcvSwchStatus}
 
 setRcvQueueDeleted :: DB.Connection -> RcvQueue -> IO ()
-setRcvQueueDeleted db RcvQueue {rcvId, server = ProtocolServer {host, port}} = do
-  DB.execute
-    db
-    [sql|
-      UPDATE rcv_queues
-      SET deleted = 1
-      WHERE host = ? AND port = ? AND rcv_id = ?
-    |]
-    (host, port, rcvId)
+setRcvQueueDeleted db RcvQueue {connId, dbQueueId, rcvId} = do
+  q_ :: Maybe BoolInt <-
+    maybeFirstRow fromOnly $
+      DB.query db "SELECT rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ? AND deleted = 0" (connId, dbQueueId)
+  forM_ q_ $ \(BI rcvServiceAssoc) -> do
+    DB.execute db "UPDATE rcv_queues SET deleted = 1 WHERE host = ? AND port = ? AND rcv_id = ?" (host, port, rcvId)
+    when rcvServiceAssoc $ updateServiceAggregates db userId server (-1) [rcvId]
 
 setRcvQueueConfirmedE2E :: DB.Connection -> RcvQueue -> C.DhSecretX25519 -> VersionSMPC -> IO ()
 setRcvQueueConfirmedE2E db RcvQueue {rcvId, server = ProtocolServer {host, port}} e2eDhSecret smpClientVersion =
@@ -677,8 +697,13 @@ incRcvDeleteErrors db RcvQueue {connId, dbQueueId} =
   DB.execute db "UPDATE rcv_queues SET delete_errors = delete_errors + 1 WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId)
 
 deleteConnRcvQueue :: DB.Connection -> RcvQueue -> IO ()
-deleteConnRcvQueue db RcvQueue {connId, dbQueueId} =
-  DB.execute db "DELETE FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId)
+deleteConnRcvQueue db RcvQueue {connId, dbQueueId, server} = do
+  q_ :: Maybe (BoolInt, BoolInt) <-
+    maybeFirstRow id $
+      DB.query db "SELECT deleted, rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId)
+  forM_ q_ $ \(BI deleted, BI rcvServiceAssoc) -> do
+    DB.execute db "DELETE FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId)
+    when (not deleted && rcvServiceAssoc) $ updateServiceAggregates db userId server (-1) [rcvId]
 
 deleteConnSndQueue :: DB.Connection -> ConnId -> SndQueue -> IO ()
 deleteConnSndQueue db connId SndQueue {dbQueueId} = do
@@ -687,7 +712,7 @@ deleteConnSndQueue db connId SndQueue {dbQueueId} = do
 
 getPrimaryRcvQueue :: DB.Connection -> ConnId -> IO (Either StoreError RcvQueue)
 getPrimaryRcvQueue db connId =
-  maybe (Left SEConnNotFound) (Right . L.head) <$> getRcvQueuesByConnId_ db connId
+  maybe (Left SEConnNotFound) (Right . L.head) <$> getRcvQueuesByConnId_ db connId False
 
 getRcvQueue :: DB.Connection -> ConnId -> SMPServer -> SMP.RecipientId -> IO (Either StoreError RcvQueue)
 getRcvQueue db connId (SMPServer host port _) rcvId =
@@ -1015,7 +1040,7 @@ getPendingQueueMsg db connId SndQueue {dbQueueId} =
     getMsgData :: InternalId -> IO (Either StoreError (Maybe RcvQueue, PendingMsgData))
     getMsgData msgId = runExceptT $ do
       msg <- ExceptT $ firstRow' pendingMsgData err getMsgData_
-      rq_ <- liftIO $ L.head <$$> getRcvQueuesByConnId_ db connId
+      rq_ <- liftIO $ L.head <$$> getRcvQueuesByConnId_ db connId False
       pure (rq_, msg)
       where
         getMsgData_ =
@@ -2061,6 +2086,7 @@ insertRcvQueue_ db connId' rq@RcvQueue {..} serverKeyHash_ = do
         :. (shortLinkId <$> shortLink, shortLinkKey <$> shortLink, linkPrivSigKey <$> shortLink, linkEncFixedData <$> shortLink)
         :. ntfCredsFields
     )
+  when rcvServiceAssoc $ updateServiceAggregates db userId server 1 [rcvId]
   pure (rq :: NewRcvQueue) {connId = connId', dbQueueId = qId}
   where
     ntfCredsFields = case clientNtfCreds of
@@ -2119,21 +2145,19 @@ getDeletedConn = getAnyConn True
 {-# INLINE getDeletedConn #-}
 
 getAnyConn :: Bool -> DB.Connection -> ConnId -> IO (Either StoreError SomeConn)
-getAnyConn deleted' dbConn connId =
-  getConnData dbConn connId >>= \case
+getAnyConn deleted dbConn connId =
+  getConnData deleted dbConn connId >>= \case
     Nothing -> pure $ Left SEConnNotFound
-    Just (cData@ConnData {deleted}, cMode)
-      | deleted /= deleted' -> pure $ Left SEConnNotFound
-      | otherwise -> do
-          rQ <- getRcvQueuesByConnId_ dbConn connId
-          sQ <- getSndQueuesByConnId_ dbConn connId
-          pure $ case (rQ, sQ, cMode) of
-            (Just rqs, Just sqs, CMInvitation) -> Right $ SomeConn SCDuplex (DuplexConnection cData rqs sqs)
-            (Just (rq :| _), Nothing, CMInvitation) -> Right $ SomeConn SCRcv (RcvConnection cData rq)
-            (Nothing, Just (sq :| _), CMInvitation) -> Right $ SomeConn SCSnd (SndConnection cData sq)
-            (Just (rq :| _), Nothing, CMContact) -> Right $ SomeConn SCContact (ContactConnection cData rq)
-            (Nothing, Nothing, _) -> Right $ SomeConn SCNew (NewConnection cData)
-            _ -> Left SEConnNotFound
+    Just (cData, cMode) -> do
+      rQ <- getRcvQueuesByConnId_ dbConn connId deleted
+      sQ <- getSndQueuesByConnId_ dbConn connId
+      pure $ case (rQ, sQ, cMode) of
+        (Just rqs, Just sqs, CMInvitation) -> Right $ SomeConn SCDuplex (DuplexConnection cData rqs sqs)
+        (Just (rq :| _), Nothing, CMInvitation) -> Right $ SomeConn SCRcv (RcvConnection cData rq)
+        (Nothing, Just (sq :| _), CMInvitation) -> Right $ SomeConn SCSnd (SndConnection cData sq)
+        (Just (rq :| _), Nothing, CMContact) -> Right $ SomeConn SCContact (ContactConnection cData rq)
+        (Nothing, Nothing, _) -> Right $ SomeConn SCNew (NewConnection cData)
+        _ -> Left SEConnNotFound
 
 getConns :: DB.Connection -> [ConnId] -> IO [Either StoreError SomeConn]
 getConns = getAnyConns_ False
@@ -2149,8 +2173,8 @@ getAnyConns_ deleted' db connIds = forM connIds $ E.handle handleDBError . getAn
     handleDBError :: E.SomeException -> IO (Either StoreError SomeConn)
     handleDBError = pure . Left . SEInternal . bshow
 
-getConnData :: DB.Connection -> ConnId -> IO (Maybe (ConnData, ConnectionMode))
-getConnData db connId' =
+getConnData :: Bool -> DB.Connection -> ConnId -> IO (Maybe (ConnData, ConnectionMode))
+getConnData deleted db connId' =
   maybeFirstRow cData $
     DB.query
       db
@@ -2159,9 +2183,9 @@ getConnData db connId' =
           user_id, conn_id, conn_mode, smp_agent_version, enable_ntfs,
           last_external_snd_msg_id, deleted, ratchet_sync_state, pq_support
         FROM connections
-        WHERE conn_id = ?
+        WHERE conn_id = ? AND deleted = ?
       |]
-      (Only connId')
+      (connId', BI deleted)
   where
     cData (userId, connId, cMode, connAgentVersion, enableNtfs_, lastExternalSndId, BI deleted, ratchetSyncState, pqSupport) =
       (ConnData {userId, connId, connAgentVersion, enableNtfs = maybe True unBI enableNtfs_, lastExternalSndId, deleted, ratchetSyncState, pqSupport}, cMode)
@@ -2171,8 +2195,21 @@ setConnDeleted db waitDelivery connId
   | waitDelivery = do
       currentTs <- getCurrentTime
       DB.execute db "UPDATE connections SET deleted_at_wait_delivery = ? WHERE conn_id = ?" (currentTs, connId)
-  | otherwise =
+  | otherwise = do
       DB.execute db "UPDATE connections SET deleted = ? WHERE conn_id = ?" (BI True, connId)
+      -- TODO [certs rcv] it needs to be grouped per server
+      qs :: [(Int64, RecipientId, BoolInt)] <-
+        DB.query db "SELECT rcv_queue_id, rcv_id, rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND deleted = 0" (Only connId)
+      unless (null qs) $ do
+#if defined(dbPostgres)
+        let dbQIds = In (map (\(dbQId, _, _) -> dbQId) qs)
+        DB.execute db "UPDATE rcv_queues SET deleted = 1 WHERE conn_id = ? AND rcv_queue_id IN ?" (connId, dbQIds)
+#else
+        let ids = map (\(dbQId, _, _) -> (connId, dbQId)) qs
+        DB.executeMany db "UPDATE rcv_queues SET deleted = 1 WHERE conn_id = ? AND rcv_queue_id = ?" ids
+#endif
+        let rIds' = mapMaybe (\case (_, rId, BI True) -> Just rId; _ -> Nothing) qs
+        unless (null rIds') $ updateServiceAggregates db userId server (- length rIds') rIds'
 
 setConnUserId :: DB.Connection -> UserId -> ConnId -> UserId -> IO ()
 setConnUserId db oldUserId connId newUserId =
@@ -2218,10 +2255,10 @@ deleteRatchetKeyHashesExpired db ttl = do
   DB.execute db "DELETE FROM processed_ratchet_key_hashes WHERE created_at < ?" (Only cutoffTs)
 
 -- | returns all connection queues, the first queue is the primary one
-getRcvQueuesByConnId_ :: DB.Connection -> ConnId -> IO (Maybe (NonEmpty RcvQueue))
-getRcvQueuesByConnId_ db connId =
+getRcvQueuesByConnId_ :: DB.Connection -> ConnId -> Bool -> IO (Maybe (NonEmpty RcvQueue))
+getRcvQueuesByConnId_ db connId deleted =
   L.nonEmpty . sortBy primaryFirst . map toRcvQueue
-    <$> DB.query db (rcvQueueQuery <> " WHERE q.conn_id = ? AND q.deleted = 0") (Only connId)
+    <$> DB.query db (rcvQueueQuery <> " WHERE q.conn_id = ? AND q.deleted = ?") (connId, deleted)
   where
     primaryFirst RcvQueue {primary = p, dbReplaceQueueId = i} RcvQueue {primary = p', dbReplaceQueueId = i'} =
       -- the current primary queue is ordered first, the next primary - second

src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250815_service_certs.hs

@@ -17,7 +17,7 @@ CREATE TABLE client_services(
   service_cert BYTEA NOT NULL,
   service_cert_hash BYTEA NOT NULL,
   service_priv_key BYTEA NOT NULL,
-  service_id BYTEA,
+  rcv_service_id BYTEA,
   service_queue_count BIGINT NOT NULL DEFAULT 0,
   service_queue_ids_hash BYTEA NOT NULL DEFAULT '\x00000000000000000000000000000000',
   FOREIGN KEY(host, port) REFERENCES servers ON UPDATE CASCADE ON DELETE RESTRICT

src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250815_service_certs.hs

@@ -15,7 +15,7 @@ CREATE TABLE client_services(
   service_cert BLOB NOT NULL,
   service_cert_hash BLOB NOT NULL,
   service_priv_key BLOB NOT NULL,
-  service_id BLOB,
+  rcv_service_id BLOB,
   service_queue_count INTEGER NOT NULL DEFAULT 0,
   service_queue_ids_hash BLOB NOT NULL DEFAULT x'00000000000000000000000000000000',
   FOREIGN KEY(host, port) REFERENCES servers ON UPDATE CASCADE ON DELETE RESTRICT

src/Simplex/Messaging/Protocol.hs

@@ -1465,7 +1465,7 @@ type MsgId = ByteString
 type MsgBody = ByteString
 
 data ServiceSub = ServiceSub
-  { serviceId :: ServiceId,
+  { smpServiceId :: ServiceId,
     smpQueueCount :: Int64,
     smpQueueIdsHash :: IdsHash
   }