From e04705d9c5e6b3d3652f909a5176c375acf29411 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com> Date: Wed, 6 Mar 2024 11:10:49 +0200 Subject: [PATCH] utils: add generic batching and compression (#1018) * extract batchTransmissions_ * add Simplex.Messaging.Compression * add combined compression/batching * force NonEmpty for batching * hide FFI and allocation related IO * split packing * remove batch compression, tweak API * OCD over API * remove Empty, extract passthrough const --------- Co-authored-by: Evgeny Poberezkin --- package.yaml | 1 + simplexmq.cabal | 8 +++ src/Simplex/Messaging/Compression.hs | 74 ++++++++++++++++++++++++++++ src/Simplex/Messaging/Protocol.hs | 23 ++++++--- 4 files changed, 98 insertions(+), 8 deletions(-) create mode 100644 src/Simplex/Messaging/Compression.hs diff --git a/package.yaml b/package.yaml index 4dbc971a1..76de72ac6 100644 --- a/package.yaml +++ b/package.yaml @@ -73,6 +73,7 @@ dependencies: - unliftio-core == 0.2.* - websockets == 0.12.* - yaml == 0.11.* + - zstd == 0.1.3.* flags: swift: diff --git a/simplexmq.cabal b/simplexmq.cabal index 9242f6168..35f916cc8 100644 --- a/simplexmq.cabal +++ b/simplexmq.cabal @@ -108,6 +108,7 @@ library Simplex.Messaging.Agent.TRcvQueues Simplex.Messaging.Client Simplex.Messaging.Client.Agent + Simplex.Messaging.Compression Simplex.Messaging.Crypto Simplex.Messaging.Crypto.File Simplex.Messaging.Crypto.Lazy @@ -228,6 +229,7 @@ library , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -300,6 +302,7 @@ executable ntf-server , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -372,6 +375,7 @@ executable smp-agent , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -444,6 +448,7 @@ executable smp-server , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -516,6 +521,7 @@ executable xftp , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -588,6 +594,7 @@ executable xftp-server , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON @@ -701,6 +708,7 @@ test-suite simplexmq-test , unliftio-core ==0.2.* , websockets ==0.12.* , yaml ==0.11.* + , zstd ==0.1.3.* default-language: Haskell2010 if flag(swift) cpp-options: -DswiftJSON diff --git a/src/Simplex/Messaging/Compression.hs b/src/Simplex/Messaging/Compression.hs new file mode 100644 index 000000000..c6664a179 --- /dev/null +++ b/src/Simplex/Messaging/Compression.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Simplex.Messaging.Compression where + +import qualified Codec.Compression.Zstd.FFI as Z +import Control.Monad (forM) +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import qualified Data.ByteString.Unsafe as B +import Data.List.NonEmpty (NonEmpty) +import Foreign +import Foreign.C.Types +import GHC.IO (unsafePerformIO) +import Simplex.Messaging.Encoding +import UnliftIO.Exception (bracket) + +data Compressed + = -- | Short messages are left intact to skip copying and FFI festivities. + Passthrough ByteString + | -- | Generic compression using no extra context. + Compressed Large + +-- | Messages below this length are not encoded to avoid compression overhead. +maxLengthPassthrough :: Int +maxLengthPassthrough = 181 -- Sampled from real client data. Messages with length >=181 rapidly gain compression ratio. + +instance Encoding Compressed where + smpEncode = \case + Passthrough bytes -> "0" <> smpEncode bytes + Compressed bytes -> "1" <> smpEncode bytes + smpP = + smpP >>= \case + '0' -> Passthrough <$> smpP + '1' -> Compressed <$> smpP + x -> fail $ "unknown Compressed tag: " <> show x + +type CompressCtx = (Ptr Z.CCtx, Ptr CChar, Int) + +withCompressCtx :: Int -> (CompressCtx -> IO a) -> IO a +withCompressCtx scratchSize action = + bracket Z.createCCtx Z.freeCCtx $ \cctx -> + allocaBytes scratchSize $ \scratchPtr -> + action (cctx, scratchPtr, scratchSize) + +compress :: CompressCtx -> ByteString -> IO (Either String Compressed) +compress (cctx, scratchPtr, scratchSize) bs + | B.length bs < maxLengthPassthrough = pure . Right $ Passthrough bs + | otherwise = + B.unsafeUseAsCStringLen bs $ \(sourcePtr, sourceSize) -> do + res <- Z.checkError $ Z.compressCCtx cctx scratchPtr (fromIntegral scratchSize) sourcePtr (fromIntegral sourceSize) 3 + case res of + Left e -> pure $ Left e -- should not happen, unless input buffer is too short + Right dstSize -> Right . Compressed . Large <$> B.packCStringLen (scratchPtr, fromIntegral dstSize) + +type DecompressCtx = (Ptr Z.DCtx, Ptr CChar, CSize) + +withDecompressCtx :: Int -> (DecompressCtx -> IO a) -> IO a +withDecompressCtx maxUnpackedSize action = + bracket Z.createDCtx Z.freeDCtx $ \dctx -> + allocaBytes maxUnpackedSize $ \scratchPtr -> + action (dctx, scratchPtr, fromIntegral maxUnpackedSize) + +decompress :: DecompressCtx -> Compressed -> IO (Either String ByteString) +decompress (dctx, scratchPtr, scratchSize) = \case + Passthrough bs -> pure $ Right bs + Compressed (Large bs) -> + B.unsafeUseAsCStringLen bs $ \(sourcePtr, sourceSize) -> do + res <- Z.checkError $ Z.decompressDCtx dctx scratchPtr scratchSize sourcePtr (fromIntegral sourceSize) + forM res $ \dstSize -> B.packCStringLen (scratchPtr, fromIntegral dstSize) + +decompressBatch :: Int -> NonEmpty Compressed -> NonEmpty (Either String ByteString) +decompressBatch maxUnpackedSize items = unsafePerformIO $ withDecompressCtx maxUnpackedSize $ forM items . decompress +{-# NOINLINE decompressBatch #-} -- prevent double-evaluation under unsafePerformIO diff --git a/src/Simplex/Messaging/Protocol.hs b/src/Simplex/Messaging/Protocol.hs index 8f36d8e7a..3aef08622 100644 --- a/src/Simplex/Messaging/Protocol.hs +++ b/src/Simplex/Messaging/Protocol.hs @@ -1,4 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE DeriveAnyClass #-} {-# LANGUAGE DuplicateRecordFields #-} @@ -157,6 +158,7 @@ module Simplex.Messaging.Protocol tEncodeBatch1, batchTransmissions, batchTransmissions', + batchTransmissions_, -- * exports for tests CommandTag (..), @@ -171,6 +173,7 @@ import Data.Aeson (FromJSON (..), ToJSON (..)) import qualified Data.Aeson.TH as J import Data.Attoparsec.ByteString.Char8 (Parser, ()) import qualified Data.Attoparsec.ByteString.Char8 as A +import Data.Bifunctor (first) import qualified Data.ByteString.Base64 as B64 import Data.ByteString.Char8 (ByteString) import qualified Data.ByteString.Char8 as B @@ -1325,11 +1328,11 @@ data TransportBatch r = TBTransmissions ByteString Int [r] | TBTransmission Byte batchTransmissions :: Bool -> Int -> NonEmpty (Either TransportError SentRawTransmission) -> [TransportBatch ()] batchTransmissions batch bSize = batchTransmissions' batch bSize . L.map (,()) --- | encodes and batches transmissions into blocks, +-- | encodes and batches transmissions into blocks batchTransmissions' :: forall r. Bool -> Int -> NonEmpty (Either TransportError SentRawTransmission, r) -> [TransportBatch r] -batchTransmissions' batch bSize - | batch = addBatch . foldr addTransmission ([], 0, 0, [], []) - | otherwise = map mkBatch1 . L.toList +batchTransmissions' batch bSize ts + | batch = batchTransmissions_ bSize $ L.map (first $ fmap tEncodeForBatch) ts + | otherwise = map mkBatch1 $ L.toList ts where mkBatch1 :: (Either TransportError SentRawTransmission, r) -> TransportBatch r mkBatch1 (t_, r) = case t_ of @@ -1340,17 +1343,21 @@ batchTransmissions' batch bSize | otherwise -> TBError TELargeMsg r where s = tEncode t + +-- | Pack encoded transmissions into batches +batchTransmissions_ :: Int -> NonEmpty (Either TransportError ByteString, r) -> [TransportBatch r] +batchTransmissions_ bSize = addBatch . foldr addTransmission ([], 0, 0, [], []) + where -- 3 = 2 bytes reserved for pad size + 1 for transmission count bSize' = bSize - 3 - addTransmission :: (Either TransportError SentRawTransmission, r) -> ([TransportBatch r], Int, Int, [ByteString], [r]) -> ([TransportBatch r], Int, Int, [ByteString], [r]) - addTransmission (t_, r) acc@(bs, len, n, ss, rs) = case t_ of + addTransmission :: (Either TransportError ByteString, r) -> ([TransportBatch r], Int, Int, [ByteString], [r]) -> ([TransportBatch r], Int, Int, [ByteString], [r]) + addTransmission (t_, r) acc@(bs, !len, !n, ss, rs) = case t_ of Left e -> (TBError e r : addBatch acc, 0, 0, [], []) - Right t + Right s | len' <= bSize' && n < 255 -> (bs, len', 1 + n, s : ss, r : rs) | sLen <= bSize' -> (addBatch acc, sLen, 1, [s], [r]) | otherwise -> (TBError TELargeMsg r : addBatch acc, 0, 0, [], []) where - s = tEncodeForBatch t sLen = B.length s len' = len + sLen addBatch :: ([TransportBatch r], Int, Int, [ByteString], [r]) -> [TransportBatch r]