agent: lock rows for concurrent queries in PostgreSQL (#1688)

* agent: lock rows for concurrent queries in PostgreSQL

* fix race conditions in workers

* refactor

Author: epoberezkin

src/Simplex/FileTransfer/Agent.hs

@@ -223,6 +223,7 @@ runXFTPRcvWorker c srv Worker {doWork} = do
       agentXFTPDownloadChunk c userId digest replica chunkSpec
       liftIO $ waitUntilForeground c
       (entityId, complete, progress) <- withStore c $ \db -> runExceptT $ do
+        liftIO $ lockRcvFileForUpdate db rcvFileId
         liftIO $ updateRcvFileChunkReceived db (rcvChunkReplicaId replica) rcvChunkId relChunkPath
         RcvFile {size = FileSize currentSize, chunks, redirect} <- ExceptT $ getRcvFile db rcvFileId
         let rcvd = receivedSize chunks
@@ -413,6 +414,7 @@ runXFTPSndPrepareWorker c Worker {doWork} = do
             withStore' c $ \db -> updateSndFileStatus db sndFileId SFSEncrypting
             (digest, chunkSpecsDigests) <- encryptFileForUpload sndFile fsEncPath
             withStore c $ \db -> do
+              lockSndFileForUpdate db sndFileId
               updateSndFileEncrypted db sndFileId digest chunkSpecsDigests
               getSndFile db sndFileId
           else pure sndFile
@@ -530,6 +532,7 @@ runXFTPSndWorker c srv Worker {doWork} = do
       agentXFTPUploadChunk c userId chunkDigest replica' chunkSpec'
       liftIO $ waitUntilForeground c
       sf@SndFile {sndFileEntityId, prefixPath, chunks} <- withStore c $ \db -> do
+        lockSndFileForUpdate db sndFileId
         updateSndChunkReplicaStatus db sndChunkReplicaId SFRSUploaded
         getSndFile db sndFileId
       let uploaded = uploadedSize chunks

src/Simplex/Messaging/Agent.hs

@@ -1145,7 +1145,8 @@ startJoinInvitation c userId connId sq_ enableNtfs cReqUri pqSup =
       let cData = ConnData {userId, connId, connAgentVersion, enableNtfs, lastExternalSndId = 0, deleted = False, ratchetSyncState = RSOk, pqSupport}
       case sq_ of
         Just sq@SndQueue {e2ePubKey = Just _k} -> do
-          e2eSndParams <- withStore c $ \db ->
+          e2eSndParams <- withStore c $ \db -> do
+            lockConnForUpdate db connId
             getSndRatchet db connId v >>= \case
               Right r -> pure $ Right $ snd r
               Left e -> do
@@ -1159,6 +1160,7 @@ startJoinInvitation c userId connId sq_ enableNtfs cReqUri pqSup =
               sndKey_ = snd <$> invLink_
           (q, _) <- lift $ newSndQueue userId "" qInfo sndKey_
           withStore c $ \db -> runExceptT $ do
+            liftIO $ lockConnForUpdate db connId
             e2eSndParams <- createRatchet_ db g maxSupported pqSupport e2eRcvParams
             sq' <- maybe (ExceptT $ updateNewConnSnd db connId q) pure sq_
             pure (cData, sq', e2eSndParams, lnkId_)
@@ -1237,7 +1239,8 @@ joinConnSrv c nm userId connId enableNtfs cReqUri@CRContactUri {} cInfo pqSup su
           AgentConfig {smpClientVRange = vr, smpAgentVRange, e2eEncryptVRange = e2eVR} <- asks config
           let qUri = SMPQueueUri vr $ (rcvSMPQueueAddress rq) {queueMode = Just QMMessaging}
               crData = ConnReqUriData SSSimplex smpAgentVRange [qUri] Nothing
-          e2eRcvParams <- withStore' c $ \db ->
+          e2eRcvParams <- withStore' c $ \db -> do
+            lockConnForUpdate db connId
             getRatchetX3dhKeys db connId >>= \case
               Right keys -> pure $ CR.mkRcvE2ERatchetParams (maxVersion e2eVR) keys
               Left e -> do
@@ -1957,7 +1960,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server,
         withRetryLock2 ri' qLock $ \riState loop -> do
           liftIO $ waitWhileSuspended c
           liftIO $ waitForUserNetwork c
-          resp <- tryError $ case msgType of
+          resp <- tryAllErrors $ case msgType of
             AM_CONN_INFO -> sendConfirmation c NRMBackground sq msgBody
             AM_CONN_INFO_REPLY -> sendConfirmation c NRMBackground sq msgBody
             _ -> case pendingMsgPrepData_ of
@@ -2097,10 +2100,12 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server,
     notifyDelMsgs :: InternalId -> AgentErrorType -> UTCTime -> AM ()
     notifyDelMsgs msgId err expireTs = do
       notifyDel msgId $ MERR (unId msgId) err
-      msgIds_ <- withStore' c $ \db -> getExpiredSndMessages db connId sq expireTs
+      msgIds_ <- withStore' c $ \db -> do
+        msgIds_ <- getExpiredSndMessages db connId sq expireTs
+        forM_ msgIds_ $ \msgId' -> deleteSndMsgDelivery db connId sq msgId' False `catchAll_` pure ()
+        pure msgIds_
       forM_ (L.nonEmpty msgIds_) $ \msgIds -> do
         notify $ MERRS (L.map unId msgIds) err
-        withStore' c $ \db -> forM_ msgIds $ \msgId' -> deleteSndMsgDelivery db connId sq msgId' False `catchAll_` pure ()
       atomically $ incSMPServerStat' c userId server sentExpiredErrs (length msgIds_ + 1)
     delMsg :: InternalId -> AM ()
     delMsg = delMsgKeep False
@@ -3025,7 +3030,8 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), _v, sessId
                               throwE e
                           agentClientMsg :: TVar ChaChaDRG -> ByteString -> AM (Maybe (InternalId, MsgMeta, AMessage, CR.RatchetX448))
                           agentClientMsg g encryptedMsgHash = withStore c $ \db -> runExceptT $ do
-                            rc <- ExceptT $ getRatchet db connId -- ratchet state pre-decryption - required for processing EREADY
+                            liftIO $ lockConnForUpdate db connId
+                            rc <- ExceptT $ getRatchetForUpdate db connId -- ratchet state pre-decryption - required for processing EREADY
                             (agentMsgBody, pqEncryption) <- agentRatchetDecrypt' g db connId rc encAgentMessage
                             liftEither (parse smpP (SEAgentError $ AGENT A_MESSAGE) agentMsgBody) >>= \case
                               agentMsg@(AgentMessage APrivHeader {sndMsgId, prevMsgHash} aMessage) -> do
@@ -3260,6 +3266,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), _v, sessId
                       Just sqs' -> do
                         (sq_@SndQueue {sndPrivateKey}, dhPublicKey) <- lift $ newSndQueue userId connId qInfo Nothing
                         sq2 <- withStore c $ \db -> do
+                          lockConnForUpdate db connId
                           liftIO $ mapM_ (deleteConnSndQueue db connId) delSqs
                           addConnSndQueue db connId (sq_ :: NewSndQueue) {primary = True, dbReplaceQueueId = Just dbQueueId}
                         logServer "<--" c srv rId $ "MSG <QADD>:" <> logSecret' srvMsgId <> " " <> logSecret (senderId queueAddress)
@@ -3564,7 +3571,7 @@ agentRatchetEncrypt db cData msg getPaddedLen pqEnc_ currentE2EVersion = do
 
 agentRatchetEncryptHeader :: DB.Connection -> ConnData -> (VersionSMPA -> PQSupport -> Int) -> Maybe PQEncryption -> CR.VersionE2E -> ExceptT StoreError IO (CR.MsgEncryptKeyX448, Int, PQEncryption)
 agentRatchetEncryptHeader db ConnData {connId, connAgentVersion = v, pqSupport} getPaddedLen pqEnc_ currentE2EVersion = do
-  rc <- ExceptT $ getRatchet db connId
+  rc <- ExceptT $ getRatchetForUpdate db connId
   let paddedLen = getPaddedLen v pqSupport
   (mek, rc') <- withExceptT (SEAgentError . cryptoError) $ CR.rcEncryptHeader rc pqEnc_ currentE2EVersion
   liftIO $ updateRatchet db connId rc' CR.SMDNoChange
@@ -3573,7 +3580,7 @@ agentRatchetEncryptHeader db ConnData {connId, connAgentVersion = v, pqSupport}
 -- encoded EncAgentMessage -> encoded AgentMessage
 agentRatchetDecrypt :: TVar ChaChaDRG -> DB.Connection -> ConnId -> ByteString -> ExceptT StoreError IO (ByteString, PQEncryption)
 agentRatchetDecrypt g db connId encAgentMsg = do
-  rc <- ExceptT $ getRatchet db connId
+  rc <- ExceptT $ getRatchetForUpdate db connId
   agentRatchetDecrypt' g db connId rc encAgentMsg
 
 agentRatchetDecrypt' :: TVar ChaChaDRG -> DB.Connection -> ConnId -> CR.RatchetX448 -> ByteString -> ExceptT StoreError IO (ByteString, PQEncryption)

src/Simplex/Messaging/Agent/Client.hs

@@ -2114,39 +2114,42 @@ withWork :: AgentClient -> TMVar () -> (DB.Connection -> IO (Either StoreError (
 withWork c doWork = withWork_ c doWork . withStore' c
 {-# INLINE withWork #-}
 
+-- setting doWork flag to "no work" before getWork rather than after prevents race condition when flag is set to "has work" by another thread after getWork call.
 withWork_ :: (AnyStoreError e', MonadIO m) => AgentClient -> TMVar () -> ExceptT e m (Either e' (Maybe a)) -> (a -> ExceptT e m ()) -> ExceptT e m ()
 withWork_ c doWork getWork action =
-  getWork >>= \case
-    Right (Just r) -> action r
-    Right Nothing -> noWork
-    -- worker is stopped here (noWork) because the next iteration is likely to produce the same result
+  noWork >> getWork >>= \case
+    Right (Just r) -> hasWork >> action r
+    Right Nothing -> pure ()
     Left e
-      | isWorkItemError e -> noWork >> notifyErr (CRITICAL False) e
-      | otherwise -> notifyErr INTERNAL e
+      | isWorkItemError e -> notifyErr (CRITICAL False) e -- worker remains stopped here because the next iteration is likely to produce the same result
+      | otherwise -> hasWork >> notifyErr INTERNAL e
   where
+    hasWork = atomically $ hasWorkToDo' doWork
     noWork = liftIO $ noWorkToDo doWork
     notifyErr err e = do
       logError $ "withWork_ error: " <> tshow e
       atomically $ writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR $ err $ show e)
 
 withWorkItems :: (AnyStoreError e', MonadIO m) => AgentClient -> TMVar () -> ExceptT e m (Either e' [Either e' a]) -> (NonEmpty a -> ExceptT e m ()) -> ExceptT e m ()
 withWorkItems c doWork getWork action = do
-  getWork >>= \case
-    Right [] -> noWork
+  noWork >> getWork >>= \case
+    Right [] -> pure ()
     Right rs -> do
       let (errs, items) = partitionEithers rs
       case L.nonEmpty items of
-        Just items' -> action items'
+        Just items' -> hasWork >> action items'
         Nothing -> do
-          let criticalErr = find isWorkItemError errs
-          forM_ criticalErr $ \err -> do
-            notifyErr (CRITICAL False) err
-            when (all isWorkItemError errs) noWork
+          case find isWorkItemError errs of
+            Nothing -> hasWork
+            Just err -> do
+              notifyErr (CRITICAL False) err
+              unless (all isWorkItemError errs) hasWork
       forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (\e -> ("", INTERNAL $ show e))
     Left e
-      | isWorkItemError e -> noWork >> notifyErr (CRITICAL False) e
-      | otherwise -> notifyErr INTERNAL e
+      | isWorkItemError e -> notifyErr (CRITICAL False) e
+      | otherwise -> hasWork >> notifyErr INTERNAL e
   where
+    hasWork = atomically $ hasWorkToDo' doWork
     noWork = liftIO $ noWorkToDo doWork
     notifyErr err e = do
       logError $ "withWorkItems error: " <> tshow e