From f3111f45593c9f64c39b279a36c39182751b2eda Mon Sep 17 00:00:00 2001 From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com> Date: Fri, 18 Aug 2023 21:02:47 +0100 Subject: [PATCH] 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 | 1 + src/Simplex/Messaging/Client.hs | 172 ++++++++++++++++++++-------- src/Simplex/Messaging/Protocol.hs | 62 ++++++---- src/Simplex/Messaging/Server.hs | 1 + tests/CoreTests/BatchingTests.hs | 180 ++++++++++++++++++++++++++++++ tests/Test.hs | 2 + tests/XFTPAgent.hs | 44 ++++---- 7 files changed, 377 insertions(+), 85 deletions(-) create mode 100644 tests/CoreTests/BatchingTests.hs diff --git a/simplexmq.cabal b/simplexmq.cabal index 2a3bd54d8..ad211c8e8 100644 --- a/simplexmq.cabal +++ b/simplexmq.cabal @@ -534,6 +534,7 @@ test-suite simplexmq-test AgentTests.SchemaDump AgentTests.SQLiteTests CLITests + CoreTests.BatchingTests CoreTests.CryptoTests CoreTests.EncodingTests CoreTests.ProtocolErrorTests diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs index cbd7be42f..b1f86f550 100644 --- a/src/Simplex/Messaging/Client.hs +++ b/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 diff --git a/src/Simplex/Messaging/Protocol.hs b/src/Simplex/Messaging/Protocol.hs index 31bf0ee4c..2de573349 100644 --- a/src/Simplex/Messaging/Protocol.hs +++ b/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 diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index 5826bcf7b..5651a883f 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/src/Simplex/Messaging/Server.hs @@ -31,6 +31,7 @@ module Simplex.Messaging.Server disconnectTransport, verifyCmdSignature, dummyVerifyCmd, + randomId, ) where diff --git a/tests/CoreTests/BatchingTests.hs b/tests/CoreTests/BatchingTests.hs new file mode 100644 index 000000000..34936d62b --- /dev/null +++ b/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 diff --git a/tests/Test.hs b/tests/Test.hs index c948c5437..21c6453e5 100644 --- a/tests/Test.hs +++ b/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 diff --git a/tests/XFTPAgent.hs b/tests/XFTPAgent.hs index 53a0184f3..abec05c90 100644 --- a/tests/XFTPAgent.hs +++ b/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 =