agent: batch sending messages (attempt 4)

package.yaml

@@ -65,8 +65,7 @@ dependencies:
   - sqlcipher-simple == 0.4.*
   - stm == 2.5.*
   - temporary == 1.3.*
-  - time == 1.9.*
-  - time-compat == 1.9.*
+  - time == 1.12.*
   - time-manager == 0.0.*
   - tls >= 1.7.0 && < 1.8
   - transformers == 0.6.*

simplexmq.cabal

@@ -212,8 +212,7 @@ library
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -285,8 +284,7 @@ executable ntf-server
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -358,8 +356,7 @@ executable smp-agent
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -431,8 +428,7 @@ executable smp-server
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -504,8 +500,7 @@ executable xftp
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -577,8 +572,7 @@ executable xftp-server
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , tls >=1.7.0 && <1.8
     , transformers ==0.6.*
@@ -687,8 +681,7 @@ test-suite simplexmq-test
     , sqlcipher-simple ==0.4.*
     , stm ==2.5.*
     , temporary ==1.3.*
-    , time ==1.9.*
-    , time-compat ==1.9.*
+    , time ==1.12.*
     , time-manager ==0.0.*
     , timeit ==2.0.*
     , tls >=1.7.0 && <1.8

src/Simplex/Messaging/Agent.hs

@@ -64,6 +64,7 @@ module Simplex.Messaging.Agent
     resubscribeConnection,
     resubscribeConnections,
     sendMessage,
+    sendMessages,
     ackMessage,
     switchConnection,
     abortConnectionSwitch,
@@ -166,6 +167,7 @@ import Simplex.RemoteControl.Invitation
 import Simplex.RemoteControl.Types
 import UnliftIO.Async (async, race_)
 import UnliftIO.Concurrent (forkFinally, forkIO, threadDelay)
+import UnliftIO.IORef
 import UnliftIO.STM
 
 -- import GHC.Conc (unsafeIOToSTM)
@@ -277,6 +279,12 @@ resubscribeConnections c = withAgentEnv c . resubscribeConnections' c
 sendMessage :: AgentErrorMonad m => AgentClient -> ConnId -> MsgFlags -> MsgBody -> m AgentMsgId
 sendMessage c = withAgentEnv c .:. sendMessage' c
 
+type MsgReq = (ConnId, MsgFlags, MsgBody)
+
+-- | Send multiple messages to different connections (SEND command)
+sendMessages :: AgentErrorMonad m => AgentClient -> [MsgReq] -> m [Either AgentErrorType AgentMsgId]
+sendMessages c = withAgentEnv c . sendMessages' c
+
 ackMessage :: AgentErrorMonad m => AgentClient -> ConnId -> AgentMsgId -> Maybe MsgReceiptInfo -> m ()
 ackMessage c = withAgentEnv c .:. ackMessage' c
 
@@ -865,19 +873,38 @@ getNotificationMessage' c nonce encNtfInfo = do
           Just SMP.NMsgMeta {msgId = msgId', msgTs = msgTs'} -> msgId == msgId' || msgTs > msgTs'
           Nothing -> SMP.notification msgFlags
 
+type EIORef a = IORef (Either AgentErrorType a)
+
 -- | Send message to the connection (SEND command) in Reader monad
 sendMessage' :: forall m. AgentMonad m => AgentClient -> ConnId -> MsgFlags -> MsgBody -> m AgentMsgId
-sendMessage' c connId msgFlags msg = withConnLock c connId "sendMessage" $ do
-  SomeConn _ conn <- withStore c (`getConn` connId)
-  case conn of
-    DuplexConnection cData _ sqs -> enqueueMsgs cData sqs
-    SndConnection cData sq -> enqueueMsgs cData [sq]
-    _ -> throwError $ CONN SIMPLEX
+sendMessage' c connId msgFlags msg =
+  oneResult $ \r -> sendMessagesB c [(r, (connId, msgFlags, msg))]
+
+-- | Send multiple messages to different connections (SEND command) in Reader monad
+sendMessages' :: forall m. AgentMonad m => AgentClient -> [MsgReq] -> m [Either AgentErrorType AgentMsgId]
+sendMessages' c msgReqs = do
+  rs <- replicateM (length msgReqs) (newIORef $ Left $ INTERNAL "skipped in batch")
+  sendMessagesB c $ zip rs msgReqs
+  mapM readIORef rs
+
+sendMessagesB :: forall m. AgentMonad m => AgentClient -> [(EIORef AgentMsgId, MsgReq)] -> m ()
+sendMessagesB c reqs = withConnLocks c connIds "sendMessages" $ do
+  reqs' <- zip reqs <$> withStoreBatch c (\db -> map (getConn db) connIds)
+  reqs'' <- catMaybes <$> mapM prepareConn reqs'
+  enqueueMessagesB c reqs''
   where
-    enqueueMsgs :: ConnData -> NonEmpty SndQueue -> m AgentMsgId
-    enqueueMsgs cData sqs = do
-      when (ratchetSyncSendProhibited cData) $ throwError $ CMD PROHIBITED
-      enqueueMessages c cData sqs msgFlags $ A_MSG msg
+    prepareConn :: ((EIORef AgentMsgId, MsgReq), Either AgentErrorType SomeConn) -> m (Maybe (EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage)))
+    prepareConn (req@(r, _), conn_) = case conn_ of
+      Left e -> Nothing <$ writeIORef r (Left e)
+      Right (SomeConn _ conn) -> case conn of
+        DuplexConnection cData _ sqs -> enqueueMsgs cData sqs req
+        SndConnection cData sq -> enqueueMsgs cData [sq] req
+        _ -> Nothing <$ writeIORef r (Left $ CONN SIMPLEX)
+    enqueueMsgs :: ConnData -> NonEmpty SndQueue -> (EIORef AgentMsgId, MsgReq) -> m (Maybe (EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage)))
+    enqueueMsgs cData sqs (r, (_, msgFlags, msg))
+      | ratchetSyncSendProhibited cData = Nothing <$ writeIORef r (Left $ CMD PROHIBITED)
+      | otherwise = pure $ Just (r, (cData, sqs, msgFlags, A_MSG msg))
+    connIds = map (\(_, (connId, _, _)) -> connId) reqs
 
 -- / async command processing v v v
 
@@ -1056,22 +1083,32 @@ enqueueMessages c cData sqs msgFlags aMessage = do
   enqueueMessages' c cData sqs msgFlags aMessage
 
 enqueueMessages' :: AgentMonad m => AgentClient -> ConnData -> NonEmpty SndQueue -> MsgFlags -> AMessage -> m AgentMsgId
-enqueueMessages' c cData (sq :| sqs) msgFlags aMessage = do
-  msgId <- enqueueMessage c cData sq msgFlags aMessage
-  mapM_ (enqueueSavedMessage c cData msgId) $
-    filter (\SndQueue {status} -> status == Secured || status == Active) sqs
-  pure msgId
+enqueueMessages' c cData sqs msgFlags aMessage =
+  oneResult $ \r -> enqueueMessagesB c [(r, (cData, sqs, msgFlags, aMessage))]
+
+enqueueMessagesB :: AgentMonad m => AgentClient -> [(EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage))] -> m ()
+enqueueMessagesB _ [] = pure ()
+enqueueMessagesB c reqs = enqueueMessageB c reqs >>= enqueueSavedMessageB c
+
+isActiveSndQ :: SndQueue -> Bool
+isActiveSndQ SndQueue {status} = status == Secured || status == Active
 
 enqueueMessage :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> MsgFlags -> AMessage -> m AgentMsgId
-enqueueMessage c cData@ConnData {connId} sq msgFlags aMessage = do
-  resumeMsgDelivery c cData sq
+enqueueMessage c cData sq msgFlags aMessage =
+  oneResult $ \r -> enqueueMessageB c [(r, (cData, [sq], msgFlags, aMessage))]
+
+-- this function is used only for sending messages in batch, it returns the list of successes to enqueue additional deliveries
+enqueueMessageB :: forall m. AgentMonad m => AgentClient -> [(EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage))] -> m [(ConnData, [SndQueue], AgentMsgId)]
+enqueueMessageB c reqs = do
+  forM_ reqs $ \(_, (cData, sq :| _, _, _)) ->
+    resumeMsgDelivery c cData sq
   aVRange <- asks $ smpAgentVRange . config
-  msgId <- storeSentMsg $ maxVersion aVRange
-  queuePendingMsgs c sq [msgId]
-  pure $ unId msgId
+  mIds <- withStoreBatch c $ \db ->
+    map (storeSentMsg db $ maxVersion aVRange) reqs
+  catMaybes <$> mapM processResults (zip reqs mIds)
   where
-    storeSentMsg :: Version -> m InternalId
-    storeSentMsg agentVersion = withStore c $ \db -> runExceptT $ do
+    storeSentMsg :: DB.Connection -> Version -> (EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage)) -> IO (Either StoreError InternalId)
+    storeSentMsg db agentVersion (_, (ConnData {connId}, sq :| _, msgFlags, aMessage)) = runExceptT $ do
       internalTs <- liftIO getCurrentTime
       (internalId, internalSndId, prevMsgHash) <- liftIO $ updateSndIds db connId
       let privHeader = APrivHeader (unSndId internalSndId) prevMsgHash
@@ -1085,13 +1122,39 @@ enqueueMessage c cData@ConnData {connId} sq msgFlags aMessage = do
       liftIO $ createSndMsg db connId msgData
       liftIO $ createSndMsgDelivery db connId sq internalId
       pure internalId
+    processResults :: ((EIORef AgentMsgId, (ConnData, NonEmpty SndQueue, MsgFlags, AMessage)), Either AgentErrorType InternalId) -> m (Maybe (ConnData, [SndQueue], AgentMsgId))
+    processResults ((r, (cData, sq :| sqs, _, _)), mId_) = case mId_ of
+      Left e -> Nothing <$ writeIORef r (Left e)
+      Right mId -> do
+        let InternalId msgId = mId
+        writeIORef r $ Right msgId
+        queuePendingMsgs c sq [mId]
+        let sqs' = filter isActiveSndQ sqs
+        pure $ if null sqs' then Nothing else Just (cData, sqs', msgId)
 
 enqueueSavedMessage :: AgentMonad m => AgentClient -> ConnData -> AgentMsgId -> SndQueue -> m ()
-enqueueSavedMessage c cData@ConnData {connId} msgId sq = do
-  resumeMsgDelivery c cData sq
-  let mId = InternalId msgId
-  queuePendingMsgs c sq [mId]
-  withStore' c $ \db -> createSndMsgDelivery db connId sq mId
+enqueueSavedMessage c cData msgId sq = enqueueSavedMessageB c [(cData, [sq], msgId)]
+
+enqueueSavedMessageB :: AgentMonad m => AgentClient -> [(ConnData, [SndQueue], AgentMsgId)] -> m ()
+enqueueSavedMessageB c reqs = do
+  -- saving to the database moved to the start to avoid race conditions when delivery is read from queue before it is saved
+  void $ withStoreBatch' c $ \db -> concatMap (storeDeliveries db) reqs
+  forM_ reqs $ \(cData, sqs, msgId) ->
+    forM sqs $ \sq -> do
+      resumeMsgDelivery c cData sq
+      let mId = InternalId msgId
+      queuePendingMsgs c sq [mId]
+  where
+    storeDeliveries :: DB.Connection -> (ConnData, [SndQueue], AgentMsgId) -> [IO ()]
+    storeDeliveries db (ConnData {connId}, sqs, msgId) = do
+      let mId = InternalId msgId
+       in map (\sq -> createSndMsgDelivery db connId sq mId) sqs
+
+oneResult :: AgentMonad m => (EIORef a -> m b) -> m a
+oneResult action = do
+  r <- newIORef $ Left $ INTERNAL "skipped in batch of one"
+  _ <- action r
+  readIORef r >>= liftEither
 
 resumeMsgDelivery :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> m ()
 resumeMsgDelivery c cData@ConnData {connId} sq@SndQueue {server, sndId} = do
@@ -2434,8 +2497,7 @@ storeConfirmation c ConnData {connId, connAgentVersion} sq e2eEncryption_ agentM
 enqueueRatchetKeyMsgs :: forall m. AgentMonad m => AgentClient -> ConnData -> NonEmpty SndQueue -> CR.E2ERatchetParams 'C.X448 -> m AgentMsgId
 enqueueRatchetKeyMsgs c cData (sq :| sqs) e2eEncryption = do
   msgId <- enqueueRatchetKey c cData sq e2eEncryption
-  mapM_ (enqueueSavedMessage c cData msgId) $
-    filter (\SndQueue {status} -> status == Secured || status == Active) sqs
+  mapM_ (enqueueSavedMessage c cData msgId) $ filter isActiveSndQ sqs
   pure msgId
 
 enqueueRatchetKey :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> CR.E2ERatchetParams 'C.X448 -> m AgentMsgId

src/Simplex/Messaging/Agent/Client.hs

@@ -24,6 +24,7 @@ module Simplex.Messaging.Agent.Client
     ProtocolTestStep (..),
     newAgentClient,
     withConnLock,
+    withConnLocks,
     withInvLock,
     closeAgentClient,
     closeProtocolServerClients,
@@ -99,6 +100,8 @@ module Simplex.Messaging.Agent.Client
     withStore',
     withStoreCtx,
     withStoreCtx',
+    withStoreBatch,
+    withStoreBatch',
     storeError,
     userServers,
     pickServer,
@@ -658,8 +661,17 @@ withConnLock AgentClient {connLocks} connId name = withLockMap_ connLocks connId
 withInvLock :: MonadUnliftIO m => AgentClient -> ByteString -> String -> m a -> m a
 withInvLock AgentClient {invLocks} = withLockMap_ invLocks
 
+withConnLocks :: MonadUnliftIO m => AgentClient -> [ConnId] -> String -> m a -> m a
+withConnLocks AgentClient {connLocks} = withLocksMap_ connLocks . filter (not . B.null)
+
 withLockMap_ :: (Ord k, MonadUnliftIO m) => TMap k Lock -> k -> String -> m a -> m a
-withLockMap_ locks key = withGetLock $ TM.lookup key locks >>= maybe newLock pure
+withLockMap_ = withGetLock . getMapLock
+
+withLocksMap_ :: (Ord k, MonadUnliftIO m) => TMap k Lock -> [k] -> String -> m a -> m a
+withLocksMap_ = withGetLocks . getMapLock
+
+getMapLock :: Ord k => TMap k Lock -> k -> STM Lock
+getMapLock locks key = TM.lookup key locks >>= maybe newLock pure
   where
     newLock = createLock >>= \l -> TM.insert key l locks $> l
 
@@ -1291,6 +1303,20 @@ withStoreCtx_ ctx_ c action = do
     handleInternal :: String -> E.SomeException -> IO (Either StoreError a)
     handleInternal ctxStr e = pure . Left . SEInternal . B.pack $ show e <> ctxStr
 
+withStoreBatch :: AgentMonad' m => AgentClient -> (DB.Connection -> [IO (Either StoreError a)]) -> m [Either AgentErrorType a]
+withStoreBatch c actions = do
+  st <- asks store
+  rs <-
+    liftIO $ agentOperationBracket c AODatabase (\_ -> pure ()) $
+      withTransaction st $ mapM (`E.catch` handleInternal) . actions
+  pure $ map (first storeError) rs
+  where
+    handleInternal :: E.SomeException -> IO (Either StoreError a)
+    handleInternal = pure . Left . SEInternal . B.pack . show
+
+withStoreBatch' :: AgentMonad' m => AgentClient -> (DB.Connection -> [IO a]) -> m [Either AgentErrorType a]
+withStoreBatch' c actions = withStoreBatch c $ map (Right <$>) . actions
+
 storeError :: StoreError -> AgentErrorType
 storeError = \case
   SEConnNotFound -> CONN NOT_FOUND

src/Simplex/Messaging/Agent/Lock.hs

@@ -1,8 +1,18 @@
-module Simplex.Messaging.Agent.Lock where
+{-# LANGUAGE NamedFieldPuns #-}
+
+module Simplex.Messaging.Agent.Lock
+  ( Lock,
+    createLock,
+    withLock,
+    withGetLock,
+    withGetLocks,
+  )
+where
 
 import Control.Monad (void)
 import Control.Monad.IO.Unlift
 import Data.Functor (($>))
+import UnliftIO.Async (forConcurrently)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
 
@@ -18,9 +28,22 @@ withLock lock name =
     (atomically $ putTMVar lock name)
     (void . atomically $ takeTMVar lock)
 
-withGetLock :: MonadUnliftIO m => STM Lock -> String -> m a -> m a
-withGetLock getLock name a =
+withGetLock :: MonadUnliftIO m => (k -> STM Lock) -> k -> String -> m a -> m a
+withGetLock getLock key name a =
   E.bracket
-    (atomically $ getLock >>= \l -> putTMVar l name $> l)
+    (atomically $ getPutLock getLock key name)
     (atomically . takeTMVar)
     (const a)
+
+withGetLocks :: MonadUnliftIO m => (k -> STM Lock) -> [k] -> String -> m a -> m a
+withGetLocks getLock keys name = E.bracket holdLocks releaseLocks . const
+  where
+    holdLocks = forConcurrently keys $ \key -> atomically $ getPutLock getLock key name
+    -- only this withGetLocks would be holding the locks,
+    -- so it's safe to combine all lock releases into one transaction
+    releaseLocks = atomically . mapM_ takeTMVar
+
+-- getLock and putTMVar can be in one transaction on the assumption that getLock doesn't write in case the lock already exists,
+-- and in case it is created and added to some shared resource (we use TMap) it also helps avoid contention for the newly created lock.
+getPutLock :: (k -> STM Lock) -> k -> String -> STM Lock
+getPutLock getLock key name = getLock key >>= \l -> putTMVar l name $> l

src/Simplex/Messaging/Server/Stats.hs

@@ -11,7 +11,7 @@ import qualified Data.Attoparsec.ByteString.Char8 as A
 import qualified Data.ByteString.Char8 as B
 import Data.Set (Set)
 import qualified Data.Set as S
-import Data.Time.Calendar.Month.Compat (pattern MonthDay)
+import Data.Time.Calendar.Month (pattern MonthDay)
 import Data.Time.Calendar.OrdinalDate (mondayStartWeek)
 import Data.Time.Clock (UTCTime (..))
 import Simplex.Messaging.Encoding.String