agent: delay connection deletion to finish delivery of pending messages (#1015)

* agent: delay connection deletion to finish delivery of pending messages (wip)

* fixes, test

* notify, test

* add tests

* comment

* add test

* timeout

* test timeout

* up

* more tests

* rename

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>

src/Simplex/Messaging/Agent.hs

@@ -141,7 +141,7 @@ import qualified Data.Text as T
 import Data.Time.Clock
 import Data.Time.Clock.System (systemToUTCTime)
 import Data.Word (Word16)
-import Simplex.FileTransfer.Agent (closeXFTPAgent, deleteSndFileInternal, deleteSndFilesInternal, deleteSndFileRemote, deleteSndFilesRemote, startXFTPWorkers, toFSFilePath, xftpDeleteRcvFile', xftpDeleteRcvFiles', xftpReceiveFile', xftpSendDescription', xftpSendFile')
+import Simplex.FileTransfer.Agent (closeXFTPAgent, deleteSndFileInternal, deleteSndFileRemote, deleteSndFilesInternal, deleteSndFilesRemote, startXFTPWorkers, toFSFilePath, xftpDeleteRcvFile', xftpDeleteRcvFiles', xftpReceiveFile', xftpSendDescription', xftpSendFile')
 import Simplex.FileTransfer.Description (ValidFileDescription)
 import Simplex.FileTransfer.Protocol (FileParty (..))
 import Simplex.FileTransfer.Util (removePath)
@@ -242,12 +242,12 @@ switchConnectionAsync :: AgentErrorMonad m => AgentClient -> ACorrId -> ConnId -
 switchConnectionAsync c = withAgentEnv c .: switchConnectionAsync' c
 
 -- | Delete SMP agent connection (DEL command) asynchronously, no synchronous response
-deleteConnectionAsync :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
-deleteConnectionAsync c = withAgentEnv c . deleteConnectionAsync' c
+deleteConnectionAsync :: AgentErrorMonad m => AgentClient -> Bool -> ConnId -> m ()
+deleteConnectionAsync c waitDelivery = withAgentEnv c . deleteConnectionAsync' c waitDelivery
 
 -- | Delete SMP agent connections using batch commands asynchronously, no synchronous response
-deleteConnectionsAsync :: AgentErrorMonad m => AgentClient -> [ConnId] -> m ()
-deleteConnectionsAsync c = withAgentEnv c . deleteConnectionsAsync' c
+deleteConnectionsAsync :: AgentErrorMonad m => AgentClient -> Bool -> [ConnId] -> m ()
+deleteConnectionsAsync c waitDelivery = withAgentEnv c . deleteConnectionsAsync' c waitDelivery
 
 -- | Create SMP agent connection (NEW command)
 createConnection :: AgentErrorMonad m => AgentClient -> UserId -> Bool -> SConnectionMode c -> Maybe CRClientData -> SubscriptionMode -> m (ConnId, ConnectionRequestUri c)
@@ -541,7 +541,7 @@ createUser' c smp xftp = do
 deleteUser' :: AgentMonad m => AgentClient -> UserId -> Bool -> m ()
 deleteUser' c userId delSMPQueues = do
   if delSMPQueues
-    then withStore c (`setUserDeleted` userId) >>= deleteConnectionsAsync_ delUser c
+    then withStore c (`setUserDeleted` userId) >>= deleteConnectionsAsync_ delUser c False
     else withStore c (`deleteUserRecord` userId)
   atomically $ TM.delete userId $ smpServers c
   where
@@ -613,21 +613,21 @@ ackMessageAsync' c corrId connId msgId rcptInfo_ = do
       (RcvQueue {server}, _) <- withStoreCtx "ackMessageAsync': setMsgUserAck" c $ \db -> setMsgUserAck db connId mId
       enqueueCommand c corrId connId (Just server) . AClientCommand $ APC SAEConn $ ACK msgId rcptInfo_
 
-deleteConnectionAsync' :: forall m. AgentMonad m => AgentClient -> ConnId -> m ()
-deleteConnectionAsync' c connId = deleteConnectionsAsync' c [connId]
+deleteConnectionAsync' :: forall m. AgentMonad m => AgentClient -> Bool -> ConnId -> m ()
+deleteConnectionAsync' c waitDelivery connId = deleteConnectionsAsync' c waitDelivery [connId]
 
-deleteConnectionsAsync' :: AgentMonad m => AgentClient -> [ConnId] -> m ()
+deleteConnectionsAsync' :: AgentMonad m => AgentClient -> Bool -> [ConnId] -> m ()
 deleteConnectionsAsync' = deleteConnectionsAsync_ $ pure ()
 
-deleteConnectionsAsync_ :: forall m. AgentMonad m => m () -> AgentClient -> [ConnId] -> m ()
-deleteConnectionsAsync_ onSuccess c connIds = case connIds of
+deleteConnectionsAsync_ :: forall m. AgentMonad m => m () -> AgentClient -> Bool -> [ConnId] -> m ()
+deleteConnectionsAsync_ onSuccess c waitDelivery connIds = case connIds of
   [] -> onSuccess
   _ -> do
-    (_, rqs, connIds') <- prepareDeleteConnections_ getConns c connIds
-    withStore' c $ forM_ connIds' . setConnDeleted
+    (_, rqs, connIds') <- prepareDeleteConnections_ getConns c waitDelivery connIds
+    withStore' c $ \db -> forM_ connIds' $ setConnDeleted db waitDelivery
     void . forkIO $
       withLock (deleteLock c) "deleteConnectionsAsync" $
-        deleteConnQueues c True rqs >> onSuccess
+        deleteConnQueues c waitDelivery True rqs >> onSuccess
 
 -- | Add connection to the new receive queue
 switchConnectionAsync' :: AgentMonad m => AgentClient -> ACorrId -> ConnId -> m ConnectionStats
@@ -712,7 +712,7 @@ joinConnSrv c userId connId enableNtfs inv@CRInvitationUri {} cInfo subMode srv
       Right _ -> pure connId'
       Left e -> do
         -- possible improvement: recovery for failure on network timeout, see rfcs/2022-04-20-smp-conf-timeout-recovery.md
-        withStore' c (`deleteConn` connId')
+        void $ withStore' c $ \db -> deleteConn db Nothing connId'
         throwError e
 joinConnSrv c userId connId enableNtfs (CRContactUri ConnReqUriData {crAgentVRange, crSmpQueues = (qUri :| _)}) cInfo subMode srv = do
   aVRange <- asks $ smpAgentVRange . config
@@ -1452,19 +1452,23 @@ disableConn c connId = do
 
 -- Unlike deleteConnectionsAsync, this function does not mark connections as deleted in case of deletion failure.
 deleteConnections' :: forall m. AgentMonad m => AgentClient -> [ConnId] -> m (Map ConnId (Either AgentErrorType ()))
-deleteConnections' = deleteConnections_ getConns False
+deleteConnections' = deleteConnections_ getConns False False
 
 deleteDeletedConns :: forall m. AgentMonad m => AgentClient -> [ConnId] -> m (Map ConnId (Either AgentErrorType ()))
-deleteDeletedConns = deleteConnections_ getDeletedConns True
+deleteDeletedConns = deleteConnections_ getDeletedConns True False
+
+deleteDeletedWaitingDeliveryConns :: forall m. AgentMonad m => AgentClient -> [ConnId] -> m (Map ConnId (Either AgentErrorType ()))
+deleteDeletedWaitingDeliveryConns = deleteConnections_ getConns True True
 
 prepareDeleteConnections_ ::
   forall m.
   AgentMonad m =>
   (DB.Connection -> [ConnId] -> IO [Either StoreError SomeConn]) ->
   AgentClient ->
+  Bool ->
   [ConnId] ->
   m (Map ConnId (Either AgentErrorType ()), [RcvQueue], [ConnId])
-prepareDeleteConnections_ getConnections c connIds = do
+prepareDeleteConnections_ getConnections c waitDelivery connIds = do
   conns :: Map ConnId (Either StoreError SomeConn) <- M.fromList . zip connIds <$> withStore' c (`getConnections` connIds)
   let (errs, cs) = M.mapEither id conns
       errs' = M.map (Left . storeError) errs
@@ -1472,19 +1476,27 @@ prepareDeleteConnections_ getConnections c connIds = do
       rqs = concat $ M.elems rcvQs
       connIds' = M.keys rcvQs
   forM_ connIds' $ disableConn c
-  withStore' c $ forM_ (M.keys delRs) . deleteConn
+  -- ! delRs is not used to notify about the result in any of the calling functions,
+  -- ! it is only used to check results count in deleteConnections_;
+  -- ! 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
+  rs' <- catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) (M.keys delRs))
+  forM_ rs' $ \cId -> notify ("", cId, APC SAEConn DEL_CONN)
   pure (errs' <> delRs, rqs, connIds')
   where
     rcvQueues :: SomeConn -> Either (Either AgentErrorType ()) [RcvQueue]
     rcvQueues (SomeConn _ conn) = case connRcvQueues conn of
       [] -> Left $ Right ()
       rqs -> Right rqs
+    notify = atomically . writeTBQueue (subQ c)
 
-deleteConnQueues :: forall m. AgentMonad m => AgentClient -> Bool -> [RcvQueue] -> m (Map ConnId (Either AgentErrorType ()))
-deleteConnQueues c ntf rqs = do
+deleteConnQueues :: forall m. AgentMonad m => AgentClient -> Bool -> Bool -> [RcvQueue] -> m (Map ConnId (Either AgentErrorType ()))
+deleteConnQueues c waitDelivery ntf rqs = do
   rs <- connResults <$> (deleteQueueRecs =<< deleteQueues c rqs)
   let connIds = M.keys $ M.filter isRight rs
-  rs' <- rights <$> withStoreBatch' c (\db -> map (\cId -> deleteConn db cId $> cId) connIds)
+  deliveryTimeout <- if waitDelivery then asks (Just . connDeleteDeliveryTimeout . config) else pure Nothing
+  rs' <- catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) connIds)
   forM_ rs' $ \cId -> notify ("", cId, APC SAEConn DEL_CONN)
   pure rs
   where
@@ -1527,13 +1539,14 @@ deleteConnections_ ::
   AgentMonad m =>
   (DB.Connection -> [ConnId] -> IO [Either StoreError SomeConn]) ->
   Bool ->
+  Bool ->
   AgentClient ->
   [ConnId] ->
   m (Map ConnId (Either AgentErrorType ()))
-deleteConnections_ _ _ _ [] = pure M.empty
-deleteConnections_ getConnections ntf c connIds = do
-  (rs, rqs, _) <- prepareDeleteConnections_ getConnections c connIds
-  rcvRs <- deleteConnQueues c ntf rqs
+deleteConnections_ _ _ _ _ [] = pure M.empty
+deleteConnections_ getConnections ntf waitDelivery c connIds = do
+  (rs, rqs, _) <- prepareDeleteConnections_ getConnections c waitDelivery connIds
+  rcvRs <- deleteConnQueues c waitDelivery ntf rqs
   let rs' = M.union rs rcvRs
   notifyResultError rs'
   pure rs'
@@ -1862,6 +1875,7 @@ cleanupManager c@AgentClient {subQ} = do
     deleteConns =
       withLock (deleteLock c) "cleanupManager" $ do
         void $ withStore' c getDeletedConnIds >>= deleteDeletedConns c
+        void $ withStore' c getDeletedWaitingDeliveryConnIds >>= deleteDeletedWaitingDeliveryConns c
         withStore' c deleteUsersWithoutConns >>= mapM_ (notify "" . DEL_USER)
     deleteRcvFilesExpired = do
       rcvFilesTTL <- asks $ rcvFilesTTL . config

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

@@ -93,6 +93,7 @@ data AgentConfig = AgentConfig
     reconnectInterval :: RetryInterval,
     messageRetryInterval :: RetryInterval2,
     messageTimeout :: NominalDiffTime,
+    connDeleteDeliveryTimeout :: NominalDiffTime,
     helloTimeout :: NominalDiffTime,
     quotaExceededTimeout :: NominalDiffTime,
     initialCleanupDelay :: Int64,
@@ -161,6 +162,7 @@ defaultAgentConfig =
       reconnectInterval = defaultReconnectInterval,
       messageRetryInterval = defaultMessageRetryInterval,
       messageTimeout = 2 * nominalDay,
+      connDeleteDeliveryTimeout = 2 * nominalDay,
       helloTimeout = 2 * nominalDay,
       quotaExceededTimeout = 7 * nominalDay,
       initialCleanupDelay = 30 * 1000000, -- 30 seconds

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

@@ -60,6 +60,7 @@ module Simplex.Messaging.Agent.Store.SQLite
     setConnDeleted,
     setConnAgentVersion,
     getDeletedConnIds,
+    getDeletedWaitingDeliveryConnIds,
     setConnRatchetSync,
     addProcessedRatchetKeyHash,
     checkRatchetKeyHashExists,
@@ -241,7 +242,7 @@ import Data.List (foldl', intercalate, sortBy)
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import qualified Data.Map.Strict as M
-import Data.Maybe (fromMaybe, isJust, listToMaybe, catMaybes)
+import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing, listToMaybe)
 import Data.Ord (Down (..))
 import Data.Text (Text)
 import qualified Data.Text as T
@@ -602,12 +603,32 @@ getRcvConn db ProtocolServer {host, port} rcvId = runExceptT $ do
       DB.query db (rcvQueueQuery <> " WHERE q.host = ? AND q.port = ? AND q.rcv_id = ? AND q.deleted = 0") (host, port, rcvId)
   (rq,) <$> ExceptT (getConn db connId)
 
-deleteConn :: DB.Connection -> ConnId -> IO ()
-deleteConn db connId =
-  DB.executeNamed
-    db
-    "DELETE FROM connections WHERE conn_id = :conn_id;"
-    [":conn_id" := 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
+  Nothing -> delete
+  Just timeout ->
+    ifM
+      checkNoPendingDeliveries_
+      delete
+      ( ifM
+          (checkWaitDeliveryTimeout_ timeout)
+          delete
+          (pure Nothing)
+      )
+  where
+    delete = DB.execute db "DELETE FROM connections WHERE conn_id = ?" (Only connId) $> Just connId
+    checkNoPendingDeliveries_ = do
+      r :: (Maybe Int64) <-
+        maybeFirstRow fromOnly $
+          DB.query db "SELECT 1 FROM snd_message_deliveries WHERE conn_id = ? AND failed = 0 LIMIT 1" (Only connId)
+      pure $ isNothing r
+    checkWaitDeliveryTimeout_ timeout = do
+      cutoffTs <- addUTCTime (-timeout) <$> getCurrentTime
+      r :: (Maybe Int64) <-
+        maybeFirstRow fromOnly $
+          DB.query db "SELECT 1 FROM connections WHERE conn_id = ? AND deleted_at_wait_delivery < ? LIMIT 1" (connId, cutoffTs)
+      pure $ isJust r
 
 upgradeRcvConnToDuplex :: DB.Connection -> ConnId -> NewSndQueue -> IO (Either StoreError SndQueue)
 upgradeRcvConnToDuplex db connId sq =
@@ -1912,8 +1933,13 @@ getConnData db connId' =
     cData (userId, connId, cMode, connAgentVersion, enableNtfs_, lastExternalSndId, deleted, ratchetSyncState) =
       (ConnData {userId, connId, connAgentVersion, enableNtfs = fromMaybe True enableNtfs_, lastExternalSndId, deleted, ratchetSyncState}, cMode)
 
-setConnDeleted :: DB.Connection -> ConnId -> IO ()
-setConnDeleted db connId = DB.execute db "UPDATE connections SET deleted = ? WHERE conn_id = ?" (True, connId)
+setConnDeleted :: DB.Connection -> Bool -> ConnId -> IO ()
+setConnDeleted db waitDelivery connId
+  | waitDelivery = do
+      currentTs <- getCurrentTime
+      DB.execute db "UPDATE connections SET deleted_at_wait_delivery = ? WHERE conn_id = ?" (currentTs, connId)
+  | otherwise =
+      DB.execute db "UPDATE connections SET deleted = ? WHERE conn_id = ?" (True, connId)
 
 setConnAgentVersion :: DB.Connection -> ConnId -> Version -> IO ()
 setConnAgentVersion db connId aVersion =
@@ -1922,6 +1948,10 @@ setConnAgentVersion db connId aVersion =
 getDeletedConnIds :: DB.Connection -> IO [ConnId]
 getDeletedConnIds db = map fromOnly <$> DB.query db "SELECT conn_id FROM connections WHERE deleted = ?" (Only True)
 
+getDeletedWaitingDeliveryConnIds :: DB.Connection -> IO [ConnId]
+getDeletedWaitingDeliveryConnIds db =
+  map fromOnly <$> DB.query_ db "SELECT conn_id FROM connections WHERE deleted_at_wait_delivery IS NOT NULL"
+
 setConnRatchetSync :: DB.Connection -> ConnId -> RatchetSyncState -> IO ()
 setConnRatchetSync db connId ratchetSyncState =
   DB.execute db "UPDATE connections SET ratchet_sync_state = ? WHERE conn_id = ?" (ratchetSyncState, connId)
@@ -2267,17 +2297,18 @@ createRcvFileRedirect db gVar userId redirectFd@FileDescription {chunks = redire
       forM_ (zip [1 ..] replicas) $ \(rno, replica) -> insertRcvFileChunkReplica db rno replica chunkId
   pure dstEntityId
   where
-    dummyDst = FileDescription
-      { party = SFRecipient,
-        size,
-        digest,
-        redirect = Nothing,
-        -- updated later with updateRcvFileRedirect
-        key = C.unsafeSbKey $ B.replicate 32 '#',
-        nonce = C.cbNonce "",
-        chunkSize = FileSize 0,
-        chunks = []
-      }
+    dummyDst =
+      FileDescription
+        { party = SFRecipient,
+          size,
+          digest,
+          redirect = Nothing,
+          -- updated later with updateRcvFileRedirect
+          key = C.unsafeSbKey $ B.replicate 32 '#',
+          nonce = C.cbNonce "",
+          chunkSize = FileSize 0,
+          chunks = []
+        }
 
 insertRcvFile :: DB.Connection -> TVar ChaChaDRG -> UserId -> FileDescription 'FRecipient -> FilePath -> FilePath -> CryptoFile -> Maybe DBRcvFileId -> Maybe RcvFileId -> IO (Either StoreError (RcvFileId, DBRcvFileId))
 insertRcvFile db gVar userId FileDescription {size, digest, key, nonce, chunkSize, redirect} prefixPath tmpPath (CryptoFile savePath cfArgs) redirectId_ redirectEntityId_ = runExceptT $ do
@@ -2346,10 +2377,11 @@ getRcvFile db rcvFileId = runExceptT $ do
         toFile ((rcvFileEntityId, userId, size, digest, key, nonce, chunkSize, prefixPath, tmpPath) :. (savePath, saveKey_, saveNonce_, status, deleted, redirectDbId, redirectEntityId, redirectSize_, redirectDigest_)) =
           let cfArgs = CFArgs <$> saveKey_ <*> saveNonce_
               saveFile = CryptoFile savePath cfArgs
-              redirect = RcvFileRedirect
-                <$> redirectDbId
-                <*> redirectEntityId
-                <*> (RedirectFileInfo <$> redirectSize_ <*> redirectDigest_)
+              redirect =
+                RcvFileRedirect
+                  <$> redirectDbId
+                  <*> redirectEntityId
+                  <*> (RedirectFileInfo <$> redirectSize_ <*> redirectDigest_)
            in RcvFile {rcvFileId, rcvFileEntityId, userId, size, digest, key, nonce, chunkSize, redirect, prefixPath, tmpPath, saveFile, status, deleted, chunks = []}
     getChunks :: RcvFileId -> UserId -> FilePath -> IO [RcvFileChunk]
     getChunks rcvFileEntityId userId fileTmpPath = do

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

@@ -69,6 +69,7 @@ import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20231222_command_created
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20231225_failed_work_items
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240121_message_delivery_indexes
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240124_file_redirect
+import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240223_connections_wait_delivery
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers (dropPrefix, sumTypeJSON)
 import Simplex.Messaging.Transport.Client (TransportHost)
@@ -106,7 +107,8 @@ schemaMigrations =
     ("m20231222_command_created_at", m20231222_command_created_at, Just down_m20231222_command_created_at),
     ("m20231225_failed_work_items", m20231225_failed_work_items, Just down_m20231225_failed_work_items),
     ("m20240121_message_delivery_indexes", m20240121_message_delivery_indexes, Just down_m20240121_message_delivery_indexes),
-    ("m20240124_file_redirect", m20240124_file_redirect, Just down_m20240124_file_redirect)
+    ("m20240124_file_redirect", m20240124_file_redirect, Just down_m20240124_file_redirect),
+    ("m20240223_connections_wait_delivery", m20240223_connections_wait_delivery, Just down_m20240223_connections_wait_delivery)
   ]
 
 -- | The list of migrations in ascending order by date

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

@@ -0,0 +1,18 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+module Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240223_connections_wait_delivery where
+
+import Database.SQLite.Simple (Query)
+import Database.SQLite.Simple.QQ (sql)
+
+m20240223_connections_wait_delivery :: Query
+m20240223_connections_wait_delivery =
+  [sql|
+ALTER TABLE connections ADD COLUMN deleted_at_wait_delivery TEXT;
+|]
+
+down_m20240223_connections_wait_delivery :: Query
+down_m20240223_connections_wait_delivery =
+  [sql|
+ALTER TABLE connections DROP COLUMN deleted_at_wait_delivery;
+|]

src/Simplex/Messaging/Agent/Store/SQLite/Migrations/agent_schema.sql

@@ -26,7 +26,8 @@ CREATE TABLE connections(
   deleted INTEGER DEFAULT 0 CHECK(deleted NOT NULL),
   user_id INTEGER CHECK(user_id NOT NULL)
   REFERENCES users ON DELETE CASCADE,
-  ratchet_sync_state TEXT NOT NULL DEFAULT 'ok'
+  ratchet_sync_state TEXT NOT NULL DEFAULT 'ok',
+  deleted_at_wait_delivery TEXT
 ) WITHOUT ROWID;
 CREATE TABLE rcv_queues(
   host TEXT NOT NULL,