batcher tests

simplex-chat.cabal

@@ -29,6 +29,7 @@ library
       Simplex.Chat.Archive
       Simplex.Chat.Bot
       Simplex.Chat.Bot.KnownContacts
+      Simplex.Chat.ByteStringBatcher
       Simplex.Chat.Call
       Simplex.Chat.Controller
       Simplex.Chat.Core
@@ -532,6 +533,7 @@ test-suite simplex-chat-test
   type: exitcode-stdio-1.0
   main-is: Test.hs
   other-modules:
+      ByteStringBatcherTests
       Bots.BroadcastTests
       Bots.DirectoryTests
       ChatClient

src/Simplex/Chat.hs

@@ -55,6 +55,7 @@ import Data.Time.Clock.System (systemToUTCTime)
 import Data.Word (Word32)
 import qualified Database.SQLite.Simple as SQL
 import Simplex.Chat.Archive
+import Simplex.Chat.ByteStringBatcher (BSBatch (..), batchByteStringObjects, partitionBatches)
 import Simplex.Chat.Call
 import Simplex.Chat.Controller
 import Simplex.Chat.Files
@@ -5618,11 +5619,11 @@ sendBatchedDirectMessages conn@Connection {connId} events connOrGroupId = do
   (errs, msgs) <- partitionEithers <$> createSndMessages
   unless (null errs) $ toView $ CRChatErrors Nothing errs
   forM_ (L.nonEmpty msgs) $ \msgs' -> do
-    let (largeMsgs, msgBatches) = partitionBatches $ batchChatMessages msgs'
+    let (largeMsgs, msgBatches) = partitionBatches $ batchByteStringObjects maxChatMsgSize msgs'
         -- shouldn't happen, as large messages would have caused createNewSndMessage to throw SELargeMsg
         errs' = map (\SndMessage {msgId} -> ChatError $ CEInternalError ("large message " <> show msgId)) largeMsgs
     unless (null errs') $ toView $ CRChatErrors Nothing errs'
-    forM_ msgBatches $ \(MessagesBatch batchBuilder sndMsgs) -> do
+    forM_ msgBatches $ \(BSBatch batchBuilder sndMsgs) -> do
       let batchBody = LB.toStrict $ Builder.toLazyByteString batchBuilder
       agentMsgId <- withAgent $ \a -> sendMessage a (aConnId conn) MsgFlags {notification = True} batchBody
       let sndMsgDelivery = SndMsgDelivery {connId, agentMsgId}
@@ -5638,69 +5639,6 @@ sendBatchedDirectMessages conn@Connection {connId} events connOrGroupId = do
       pure $ first ChatErrorStore r
     encodeMessage chatVRange evnt sharedMsgId =
       encodeChatMessage ChatMessage {chatVRange, msgId = Just sharedMsgId, chatMsgEvent = evnt}
-    partitionBatches :: [ChatMessageBatch] -> ([SndMessage], [MessagesBatch])
-    partitionBatches = foldr partition' ([], [])
-      where
-        partition' :: ChatMessageBatch -> ([SndMessage], [MessagesBatch]) -> ([SndMessage], [MessagesBatch])
-        partition' (CMBMessages msgBatch) (largeMsgs, msgBatches) = (largeMsgs, msgBatch : msgBatches)
-        partition' (CMBLargeMessage largeMsg) (largeMsgs, msgBatches) = (largeMsg : largeMsgs, msgBatches)
-
-data MessagesBatch = MessagesBatch Builder.Builder [SndMessage]
-
-data ChatMessageBatch
-  = CMBMessages MessagesBatch
-  | CMBLargeMessage SndMessage
-
-batchChatMessages :: NonEmpty SndMessage -> [ChatMessageBatch]
-batchChatMessages = reverse . mkBatch []
-  where
-    mkBatch :: [ChatMessageBatch] -> NonEmpty SndMessage -> [ChatMessageBatch]
-    mkBatch batches msgs =
-      let (batch, msgs_) = encodeBatch mempty 0 0 [] msgs
-          batches' = batch : batches
-       in maybe batches' (mkBatch batches') msgs_
-    encodeBatch :: Builder.Builder -> Int64 -> Int -> [SndMessage] -> NonEmpty SndMessage -> (ChatMessageBatch, Maybe (NonEmpty SndMessage))
-    encodeBatch builder len cnt batchedMsgs remainingMsgs@(msg :| msgs_)
-      -- message fits
-      | len' <= maxSize' =
-          case L.nonEmpty msgs_ of
-            Just msgs' -> encodeBatch builder' len' cnt' batchedMsgs' msgs'
-            Nothing -> completeBatchLastMsgFits
-      -- message doesn't fit
-      | cnt == 0 = (CMBLargeMessage msg, L.nonEmpty msgs_)
-      | otherwise = completeBatchMsgDoesntFit
-      where
-        SndMessage {msgBody} = msg
-        cnt' = cnt + 1
-        (len', builder')
-          | cnt' == 1 =
-              ( LB.length msgBody, -- initially len = 0
-                Builder.lazyByteString msgBody
-              )
-          | cnt' == 2 =
-              ( len + LB.length msgBody + 2, -- for opening bracket "[" and comma ","
-                "[" <> builder <> "," <> Builder.lazyByteString msgBody
-              )
-          | otherwise =
-              ( len + LB.length msgBody + 1, -- for comma ","
-                builder <> "," <> Builder.lazyByteString msgBody
-              )
-        maxSize'
-          | cnt' == 1 = maxChatMsgSize
-          | otherwise = maxChatMsgSize - 1 -- for closing bracket "]"
-        batchedMsgs' = msg : batchedMsgs
-        completeBatchLastMsgFits =
-          (CMBMessages $ MessagesBatch completeBuilder (reverse batchedMsgs'), Nothing)
-          where
-            completeBuilder
-              | cnt' == 1 = builder' -- if last message fits, we look at current cnt'
-              | otherwise = builder' <> "]"
-        completeBatchMsgDoesntFit =
-          (CMBMessages $ MessagesBatch completeBuilder (reverse batchedMsgs), Just remainingMsgs)
-          where
-            completeBuilder
-              | cnt == 1 = builder -- if message doesn't fit, we look at previous cnt
-              | otherwise = builder <> "]"
 
 directMessage :: (MsgEncodingI e, ChatMonad m) => ChatMsgEvent e -> m ByteString
 directMessage chatMsgEvent = do

src/Simplex/Chat/ByteStringBatcher.hs

@@ -0,0 +1,97 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Simplex.Chat.ByteStringBatcher
+  ( HasByteString (..),
+    BSBatch (..),
+    BSBatcherOutput (..),
+    batchByteStringObjects,
+    partitionBatches,
+  )
+where
+
+import qualified Data.ByteString.Builder as BB
+import qualified Data.ByteString.Lazy as L
+import qualified Data.ByteString.Lazy.Char8 as LB
+import Data.Int (Int64)
+import Data.List.NonEmpty (NonEmpty (..))
+import qualified Data.List.NonEmpty as L
+
+class HasByteString a where
+  getByteString :: a -> L.ByteString
+
+instance HasByteString L.ByteString where
+  getByteString = id
+
+data HasByteString a => BSBatch a = BSBatch BB.Builder [a]
+
+data HasByteString a => BSBatcherOutput a
+  = BatcherOutputBatch (BSBatch a)
+  | BatcherOutputLarge a
+
+-- | Batches instances of HasByteString into batches of ByteString builders in form of JSON arrays.
+-- Does not check if the resulting batch is a valid JSON. If it is required,
+-- getByteString should return ByteString encoded JSON object.
+-- If a single element is passed, it is returned in form of JSON object instead.
+-- If an element exceeds batchLenLimit, it is returned as BatcherOutputLarge.
+batchByteStringObjects :: forall a. HasByteString a => Int64 -> NonEmpty a -> [BSBatcherOutput a]
+batchByteStringObjects batchLenLimit = reverse . mkBatch []
+  where
+    mkBatch :: [BSBatcherOutput a] -> NonEmpty a -> [BSBatcherOutput a]
+    mkBatch batches objs =
+      let (batch, objs_) = encodeBatch mempty 0 0 [] objs
+          batches' = batch : batches
+       in maybe batches' (mkBatch batches') objs_
+    encodeBatch :: BB.Builder -> Int64 -> Int -> [a] -> NonEmpty a -> (BSBatcherOutput a, Maybe (NonEmpty a))
+    encodeBatch builder len cnt batchedObjs remainingObjs@(obj :| objs_)
+      -- batched string fits
+      | len' <= maxSize' =
+          case L.nonEmpty objs_ of
+            Just objs' -> encodeBatch builder' len' cnt' batchedObjs' objs'
+            Nothing -> completeBatchLastStrFits
+      -- batched string doesn't fit
+      | cnt == 0 = (BatcherOutputLarge obj, L.nonEmpty objs_)
+      | otherwise = completeBatchStrDoesntFit
+      where
+        bStr = getByteString obj
+        cnt' = cnt + 1
+        (len', builder')
+          | cnt' == 1 =
+              ( LB.length bStr, -- initially len = 0
+                BB.lazyByteString bStr
+              )
+          | cnt' == 2 =
+              ( len + LB.length bStr + 2, -- for opening bracket "[" and comma ","
+                "[" <> builder <> "," <> BB.lazyByteString bStr
+              )
+          | otherwise =
+              ( len + LB.length bStr + 1, -- for comma ","
+                builder <> "," <> BB.lazyByteString bStr
+              )
+        maxSize'
+          | cnt' == 1 = batchLenLimit
+          | otherwise = batchLenLimit - 1 -- for closing bracket "]"
+        batchedObjs' :: [a]
+        batchedObjs' = obj : batchedObjs
+        completeBatchLastStrFits :: (BSBatcherOutput a, Maybe (NonEmpty a))
+        completeBatchLastStrFits =
+          (BatcherOutputBatch $ BSBatch completeBuilder (reverse batchedObjs'), Nothing)
+          where
+            completeBuilder
+              | cnt' == 1 = builder' -- if last string fits, we look at current cnt'
+              | otherwise = builder' <> "]"
+        completeBatchStrDoesntFit :: (BSBatcherOutput a, Maybe (NonEmpty a))
+        completeBatchStrDoesntFit =
+          (BatcherOutputBatch $ BSBatch completeBuilder (reverse batchedObjs), Just remainingObjs)
+          where
+            completeBuilder
+              | cnt == 1 = builder -- if string doesn't fit, we look at previous cnt
+              | otherwise = builder <> "]"
+
+-- | Partitions list of batcher outputs into lists of batches and large objects.
+partitionBatches :: forall a. HasByteString a => [BSBatcherOutput a] -> ([a], [BSBatch a])
+partitionBatches = foldr partition' ([], [])
+  where
+    partition' :: BSBatcherOutput a -> ([a], [BSBatch a]) -> ([a], [BSBatch a])
+    partition' (BatcherOutputBatch bStrBatch) (largeBStrs, bStrBatches) = (largeBStrs, bStrBatch : bStrBatches)
+    partition' (BatcherOutputLarge largeBStr) (largeBStrs, bStrBatches) = (largeBStr : largeBStrs, bStrBatches)

src/Simplex/Chat/Messages.hs

@@ -35,6 +35,7 @@ import Data.Type.Equality
 import Data.Typeable (Typeable)
 import Database.SQLite.Simple.FromField (FromField (..))
 import Database.SQLite.Simple.ToField (ToField (..))
+import Simplex.Chat.ByteStringBatcher (HasByteString(..))
 import Simplex.Chat.Markdown
 import Simplex.Chat.Messages.CIContent
 import Simplex.Chat.Protocol
@@ -772,6 +773,9 @@ data SndMessage = SndMessage
     msgBody :: LazyMsgBody
   }
 
+instance HasByteString SndMessage where
+  getByteString SndMessage {msgBody} = msgBody
+
 data NewRcvMessage e = NewRcvMessage
   { chatMsgEvent :: ChatMsgEvent e,
     msgBody :: MsgBody

tests/ByteStringBatcherTests.hs

@@ -0,0 +1,65 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module ByteStringBatcherTests where
+
+import qualified Data.ByteString.Builder as BB
+import qualified Data.ByteString.Lazy as L
+import Data.Int (Int64)
+import Data.List.NonEmpty (fromList)
+import Simplex.Chat.ByteStringBatcher
+import Test.Hspec
+
+byteStringBatcherTests :: Spec
+byteStringBatcherTests = describe "ByteStringBatcher tests" $ do
+  runBatcherTest 8 ["a", "b", "c"] [] ["[a,b,c]"]
+  runBatcherTest 8 ["a", "bb", "c"] [] ["[a,bb,c]"]
+  runBatcherTest 8 ["a", "b", "c", "d"] [] ["[a,b,c]","d"]
+  runBatcherTest 8 ["a", "bb", "c", "d"] [] ["[a,bb,c]","d"]
+  runBatcherTest 8 ["a", "bb", "c", "de"] [] ["[a,bb,c]","de"]
+  runBatcherTest 8 ["a", "b", "c", "d", "e"] [] ["[a,b,c]","[d,e]"]
+  runBatcherTest 8 ["aaaaa"] [] ["aaaaa"]
+  runBatcherTest 8 ["8aaaaaaa"] [] ["8aaaaaaa"]
+  runBatcherTest 8 ["aaaa","bbbb"] [] ["aaaa","bbbb"]
+  runBatcherTest 8 ["aa","bbb","cc","dd"] [] ["[aa,bbb]","[cc,dd]"]
+  runBatcherTest 8 ["9aaaaaaaa"] ["9aaaaaaaa"] []
+  runBatcherTest 8 ["aaaaa","bbb","cc"] [] ["aaaaa","[bbb,cc]"]
+  runBatcherTest 8 ["8aaaaaaa","bbb","cc"] [] ["8aaaaaaa","[bbb,cc]"]
+  runBatcherTest 8 ["9aaaaaaaa","bbb","cc"] ["9aaaaaaaa"] ["[bbb,cc]"]
+  runBatcherTest 8 ["9aaaaaaaa","bbb","cc","dd"] ["9aaaaaaaa"] ["[bbb,cc]","dd"]
+  runBatcherTest 8 ["9aaaaaaaa","bbb","cc","dd","e"] ["9aaaaaaaa"] ["[bbb,cc]","[dd,e]"]
+  runBatcherTest 8 ["bbb","cc","aaaaa"] [] ["[bbb,cc]","aaaaa"]
+  runBatcherTest 8 ["bbb","cc","8aaaaaaa"] [] ["[bbb,cc]","8aaaaaaa"]
+  runBatcherTest 8 ["bbb","cc","9aaaaaaaa"] ["9aaaaaaaa"] ["[bbb,cc]"]
+  runBatcherTest 8 ["bbb","cc","dd","9aaaaaaaa"] ["9aaaaaaaa"] ["[bbb,cc]","dd"]
+  runBatcherTest 8 ["bbb","cc","dd","e","9aaaaaaaa"] ["9aaaaaaaa"] ["[bbb,cc]","[dd,e]"]
+  runBatcherTest 8 ["bbb","cc","aaaaa","dd"] [] ["[bbb,cc]","aaaaa","dd"]
+  runBatcherTest 8 ["bbb","cc","aaaaa","dd","e"] [] ["[bbb,cc]","aaaaa","[dd,e]"]
+  runBatcherTest 8 ["bbb","cc","8aaaaaaa","dd"] [] ["[bbb,cc]","8aaaaaaa","dd"]
+  runBatcherTest 8 ["bbb","cc","8aaaaaaa","dd","e"] [] ["[bbb,cc]","8aaaaaaa","[dd,e]"]
+  runBatcherTest 8 ["bbb","cc","9aaaaaaaa"] ["9aaaaaaaa"] ["[bbb,cc]"]
+  runBatcherTest 8 ["bbb","cc","9aaaaaaaa","dd"] ["9aaaaaaaa"] ["[bbb,cc]","dd"]
+  runBatcherTest 8 ["bbb","cc","9aaaaaaaa","dd","e"] ["9aaaaaaaa"] ["[bbb,cc]","[dd,e]"]
+  runBatcherTest 8 ["9aaaaaaaa","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] []
+  runBatcherTest 8 ["8aaaaaaa","9aaaaaaaa","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["8aaaaaaa"]
+  runBatcherTest 8 ["9aaaaaaaa","8aaaaaaa","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["8aaaaaaa"]
+  runBatcherTest 8 ["9aaaaaaaa","10aaaaaaaa","8aaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["8aaaaaaa"]
+  runBatcherTest 8 ["bb","cc","dd","9aaaaaaaa","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["[bb,cc]","dd"]
+  runBatcherTest 8 ["bb","cc","9aaaaaaaa","dd","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["[bb,cc]","dd"]
+  runBatcherTest 8 ["bb","9aaaaaaaa","cc","dd","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["bb","[cc,dd]"]
+  runBatcherTest 8 ["bb","9aaaaaaaa","cc","10aaaaaaaa","dd"] ["9aaaaaaaa","10aaaaaaaa"] ["bb","cc","dd"]
+  runBatcherTest 8 ["9aaaaaaaa","bb","cc","dd","10aaaaaaaa"] ["9aaaaaaaa","10aaaaaaaa"] ["[bb,cc]","dd"]
+  runBatcherTest 8 ["9aaaaaaaa","bb","10aaaaaaaa","cc","dd"] ["9aaaaaaaa","10aaaaaaaa"] ["bb","[cc,dd]"]
+  runBatcherTest 8 ["9aaaaaaaa","10aaaaaaaa","bb","cc","dd"] ["9aaaaaaaa","10aaaaaaaa"] ["[bb,cc]","dd"]
+
+runBatcherTest :: Int64 -> [L.ByteString] -> [L.ByteString] -> [L.ByteString] -> SpecWith ()
+runBatcherTest batchLenLimit bStrs expectedLargeStrs expectedBatchedStrs =
+  it
+    ( (show bStrs <> ", limit " <> show batchLenLimit <> ": should return ")
+        <> (show (length expectedLargeStrs) <> " large, ")
+        <> (show (length expectedBatchedStrs) <> " batches")
+    )
+    $ do
+      let (largeStrs, batches) = partitionBatches $ batchByteStringObjects batchLenLimit (fromList bStrs)
+          batchedStrs = map (\(BSBatch batchBuilder _) -> BB.toLazyByteString batchBuilder) batches
+      largeStrs `shouldBe` expectedLargeStrs
+      batchedStrs `shouldBe` expectedBatchedStrs

tests/Test.hs

@@ -1,3 +1,4 @@
+import ByteStringBatcherTests
 import Bots.BroadcastTests
 import Bots.DirectoryTests
 import ChatClient
@@ -28,6 +29,7 @@ main = do
     describe "SimpleX chat protocol" protocolTests
     describe "WebRTC encryption" webRTCTests
     describe "Valid names" validNameTests
+    describe "ByteStringBatcher tests" byteStringBatcherTests
     around testBracket $ do
       describe "Mobile API Tests" mobileTests
       describe "SimpleX chat client" chatTests