client: batch while sending commands - wait for responses before sending the next batch (#825)

* client: batch while sending commands - wait for responses before sending the next batch

* fix comments

* fix tests

simplexmq.cabal

@@ -534,6 +534,7 @@ test-suite simplexmq-test
       AgentTests.SchemaDump
       AgentTests.SQLiteTests
       CLITests
+      CoreTests.BatchingTests
       CoreTests.CryptoTests
       CoreTests.EncodingTests
       CoreTests.ProtocolErrorTests

src/Simplex/Messaging/Client.hs

@@ -11,7 +11,6 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE TupleSections #-}
 
 -- |
 -- Module      : Simplex.Messaging.Client
@@ -69,6 +68,13 @@ module Simplex.Messaging.Client
     temporaryClientError,
     ServerTransmission,
     ClientCommand,
+
+    -- * For testing
+    ClientBatch (..),
+    PCTransmission,
+    batchClientTransmissions,
+    mkTransmission,
+    clientStub,
   )
 where
 
@@ -82,12 +88,10 @@ import Data.Aeson (FromJSON (..), ToJSON (..))
 import qualified Data.Aeson as J
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
-import Data.Either (rights)
-import Data.Foldable (foldl')
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (find)
-import Data.List.NonEmpty (NonEmpty (..), (<|))
+import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import Data.Maybe (fromMaybe)
 import Data.Time.Clock (UTCTime, getCurrentTime)
@@ -95,6 +99,7 @@ import GHC.Generics (Generic)
 import Network.Socket (ServiceName)
 import Numeric.Natural
 import qualified Simplex.Messaging.Crypto as C
+import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers (dropPrefix, enumJSON)
 import Simplex.Messaging.Protocol as SMP
@@ -131,11 +136,43 @@ data PClient err msg = PClient
     pingErrorCount :: TVar Int,
     clientCorrId :: TVar Natural,
     sentCommands :: TMap CorrId (Request err msg),
-    sndQ :: TBQueue (NonEmpty SentRawTransmission),
+    sndQ :: TBQueue ByteString,
     rcvQ :: TBQueue (NonEmpty (SignedTransmission err msg)),
     msgQ :: Maybe (TBQueue (ServerTransmission msg))
   }
 
+clientStub :: ByteString -> STM (ProtocolClient err msg)
+clientStub sessionId = do
+  connected <- newTVar False
+  clientCorrId <- newTVar 0
+  sentCommands <- TM.empty
+  sndQ <- newTBQueue 100
+  rcvQ <- newTBQueue 100
+  return
+    ProtocolClient
+      { action = Nothing,
+        sessionId,
+        sessionTs = undefined,
+        thVersion = 5,
+        timeoutPerBlock = undefined,
+        blockSize = smpBlockSize,
+        batch = undefined,
+        client_ =
+          PClient
+            { connected,
+              transportSession = undefined,
+              transportHost = undefined,
+              tcpTimeout = undefined,
+              batchDelay = Nothing,
+              pingErrorCount = undefined,
+              clientCorrId,
+              sentCommands,
+              sndQ,
+              rcvQ,
+              msgQ = Nothing
+            }
+      }
+
 type SMPClient = ProtocolClient ErrorType SMP.BrokerMsg
 
 -- | Type for client command data
@@ -247,11 +284,13 @@ defaultClientConfig =
 
 data Request err msg = Request
   { queueId :: QueueId,
-    responseVar :: TMVar (Response err msg)
+    responseVar :: TResponse err msg
   }
 
 type Response err msg = Either (ProtocolClientError err) msg
 
+type TResponse err msg = TMVar (Response err msg)
+
 chooseTransportHost :: NetworkConfig -> NonEmpty TransportHost -> Either (ProtocolClientError err) TransportHost
 chooseTransportHost NetworkConfig {socksProxy, hostMode, requiredHostMode} hosts =
   firstOrError $ case hostMode of
@@ -355,7 +394,7 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
             `finally` disconnected c'
 
     send :: Transport c => ProtocolClient err msg -> THandle c -> IO ()
-    send ProtocolClient {client_ = PClient {sndQ}} h = forever $ atomically (readTBQueue sndQ) >>= tPut h batchDelay
+    send ProtocolClient {client_ = PClient {sndQ}} h = forever $ atomically (readTBQueue sndQ) >>= tPutLog h
 
     receive :: Transport c => ProtocolClient err msg -> THandle c -> IO ()
     receive ProtocolClient {client_ = PClient {rcvQ}} h = forever $ tGet h >>= atomically . writeTBQueue rcvQ
@@ -597,55 +636,100 @@ okSMPCommands cmd c qs = L.map response <$> sendProtocolCommands c cs
 sendSMPCommand :: PartyI p => SMPClient -> Maybe C.APrivateSignKey -> QueueId -> Command p -> ExceptT SMPClientError IO BrokerMsg
 sendSMPCommand c pKey qId cmd = sendProtocolCommand c pKey qId (Cmd sParty cmd)
 
-type PCTransmission err msg = (SentRawTransmission, TMVar (Response err msg))
+type PCTransmission err msg = (SentRawTransmission, TResponse err msg)
 
 -- | Send multiple commands with batching and collect responses
--- It will result in Int overflow on 32 bit platform for a large number of blocks (~13.4k blocks / ~1.2m subscriptions)
 -- TODO switch to timeout or TimeManager that supports Int64
-sendProtocolCommands :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> NonEmpty (ClientCommand msg) -> IO (NonEmpty (Either (ProtocolClientError err) msg))
-sendProtocolCommands c@ProtocolClient {client_ = PClient {sndQ, tcpTimeout, batchDelay}, batch, blockSize, timeoutPerBlock} cs = do
-  (h :| ts) <- mapM (runExceptT . mkTransmission c) cs
-  let h' :: Either (ProtocolClientError err) (PCTransmission err msg, Int) = (,timeoutPerBlock) <$> h
-      batchSz = if batch then either (const 0) tSize h else 0
-      ts' :: NonEmpty (Either (ProtocolClientError err) (PCTransmission err msg, Int)) =
-        L.reverse . fst3 $ foldl' batchTimeouts ([h'], timeoutPerBlock, batchSz) ts
-      ts_ :: (Maybe (NonEmpty SentRawTransmission)) =
-        L.nonEmpty . map (fst . fst) . rights $ L.toList ts'
-  mapM_ (atomically . writeTBQueue sndQ) ts_
-  forConcurrently ts' $ \case
-    Right ((_t, r), bt) -> withTimeout c (tcpTimeout + bt) (atomically $ takeTMVar r)
-    Left e -> pure $ Left e
+sendProtocolCommands :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> NonEmpty (ClientCommand msg) -> IO (NonEmpty (Response err msg))
+sendProtocolCommands c@ProtocolClient {client_ = PClient {sndQ, tcpTimeout}, batch, blockSize} cs = do
+  bs <- batchClientTransmissions batch blockSize <$> mapM (runExceptT . mkTransmission c) cs
+  validate . concat =<< mapM sendBatch bs
   where
-    fst3 (x, _, _) = x
-    -- tSize calculation matches the batching logic in tPut that does actual breaking of transmissions into blocks
-    tSize :: PCTransmission err msg -> Int
-    tSize ((sig, t), _) = maybe 0 C.signatureSize sig + B.length t + 3 -- 1 byte for signature size + 2 bytes for transmission size
-    batchTimeouts :: (NonEmpty (Either (ProtocolClientError err) (PCTransmission err msg, Int)), Int, Int) -> Either (ProtocolClientError err) (PCTransmission err msg) -> (NonEmpty (Either (ProtocolClientError err) (PCTransmission err msg, Int)), Int, Int)
-    batchTimeouts (ts, bt, batchSz) = \case
-      Left e -> (Left e <| ts, bt, batchSz)
-      Right t
-        | not batch ->
-          (Right (t, bt') <| ts, bt', 0)
-        | batchSz' + 1 > blockSize ->
-          (Right (t, bt') <| ts, bt', tSz)
-        | otherwise -> -- same block in the batch
-          (Right (t, bt) <| ts, bt, batchSz') -- 1 byte for the number of transmissions in the batch
+    validate :: [Response err msg] -> IO (NonEmpty (Response err msg))
+    validate rs
+      | diff == 0 = pure $ L.fromList rs
+      | diff > 0 = do
+        putStrLn "send error: fewer responses than expected"
+        pure $ L.fromList $ rs <> replicate diff (Left $ PCETransportError TEBadBlock)
+      | otherwise = do
+        putStrLn "send error: more responses than expected"
+        pure $ L.fromList $ take (L.length cs) rs
         where
-          batchSz' = batchSz + tSz
-          bt' = bt + timeoutPerBlock + fromMaybe 0 batchDelay
-          tSz = tSize t
+          diff = L.length cs - length rs
+    sendBatch :: ClientBatch err msg -> IO [Response err msg]
+    sendBatch b = do
+      case b of
+        CBLargeTransmission -> [Left (PCETransportError TELargeMsg)] <$ putStrLn "send error: large message"
+        CBTransmissions n s rs -> do
+          when (n > 0) $ atomically $ writeTBQueue sndQ $ tEncodeBatch n s
+          forConcurrently rs $ \case
+            Right r -> withTimeout c tcpTimeout (atomically $ takeTMVar r)
+            Left e -> pure $ Left e
+        CBTransmission s r -> do
+          atomically $ writeTBQueue sndQ s
+          (: []) <$> withTimeout c tcpTimeout (atomically $ takeTMVar r)
+
+type PCTransmissionOrErr err msg = Either (ProtocolClientError err) (PCTransmission err msg)
+
+type TResponseOrErr err msg = Either (ProtocolClientError err) (TResponse err msg)
+
+data ClientBatch err msg
+  -- ByteString in CBTransmissions does not include count byte, it is added by tEncodeBatch
+  = CBTransmissions Int ByteString [TResponseOrErr err msg]
+  | CBTransmission ByteString (TResponse err msg)
+  | CBLargeTransmission
+
+-- | encodes and batches transmissions into blocks
+batchClientTransmissions :: forall err msg. Bool -> Int -> NonEmpty (PCTransmissionOrErr err msg) -> [ClientBatch err msg]
+batchClientTransmissions batch bSize
+  | batch = reverse . mkBatch []
+  | otherwise = map mkBatch1 . L.toList
+  where
+    mkBatch :: [ClientBatch err msg] -> NonEmpty (PCTransmissionOrErr err msg) -> [ClientBatch err msg]
+    mkBatch bs ts =
+      let (b, ts_) = encodeBatch 0 "" [] ts
+          bs' = b : bs
+       in maybe bs' (mkBatch bs') ts_
+    mkBatch1 :: PCTransmissionOrErr err msg -> ClientBatch err msg
+    mkBatch1 = \case
+      Left e -> CBTransmissions 0 "" [Left e]
+      Right (t, r) ->
+        let s = tEncode t
+         in if B.length s > bSize - 2 then CBLargeTransmission else CBTransmission s r
+    encodeBatch :: Int -> ByteString -> [TResponseOrErr err msg] -> NonEmpty (PCTransmissionOrErr err msg) -> (ClientBatch err msg, Maybe (NonEmpty (PCTransmissionOrErr err msg)))
+    encodeBatch n s rs ts@(t_ :| ts_)
+      | n == 255 = (res, Just ts)
+      | otherwise = case t_ of
+          Left e -> next n s (Left e : rs)
+          Right (t, r)
+            | B.length s' <= bSize - 3 -> next (n + 1) s' (Right r : rs)
+            | null rs -> (CBLargeTransmission, L.nonEmpty ts_)
+            | otherwise -> (res, Just ts)
+            where
+              s' = s <> smpEncode (Large $ tEncode t)
+      where
+        res = CBTransmissions n s (reverse rs)
+        next n' s' rs' = case L.nonEmpty ts_ of
+          Just ts' -> encodeBatch n' s' rs' ts'
+          Nothing -> (CBTransmissions n' s' (reverse rs'), Nothing)
 
 -- | Send Protocol command
 sendProtocolCommand :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> Maybe C.APrivateSignKey -> QueueId -> ProtoCommand msg -> ExceptT (ProtocolClientError err) IO msg
-sendProtocolCommand c@ProtocolClient {client_ = PClient {sndQ, tcpTimeout}} pKey qId cmd = do
+sendProtocolCommand c@ProtocolClient {client_ = PClient {sndQ, tcpTimeout}, batch, blockSize} pKey qId cmd = do
   (t, r) <- mkTransmission c (pKey, qId, cmd)
   ExceptT $ sendRecv t r
   where
     -- two separate "atomically" needed to avoid blocking
-    sendRecv :: SentRawTransmission -> TMVar (Response err msg) -> IO (Response err msg)
-    sendRecv t r = atomically (writeTBQueue sndQ [t]) >> withTimeout c tcpTimeout (atomically $ takeTMVar r)
+    sendRecv :: SentRawTransmission -> TResponse err msg -> IO (Response err msg)
+    sendRecv t r
+      | B.length s > blockSize - 2 = pure $ Left $ PCETransportError TELargeMsg
+      | otherwise = atomically (writeTBQueue sndQ s) >> withTimeout c tcpTimeout (atomically $ takeTMVar r)
+      where
+        s
+          | batch = tEncodeBatch 1 . smpEncode . Large $ tEncode t
+          | otherwise = tEncode t
 
-withTimeout :: ProtocolClient err msg -> Int -> IO (Either (ProtocolClientError err) msg) -> IO (Either (ProtocolClientError err) msg)
+withTimeout :: ProtocolClient err msg -> Int -> IO (Response err msg) -> IO (Response err msg)
 withTimeout ProtocolClient {client_ = PClient {pingErrorCount}} t a = do
   timeout t a >>= \case
     Just r -> atomically (writeTVar pingErrorCount 0) >> pure r
@@ -664,7 +748,7 @@ mkTransmission ProtocolClient {sessionId, thVersion, client_ = PClient {clientCo
       pure . CorrId $ bshow i
     signTransmission :: ByteString -> SentRawTransmission
     signTransmission t = ((`C.sign` t) <$> pKey, t)
-    mkRequest :: CorrId -> STM (TMVar (Response err msg))
+    mkRequest :: CorrId -> STM (TResponse err msg)
     mkRequest corrId = do
       r <- newEmptyTMVar
       TM.insert corrId (Request qId r) sentCommands

src/Simplex/Messaging/Protocol.hs

@@ -15,6 +15,7 @@
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StrictData #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilyDependencies #-}
@@ -132,12 +133,15 @@ module Simplex.Messaging.Protocol
     noAuthSrv,
 
     -- * TCP transport functions
+    TransportBatch(..),
     tPut,
+    tPutLog,
     tGet,
     tParse,
     tDecodeParseValidate,
     tEncode,
     tEncodeBatch,
+    batchTransmissions,
 
     -- * exports for tests
     CommandTag (..),
@@ -1246,50 +1250,70 @@ instance Encoding CommandError where
 
 -- | Send signed SMP transmission to TCP transport.
 tPut :: Transport c => THandle c -> Maybe Int -> NonEmpty SentRawTransmission -> IO [Either TransportError ()]
-tPut th delay_ trs 
-  | batch th = tPutBatch [] $ L.map tEncode trs
-  | otherwise = forM (L.toList trs) $ tPutLog . tEncode
+tPut th delay_ = fmap concat . mapM tPutBatch . batchTransmissions (batch th) (blockSize th)
   where
-    tPutBatch :: [Either TransportError ()] -> NonEmpty ByteString -> IO [Either TransportError ()]
-    tPutBatch rs ts = do
+    tPutBatch :: TransportBatch -> IO [Either TransportError ()]
+    tPutBatch = \case
+      TBLargeTransmission -> [Left TELargeMsg] <$ putStrLn "tPut error: large message"
+      TBTransmissions n s -> replicate n <$> (tPutLog th (tEncodeBatch n s) <* mapM_ threadDelay delay_)
+      TBTransmission s -> (: []) <$> tPutLog th s
+
+tPutLog :: Transport c => THandle c -> ByteString -> IO (Either TransportError ())
+tPutLog th s = do
+  r <- tPutBlock th s
+  case r of
+    Left e -> putStrLn ("tPut error: " <> show e)
+    _ -> pure ()
+  pure r
+
+-- ByteString does not include length byte, it is added by tEncodeBatch
+data TransportBatch = TBTransmissions Int ByteString | TBTransmission ByteString |  TBLargeTransmission
+
+-- | encodes and batches transmissions into blocks, 
+batchTransmissions :: Bool -> Int -> NonEmpty SentRawTransmission -> [TransportBatch]
+batchTransmissions batch bSize
+  | batch = reverse . mkBatch [] . L.map tEncode
+  | otherwise = map (mkBatch1 . tEncode) . L.toList
+  where
+    mkBatch :: [TransportBatch] -> NonEmpty ByteString -> [TransportBatch]
+    mkBatch rs ts =
       let (n, s, ts_) = encodeBatch 0 "" ts
-      r <- if n == 0 then largeMsg else replicate n <$> tPutLog (tEncodeBatch n s)
-      let rs' = rs <> r
-      case ts_ of
-        Just ts' -> mapM_ threadDelay delay_ >> tPutBatch rs' ts'
-        _ -> pure rs'
-    largeMsg = putStrLn "tPut error: large message" >> pure [Left TELargeMsg]
-    tPutLog s = do
-      r <- tPutBlock th s
-      case r of
-        Left e -> putStrLn ("tPut error: " <> show e)
-        _ -> pure ()
-      pure r
+          r = if n == 0 then TBLargeTransmission else TBTransmissions n s
+          rs' = r : rs
+      in case ts_ of
+          Just ts' -> mkBatch rs' ts'
+          _ -> rs'
+    mkBatch1 :: ByteString -> TransportBatch
+    mkBatch1 s = if B.length s > bSize - 2 then TBLargeTransmission else TBTransmission s
     encodeBatch :: Int -> ByteString -> NonEmpty ByteString -> (Int, ByteString, Maybe (NonEmpty ByteString))
     encodeBatch n s ts@(t :| ts_)
       | n == 255 = (n, s, Just ts)
       | otherwise =
         let s' = s <> smpEncode (Large t)
             n' = n + 1
-         in if B.length s' > blockSize th - 1 -- one byte is reserved for the number of messages in the batch
+         in if B.length s' > bSize - 3 -- one byte is reserved for the number of messages in the batch
               then (n,s,) $ if n == 0 then L.nonEmpty ts_ else Just ts
               else case L.nonEmpty ts_ of
                 Just ts' -> encodeBatch n' s' ts'
                 _ -> (n', s', Nothing)
 
-tEncode :: (Maybe C.ASignature, ByteString) -> ByteString
+tEncode :: SentRawTransmission -> ByteString
 tEncode (sig, t) = smpEncode (C.signatureBytes sig) <> t
+{-# INLINE tEncode #-}
 
 tEncodeBatch :: Int -> ByteString -> ByteString
 tEncodeBatch n s = lenEncode n `B.cons` s
+{-# INLINE tEncodeBatch #-}
 
 encodeTransmission :: ProtocolEncoding e c => Version -> ByteString -> Transmission c -> ByteString
 encodeTransmission v sessionId (CorrId corrId, queueId, command) =
   smpEncode (sessionId, corrId, queueId) <> encodeProtocol v command
+{-# INLINE encodeTransmission #-}
 
 -- | Receive and parse transmission from the TCP transport (ignoring any trailing padding).
 tGetParse :: Transport c => THandle c -> IO (NonEmpty (Either TransportError RawTransmission))
 tGetParse th = eitherList (tParse $ batch th) <$> tGetBlock th
+{-# INLINE tGetParse #-}
 
 tParse :: Bool -> ByteString -> NonEmpty (Either TransportError RawTransmission)
 tParse batch s

src/Simplex/Messaging/Server.hs

@@ -31,6 +31,7 @@ module Simplex.Messaging.Server
     disconnectTransport,
     verifyCmdSignature,
     dummyVerifyCmd,
+    randomId,
   )
 where
 

tests/CoreTests/BatchingTests.hs

@@ -0,0 +1,180 @@
+{-# LANGUAGE LambdaCase #-}
+
+module CoreTests.BatchingTests (batchingTests) where
+
+import Control.Concurrent.STM
+import Control.Monad.Except
+import Crypto.Random (MonadRandom(..))
+import Data.ByteString.Char8 (ByteString)
+import qualified Data.ByteString.Char8 as B
+import qualified Data.List.NonEmpty as L
+import Simplex.Messaging.Client
+import qualified Simplex.Messaging.Crypto as C
+import Simplex.Messaging.Protocol
+import Simplex.Messaging.Transport
+import Simplex.Messaging.Version (VersionRange(..))
+import Test.Hspec
+
+batchingTests :: Spec
+batchingTests = do
+  describe "batchTransmissions" $ do
+    it "should batch with 90 subscriptions per batch" testBatchSubscriptions
+    it "should break on message that does not fit" testBatchWithMessage
+    it "should break on large message" testBatchWithLargeMessage
+  describe "batchClientTransmissions" $ do
+    it "should batch with 90 subscriptions per batch" testClientBatchSubscriptions
+    it "should break on message that does not fit" testClientBatchWithMessage
+    it "should break on large message" testClientBatchWithLargeMessage
+
+testBatchSubscriptions :: IO ()
+testBatchSubscriptions = do
+  sessId <- getRandomBytes 32
+  subs <- replicateM 200 $ randomSUB sessId
+  let batches1 = batchTransmissions False smpBlockSize $ L.fromList subs
+  all lenOk1 batches1 `shouldBe` True
+  length batches1 `shouldBe` 200
+  let batches = batchTransmissions True smpBlockSize $ L.fromList subs
+  length batches `shouldBe` 3
+  [TBTransmissions n1 s1, TBTransmissions n2 s2, TBTransmissions n3 s3] <- pure batches
+  (n1, n2, n3) `shouldBe` (90, 90, 20)
+  all lenOk [s1, s2, s3] `shouldBe` True
+
+testBatchWithMessage :: IO ()
+testBatchWithMessage = do
+  sessId <- getRandomBytes 32
+  subs1 <- replicateM 60 $ randomSUB sessId
+  send <- randomSEND sessId 8000
+  subs2 <- replicateM 40 $ randomSUB sessId
+  let cmds = subs1 <> [send] <> subs2
+      batches1 = batchTransmissions False smpBlockSize $ L.fromList cmds
+  all lenOk1 batches1 `shouldBe` True
+  length batches1 `shouldBe` 101
+  let batches = batchTransmissions True smpBlockSize $ L.fromList cmds
+  length batches `shouldBe` 2
+  [TBTransmissions n1 s1, TBTransmissions n2 s2] <- pure batches
+  (n1, n2) `shouldBe` (60, 41)
+  all lenOk [s1, s2] `shouldBe` True
+
+testBatchWithLargeMessage :: IO ()
+testBatchWithLargeMessage = do
+  sessId <- getRandomBytes 32
+  subs1 <- replicateM 60 $ randomSUB sessId
+  send <- randomSEND sessId 17000
+  subs2 <- replicateM 100 $ randomSUB sessId
+  let cmds = subs1 <> [send] <> subs2
+      batches1 = batchTransmissions False smpBlockSize $ L.fromList cmds
+  all lenOk1 batches1 `shouldBe` False
+  length batches1 `shouldBe` 161
+  let batches1' = take 60 batches1 <> drop 61 batches1
+  all lenOk1 batches1' `shouldBe` True
+  length batches1' `shouldBe` 160
+  let batches = batchTransmissions True smpBlockSize $ L.fromList cmds
+  length batches `shouldBe` 4
+  [TBTransmissions n1 s1, TBLargeTransmission, TBTransmissions n2 s2, TBTransmissions n3 s3] <- pure batches
+  (n1, n2, n3) `shouldBe` (60, 90, 10)
+  all lenOk [s1, s2, s3] `shouldBe` True
+
+testClientBatchSubscriptions :: IO ()
+testClientBatchSubscriptions = do
+  sessId <- getRandomBytes 32
+  client <- atomically $ clientStub sessId
+  subs <- replicateM 200 $ randomSUBCmd client
+  let batches1 = batchClientTransmissions False smpBlockSize $ L.fromList subs
+  all lenOk1' batches1 `shouldBe` True
+  let batches = batchClientTransmissions True smpBlockSize $ L.fromList subs
+  length batches `shouldBe` 3
+  [CBTransmissions n1 s1 rs1, CBTransmissions n2 s2 rs2, CBTransmissions n3 s3 rs3] <- pure batches
+  (n1, n2, n3) `shouldBe` (90, 90, 20)
+  (length rs1, length rs2, length rs3) `shouldBe` (90, 90, 20)
+  all lenOk [s1, s2, s3] `shouldBe` True
+
+testClientBatchWithMessage :: IO ()
+testClientBatchWithMessage = do
+  sessId <- getRandomBytes 32
+  client <- atomically $ clientStub sessId
+  subs1 <- replicateM 60 $ randomSUBCmd client
+  send <- randomSENDCmd client 8000
+  subs2 <- replicateM 40 $ randomSUBCmd client
+  let cmds = subs1 <> [send] <> subs2
+      batches1 = batchClientTransmissions False smpBlockSize $ L.fromList cmds
+  all lenOk1' batches1 `shouldBe` True
+  length batches1 `shouldBe` 101
+  let batches = batchClientTransmissions True smpBlockSize $ L.fromList cmds
+  length batches `shouldBe` 2
+  [CBTransmissions n1 s1 rs1, CBTransmissions n2 s2 rs2] <- pure batches
+  (n1, n2) `shouldBe` (60, 41)
+  (length rs1, length rs2) `shouldBe` (60, 41)
+  all lenOk [s1, s2] `shouldBe` True
+
+testClientBatchWithLargeMessage :: IO ()
+testClientBatchWithLargeMessage = do
+  sessId <- getRandomBytes 32
+  client <- atomically $ clientStub sessId
+  subs1 <- replicateM 60 $ randomSUBCmd client
+  send <- randomSENDCmd client 17000
+  subs2 <- replicateM 100 $ randomSUBCmd client
+  let cmds = subs1 <> [send] <> subs2
+      batches1 = batchClientTransmissions False smpBlockSize $ L.fromList cmds
+  all lenOk1' batches1 `shouldBe` False
+  length batches1 `shouldBe` 161
+  let batches1' = take 60 batches1 <> drop 61 batches1
+  all lenOk1' batches1' `shouldBe` True
+  length batches1' `shouldBe` 160
+  --
+  let batches = batchClientTransmissions True smpBlockSize $ L.fromList cmds
+  length batches `shouldBe` 4
+  [CBTransmissions n1 s1 rs1, CBLargeTransmission, CBTransmissions n2 s2 rs2, CBTransmissions n3 s3 rs3] <- pure batches
+  (n1, n2, n3) `shouldBe` (60, 90, 10)
+  (length rs1, length rs2, length rs3) `shouldBe` (60, 90, 10)
+  all lenOk [s1, s2, s3] `shouldBe` True
+  --
+  let cmds' = [send] <> subs1 <> subs2
+  let batches' = batchClientTransmissions True smpBlockSize $ L.fromList cmds'
+  length batches' `shouldBe` 3
+  [CBLargeTransmission, CBTransmissions n1' s1' rs1', CBTransmissions n2' s2' rs2'] <- pure batches'
+  (n1', n2') `shouldBe` (90, 70)
+  (length rs1', length rs2') `shouldBe` (90, 70)
+  all lenOk [s1', s2'] `shouldBe` True
+
+randomSUB :: ByteString -> IO (Maybe C.ASignature, ByteString)
+randomSUB sessId = do
+  rId <- getRandomBytes 24
+  corrId <- CorrId <$> getRandomBytes 3
+  (_, rpKey) <- C.generateSignatureKeyPair C.SEd448
+  let s = encodeTransmission (maxVersion supportedSMPServerVRange) sessId (corrId, rId, Cmd SRecipient SUB)
+  pure (Just $ C.sign rpKey s, s)
+
+randomSUBCmd :: ProtocolClient ErrorType BrokerMsg -> IO (Either (ProtocolClientError ErrorType) (PCTransmission ErrorType BrokerMsg))
+randomSUBCmd c = do
+  rId <- getRandomBytes 24
+  (_, rpKey) <- C.generateSignatureKeyPair C.SEd448
+  runExceptT $ mkTransmission c (Just rpKey, rId, Cmd SRecipient SUB)
+
+randomSEND :: ByteString -> Int -> IO (Maybe C.ASignature, ByteString)
+randomSEND sessId len = do
+  sId <- getRandomBytes 24
+  corrId <- CorrId <$> getRandomBytes 3
+  (_, rpKey) <- C.generateSignatureKeyPair C.SEd448
+  msg <- getRandomBytes len
+  let s = encodeTransmission (maxVersion supportedSMPServerVRange) sessId (corrId, sId, Cmd SSender $ SEND noMsgFlags msg)
+  pure (Just $ C.sign rpKey s, s)
+
+randomSENDCmd :: ProtocolClient ErrorType BrokerMsg -> Int -> IO (Either (ProtocolClientError ErrorType) (PCTransmission ErrorType BrokerMsg))
+randomSENDCmd c len = do
+  sId <- getRandomBytes 24
+  (_, rpKey) <- C.generateSignatureKeyPair C.SEd448
+  msg <- getRandomBytes len
+  runExceptT $ mkTransmission c (Just rpKey, sId, Cmd SSender $ SEND noMsgFlags msg)
+
+lenOk :: ByteString -> Bool
+lenOk s = 0 < B.length s && B.length s <= smpBlockSize - 2
+
+lenOk1 :: TransportBatch -> Bool
+lenOk1 = \case
+  TBTransmission s -> lenOk s
+  _ -> False
+
+lenOk1' :: ClientBatch err msg -> Bool
+lenOk1' = \case
+  CBTransmission s _ -> lenOk s
+  _ -> False

tests/Test.hs

@@ -4,6 +4,7 @@ import AgentTests (agentTests)
 import AgentTests.SchemaDump (schemaDumpTest)
 import CLITests
 import Control.Logger.Simple
+import CoreTests.BatchingTests
 import CoreTests.CryptoTests
 import CoreTests.EncodingTests
 import CoreTests.ProtocolErrorTests
@@ -37,6 +38,7 @@ main = do
       $ do
         describe "Agent SQLite schema dump" schemaDumpTest
         describe "Core tests" $ do
+          describe "Batching tests" batchingTests
           describe "Encoding tests" encodingTests
           describe "Protocol error tests" protocolErrorTests
           describe "Version range" versionRangeTests

tests/XFTPAgent.hs

@@ -15,7 +15,7 @@ import qualified Data.ByteString.Char8 as B
 import Data.Int (Int64)
 import Data.List (find, isSuffixOf)
 import Data.Maybe (fromJust)
-import SMPAgentClient (agentCfg, initAgentServers, testDB)
+import SMPAgentClient (agentCfg, initAgentServers, testDB, testDB2)
 import Simplex.FileTransfer.Description
 import Simplex.FileTransfer.Protocol (FileParty (..), XFTPErrorType (AUTH))
 import Simplex.FileTransfer.Server.Env (XFTPServerConfig (..))
@@ -93,10 +93,11 @@ testXFTPAgentSendReceive = withXFTPServer $ do
   testReceiveDelete rfd2 filePath
   where
     testReceiveDelete rfd originalFilePath = do
-      rcp <- getSMPAgentClient' agentCfg initAgentServers testDB
+      rcp <- getSMPAgentClient' agentCfg initAgentServers testDB2
       runRight_ $ do
         rfId <- testReceive rcp rfd originalFilePath
         xftpDeleteRcvFile rcp rfId
+      disconnectAgentClient rcp
 
 createRandomFile :: IO FilePath
 createRandomFile = do
@@ -145,7 +146,7 @@ testXFTPAgentReceiveRestore = withGlobalLogging logCfgNoLogs $ do
       pure rfd
 
   -- receive file - should not succeed with server down
-  rcp <- getSMPAgentClient' agentCfg initAgentServers testDB
+  rcp <- getSMPAgentClient' agentCfg initAgentServers testDB2
   rfId <- runRight $ do
     xftpStartWorkers rcp (Just recipientFiles)
     rfId <- xftpReceiveFile rcp 1 rfd
@@ -159,7 +160,7 @@ testXFTPAgentReceiveRestore = withGlobalLogging logCfgNoLogs $ do
 
   withXFTPServerStoreLogOn $ \_ -> do
     -- receive file - should start downloading with server up
-    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $ xftpStartWorkers rcp' (Just recipientFiles)
     ("", rfId', RFPROG _ _) <- rfGet rcp'
     liftIO $ rfId' `shouldBe` rfId
@@ -169,7 +170,7 @@ testXFTPAgentReceiveRestore = withGlobalLogging logCfgNoLogs $ do
 
   withXFTPServerStoreLogOn $ \_ -> do
     -- receive file - should continue downloading with server up
-    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $ xftpStartWorkers rcp' (Just recipientFiles)
     rfProgress rcp' $ mb 18
     ("", rfId', RFDONE path) <- rfGet rcp'
@@ -193,7 +194,7 @@ testXFTPAgentReceiveCleanup = withGlobalLogging logCfgNoLogs $ do
       pure rfd
 
   -- receive file - should not succeed with server down
-  rcp <- getSMPAgentClient' agentCfg initAgentServers testDB
+  rcp <- getSMPAgentClient' agentCfg initAgentServers testDB2
   rfId <- runRight $ do
     xftpStartWorkers rcp (Just recipientFiles)
     rfId <- xftpReceiveFile rcp 1 rfd
@@ -207,7 +208,7 @@ testXFTPAgentReceiveCleanup = withGlobalLogging logCfgNoLogs $ do
 
   withXFTPServerThreadOn $ \_ -> do
     -- receive file - should fail with AUTH error
-    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp' <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $ xftpStartWorkers rcp' (Just recipientFiles)
     ("", rfId', RFERR (INTERNAL "XFTP {xftpErr = AUTH}")) <- rfGet rcp'
     rfId' `shouldBe` rfId
@@ -259,7 +260,7 @@ testXFTPAgentSendRestore = withGlobalLogging logCfgNoLogs $ do
     doesFileExist encPath `shouldReturn` False
 
     -- receive file
-    rcp <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $
       void $ testReceive rcp rfd1 filePath
 
@@ -309,7 +310,7 @@ testXFTPAgentDelete = withGlobalLogging logCfgNoLogs $
     (sfId, sndDescr, rfd1, rfd2) <- runRight $ testSend sndr filePath
 
     -- receive file
-    rcp1 <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp1 <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $
       void $ testReceive rcp1 rfd1 filePath
 
@@ -321,12 +322,13 @@ testXFTPAgentDelete = withGlobalLogging logCfgNoLogs $
       xftpDeleteSndFileRemote sndr 1 sfId sndDescr
       Nothing <- liftIO $ 100000 `timeout` sfGet sndr
       pure ()
+    disconnectAgentClient rcp1
 
     threadDelay 1000000
     length <$> listDirectory xftpServerFiles `shouldReturn` 0
 
     -- receive file - should fail with AUTH error
-    rcp2 <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp2 <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight $ do
       xftpStartWorkers rcp2 (Just recipientFiles)
       rfId <- xftpReceiveFile rcp2 1 rfd2
@@ -343,10 +345,11 @@ testXFTPAgentDeleteRestore = withGlobalLogging logCfgNoLogs $ do
     (sfId, sndDescr, rfd1, rfd2) <- runRight $ testSend sndr filePath
 
     -- receive file
-    rcp1 <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp1 <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight_ $
       void $ testReceive rcp1 rfd1 filePath
-
+    disconnectAgentClient rcp1
+    disconnectAgentClient sndr
     pure (sfId, sndDescr, rfd2)
 
   -- delete file - should not succeed with server down
@@ -369,7 +372,7 @@ testXFTPAgentDeleteRestore = withGlobalLogging logCfgNoLogs $ do
     length <$> listDirectory xftpServerFiles `shouldReturn` 0
 
     -- receive file - should fail with AUTH error
-    rcp2 <- getSMPAgentClient' agentCfg initAgentServers testDB
+    rcp2 <- getSMPAgentClient' agentCfg initAgentServers testDB2
     runRight $ do
       xftpStartWorkers rcp2 (Just recipientFiles)
       rfId <- xftpReceiveFile rcp2 1 rfd2
@@ -394,15 +397,12 @@ testXFTPAgentRequestAdditionalRecipientIDs = withXFTPServer $ do
 
   -- receive file using different descriptions
   -- ! revise number of recipients and indexes if xftpMaxRecipientsPerRequest is changed
-  testReceive' (head rfds) filePath
-  testReceive' (rfds !! 99) filePath
-  testReceive' (rfds !! 299) filePath
-  testReceive' (rfds !! 499) filePath
-  where
-    testReceive' rfd originalFilePath = do
-      rcp <- getSMPAgentClient' agentCfg initAgentServers testDB
-      runRight_ $
-        void $ testReceive rcp rfd originalFilePath
+  rcp <- getSMPAgentClient' agentCfg initAgentServers testDB2
+  runRight_ $ do
+    void $ testReceive rcp (head rfds) filePath
+    void $ testReceive rcp (rfds !! 99) filePath
+    void $ testReceive rcp (rfds !! 299) filePath
+    void $ testReceive rcp (rfds !! 499) filePath
 
 testXFTPServerTest :: Maybe BasicAuth -> XFTPServerWithAuth -> IO (Maybe ProtocolTestFailure)
 testXFTPServerTest newFileBasicAuth srv =