make timeouts for batched functions dependent on the number of batches (#786)

* make timeouts for batched functions dependent on the number of batches * fix * refactor * refactor * change import * refactor * rename function * rename * refactor * remove space
2026-08-25 05:20:04 +00:00 · 2023-06-30 12:26:09 +01:00
parent 8be2505fa0
commit 16367fcb3b
2 changed files with 58 additions and 15 deletions
@@ -11,6 +11,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TupleSections #-}

 -- |
 -- Module      : Simplex.Messaging.Client
@@ -82,10 +83,11 @@ 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)
@@ -114,6 +116,9 @@ data ProtocolClient err msg = ProtocolClient
    sessionId :: SessionId,
    sessionTs :: UTCTime,
    thVersion :: Version,
+    timeoutPerBlock :: Int,
+    blockSize :: Int,
+    batch :: Bool,
    client_ :: PClient err msg
  }

@@ -122,6 +127,7 @@ data PClient err msg = PClient
    transportSession :: TransportSession msg,
    transportHost :: TransportHost,
    tcpTimeout :: Int,
+    batchDelay :: Maybe Int,
    pingErrorCount :: TVar Int,
    clientCorrId :: TVar Natural,
    sentCommands :: TMap CorrId (Request err msg),
@@ -168,6 +174,8 @@ data NetworkConfig = NetworkConfig
    tcpConnectTimeout :: Int,
    -- | timeout of protocol commands (microseconds)
    tcpTimeout :: Int,
+    -- | additional timeout per kilobyte (1024 bytes) to be sent
+    tcpTimeoutPerKb :: Int,
    -- | TCP keep-alive options, Nothing to skip enabling keep-alive
    tcpKeepAlive :: Maybe KeepAliveOpts,
    -- | period for SMP ping commands (microseconds, 0 to disable)
@@ -201,6 +209,7 @@ defaultNetworkConfig =
      sessionMode = TSMUser,
      tcpConnectTimeout = 7_500_000,
      tcpTimeout = 5_000_000,
+      tcpTimeoutPerKb = 10_000, -- 10ms, should be less than 130ms to avoid Int overflow on 32 bit systems
      tcpKeepAlive = Just defaultKeepAliveOpts,
      smpPingInterval = 600_000_000, -- 10min
      smpPingCount = 3,
@@ -286,7 +295,7 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
        `catch` \(e :: IOException) -> pure . Left $ PCEIOError e
    Left e -> pure $ Left e
  where
-    NetworkConfig {tcpConnectTimeout, tcpTimeout, smpPingInterval} = networkConfig
+    NetworkConfig {tcpConnectTimeout, tcpTimeout, tcpTimeoutPerKb, smpPingInterval} = networkConfig
    mkProtocolClient :: TransportHost -> STM (PClient err msg)
    mkProtocolClient transportHost = do
      connected <- newTVar False
@@ -301,6 +310,7 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
            transportSession,
            transportHost,
            tcpTimeout,
+            batchDelay,
            pingErrorCount,
            clientCorrId,
            sentCommands,
@@ -334,9 +344,10 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
    client _ c cVar h =
      runExceptT (protocolClientHandshake @err @msg h (keyHash srv) serverVRange) >>= \case
        Left e -> atomically . putTMVar cVar . Left $ PCETransportError e
-        Right th@THandle {sessionId, thVersion} -> do
+        Right th@THandle {sessionId, thVersion, blockSize, batch} -> do
          sessionTs <- getCurrentTime
-          let c' = ProtocolClient {action = Nothing, client_ = c, sessionId, thVersion, sessionTs}
+          let timeoutPerBlock = (blockSize * tcpTimeoutPerKb) `div` 1024
+              c' = ProtocolClient {action = Nothing, client_ = c, sessionId, thVersion, sessionTs, timeoutPerBlock, blockSize, batch}
          atomically $ do
            writeTVar (connected c) True
            putTMVar cVar $ Right c'
@@ -586,32 +597,61 @@ 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))
+
 -- | 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}} cs = do
-  ts <- mapM (runExceptT . mkTransmission c) cs
-  mapM_ (atomically . writeTBQueue sndQ . L.map fst) . L.nonEmpty . rights $ L.toList ts
-  forConcurrently ts $ \case
-    Right (_, r) -> withTimeout c $ atomically $ takeTMVar r
+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
+  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
+        where
+          batchSz' = batchSz + tSz
+          bt' = bt + timeoutPerBlock + fromMaybe 0 batchDelay
+          tSz = tSize t

 -- | 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}} pKey qId cmd = do
+sendProtocolCommand c@ProtocolClient {client_ = PClient {sndQ, tcpTimeout}} 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 (atomically $ takeTMVar r)
+    sendRecv t r = atomically (writeTBQueue sndQ [t]) >> withTimeout c tcpTimeout (atomically $ takeTMVar r)

-withTimeout :: ProtocolClient err msg -> IO (Either (ProtocolClientError err) msg) -> IO (Either (ProtocolClientError err) msg)
-withTimeout ProtocolClient {client_ = PClient {tcpTimeout, pingErrorCount}} a =
-  timeout tcpTimeout a >>= \case
+withTimeout :: ProtocolClient err msg -> Int -> IO (Either (ProtocolClientError err) msg) -> IO (Either (ProtocolClientError err) msg)
+withTimeout ProtocolClient {client_ = PClient {pingErrorCount}} t a = do
+  timeout t a >>= \case
    Just r -> atomically (writeTVar pingErrorCount 0) >> pure r
    _ -> pure $ Left PCEResponseTimeout

-mkTransmission :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> ClientCommand msg -> ExceptT (ProtocolClientError err) IO (SentRawTransmission, TMVar (Response err msg))
+mkTransmission :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> ClientCommand msg -> ExceptT (ProtocolClientError err) IO (PCTransmission err msg)
 mkTransmission ProtocolClient {sessionId, thVersion, client_ = PClient {clientCorrId, sentCommands}} (pKey, qId, cmd) = do
  corrId <- liftIO $ atomically getNextCorrId
  let t = signTransmission $ encodeTransmission thVersion sessionId (corrId, qId, cmd)
@@ -680,6 +680,9 @@ instance SignatureSize (Signature a) where
    SignatureEd25519 _ -> Ed25519.signatureSize
    SignatureEd448 _ -> Ed448.signatureSize

+instance SignatureSize ASignature where
+  signatureSize (ASignature _ s) = signatureSize s
+
 instance SignatureSize APrivateSignKey where
  signatureSize (APrivateSignKey _ k) = signatureSize k