From b9a5b7802cb6089b2b97d37378f6c1ae13b0b863 Mon Sep 17 00:00:00 2001 From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com> Date: Sat, 16 Dec 2023 18:54:07 +0000 Subject: [PATCH] agent: batch sending messages (attempt 4) --- package.yaml | 3 +- simplexmq.cabal | 21 ++--- src/Simplex/Messaging/Agent.hs | 120 +++++++++++++++++++------- src/Simplex/Messaging/Agent/Client.hs | 28 +++++- src/Simplex/Messaging/Agent/Lock.hs | 31 ++++++- src/Simplex/Messaging/Server/Stats.hs | 2 +- 6 files changed, 154 insertions(+), 51 deletions(-) diff --git a/package.yaml b/package.yaml index 83d8a1ab2..4b35d3b2f 100644 --- a/package.yaml +++ b/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.* diff --git a/simplexmq.cabal b/simplexmq.cabal index 7ed3df683..d64055ddf 100644 --- a/simplexmq.cabal +++ b/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 diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index 336d82fbf..f8d485d50 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/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 diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs index 18eb3d642..987bb52c3 100644 --- a/src/Simplex/Messaging/Agent/Client.hs +++ b/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 diff --git a/src/Simplex/Messaging/Agent/Lock.hs b/src/Simplex/Messaging/Agent/Lock.hs index 10062495d..37b63eb0e 100644 --- a/src/Simplex/Messaging/Agent/Lock.hs +++ b/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 diff --git a/src/Simplex/Messaging/Server/Stats.hs b/src/Simplex/Messaging/Server/Stats.hs index 493bd5ac1..38e1d13db 100644 --- a/src/Simplex/Messaging/Server/Stats.hs +++ b/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