diff --git a/simplex-chat.cabal b/simplex-chat.cabal index cbc83b01b2..6622b33f3b 100644 --- a/simplex-chat.cabal +++ b/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 diff --git a/src/Simplex/Chat.hs b/src/Simplex/Chat.hs index a62acd14c0..7b8a27270d 100644 --- a/src/Simplex/Chat.hs +++ b/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 diff --git a/src/Simplex/Chat/ByteStringBatcher.hs b/src/Simplex/Chat/ByteStringBatcher.hs new file mode 100644 index 0000000000..d09ed988aa --- /dev/null +++ b/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) diff --git a/src/Simplex/Chat/Messages.hs b/src/Simplex/Chat/Messages.hs index 6b9557e2c0..1ec439a72c 100644 --- a/src/Simplex/Chat/Messages.hs +++ b/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 diff --git a/tests/ByteStringBatcherTests.hs b/tests/ByteStringBatcherTests.hs new file mode 100644 index 0000000000..b2c05ebb95 --- /dev/null +++ b/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 diff --git a/tests/Test.hs b/tests/Test.hs index ee5804aa9a..d181f8a6b8 100644 --- a/tests/Test.hs +++ b/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