Allow application to burst a sequence of sdu's

Author: crocodile-dentist

network-mux/bench/socket_read_write/Main.hs

@@ -144,7 +144,7 @@ mkMiniProtocolState num = do
   mpv    <- newTVarIO StatusRunning
 
   let mpi = MiniProtocolInfo (MiniProtocolNum num) InitiatorDirectionOnly
-                             (MiniProtocolLimits maxBound) Nothing
+                             (MiniProtocolLimits maxBound Nothing) Nothing
   return $ MiniProtocolState mpi mpq mpv
 
 -- | Run a server that accept connections on `ad`.
@@ -253,7 +253,7 @@ startServerEgresss pollInterval sndSizeV ad = forever $ do
            let wasEmpty = BL.null buf
            writeTVar w (BL.append buf msg)
            when wasEmpty $
-             writeTBQueue eq (TLSRDemand mc md $ Wanton w)
+             writeTBQueue eq (TLSRDemand mc md (Wanton w) $ ProtocolBurst 1)
          else retry
 
 setupServer :: Socket -> IO Socket.SockAddr

network-mux/demo/mux-demo.hs

@@ -80,7 +80,8 @@ debugTracer = showTracing (Tracer putStrLn_)
 defaultProtocolLimits :: MiniProtocolLimits
 defaultProtocolLimits =
     MiniProtocolLimits {
-      maximumIngressQueue = 64_000
+      maximumIngressQueue = 64_000,
+      burst = Nothing
     }
 
 --

network-mux/demo/mux-leios-demo.hs

@@ -115,7 +115,8 @@ reqrespTracer tag = Tracer $ \case
 defaultProtocolLimits :: MiniProtocolLimits
 defaultProtocolLimits =
     MiniProtocolLimits {
-      maximumIngressQueue = 10_000_000
+      maximumIngressQueue = 10_000_000,
+      burst = Nothing
     }
 
 

network-mux/network-mux.cabal

@@ -67,6 +67,7 @@ library
     statistics-linreg >=0.3 && <0.4,
     strict,
     time >=1.9.1 && <1.16,
+    transformers,
     vector >=0.12 && <0.14,
 
   if os(windows)

network-mux/src/Network/Mux.hs

@@ -64,7 +64,7 @@ import Data.ByteString.Lazy qualified as BL
 import Data.Int (Int64)
 import Data.Map (Map)
 import Data.Map.Strict qualified as Map
-import Data.Maybe (isNothing)
+import Data.Maybe (fromMaybe, isNothing)
 import Data.Monoid.Synchronisation (FirstToFinish (..))
 import Data.Strict.Tuple (pattern (:!:))
 
@@ -76,6 +76,7 @@ import Control.Monad
 import Control.Monad.Class.MonadAsync
 import Control.Monad.Class.MonadFork
 import Control.Monad.Class.MonadThrow
+import Control.Monad.Class.MonadTime.SI (Time (..))
 import Control.Monad.Class.MonadTimer.SI hiding (timeout)
 import Control.Tracer
 
@@ -307,7 +308,8 @@ miniProtocolJob TracersI {
                   miniProtocolInfo =
                     MiniProtocolInfo {
                       miniProtocolNum,
-                      miniProtocolDir
+                      miniProtocolDir,
+                      miniProtocolLimits
                     },
                   miniProtocolIngressQueue,
                   miniProtocolStatusVar
@@ -323,9 +325,11 @@ miniProtocolJob TracersI {
   where
     jobAction = do
       w <- newTVarIO BL.empty
-      let chan = muxChannel channelTracer_ egressQueue (Wanton w)
+      lastSent <- newTVarIO (Time 0)
+      bucket   <- newTVarIO 0
+      let chan = muxChannel channelTracer_ egressQueue (Wanton w lastSent bucket)
                             miniProtocolNum miniProtocolDirEnum
-                            miniProtocolIngressQueue
+                            miniProtocolIngressQueue (burst miniProtocolLimits)
       (result, remainder) <- miniProtocolAction chan
       traceWith tracer_ (TraceTerminating miniProtocolNum miniProtocolDirEnum)
       atomically $ do
@@ -668,8 +672,9 @@ muxChannel
     -> MiniProtocolNum
     -> MiniProtocolDir
     -> IngressQueue m
+    -> Maybe ProtocolBurst
     -> ByteChannel m
-muxChannel tracer egressQueue want@(Wanton w) mc md q =
+muxChannel tracer egressQueue want@(Wanton w _ _) mc md q mBurst =
     Channel { send, recv }
   where
     -- A soft limit on the egress buffer (Wanton) size.
@@ -682,6 +687,8 @@ muxChannel tracer egressQueue want@(Wanton w) mc md q =
     egressSoftBufferLimit :: Int64
     egressSoftBufferLimit = 0x3ffff
 
+    burst = fromMaybe (ProtocolBurst 0 0) mBurst
+
     send :: BL.ByteString -> m ()
     send encoding = do
         -- We send CBOR encoded messages by encoding them into by ByteString
@@ -696,7 +703,7 @@ muxChannel tracer egressQueue want@(Wanton w) mc md q =
                    let wasEmpty = BL.null buf
                    writeTVar w (BL.append buf encoding)
                    when wasEmpty $
-                     writeTBQueue egressQueue (TLSRDemand mc md want)
+                     writeTBQueue egressQueue (TLSRDemand mc md want burst)
                else retry
 
         traceWith tracer $ TraceChannelSendEnd mc

network-mux/src/Network/Mux/Egress.hs

@@ -1,8 +1,12 @@
 {-# LANGUAGE BangPatterns          #-}
+{-# LANGUAGE BlockArguments        #-}
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE MultiWayIf            #-}
 {-# LANGUAGE NamedFieldPuns        #-}
 {-# LANGUAGE RankNTypes            #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TupleSections         #-}
 {-# LANGUAGE TypeFamilies          #-}
 
 module Network.Mux.Egress
@@ -14,8 +18,13 @@ module Network.Mux.Egress
   , Wanton (..)
   ) where
 
+import Control.Exception
 import Control.Monad
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.Except
+import Data.Bool
 import Data.ByteString.Lazy qualified as BL
+import Data.Word (Word32)
 
 import Control.Concurrent.Class.MonadSTM.Strict
 import Control.Monad.Class.MonadAsync
@@ -121,12 +130,16 @@ type EgressQueue m = StrictTBQueue m (TranslocationServiceRequest m)
 --  responsible for the segmentation of concrete representation into
 --  appropriate SDU's for onward transmission.
 data TranslocationServiceRequest m =
-     TLSRDemand !MiniProtocolNum !MiniProtocolDir !(Wanton m)
+     TLSRDemand !MiniProtocolNum !MiniProtocolDir !(Wanton m) !ProtocolBurst
 
 -- | A Wanton represent the concrete data to be translocated, note that the
 --  TVar becoming empty indicates -- that the last fragment of the data has
 --  been enqueued on the -- underlying bearer.
-newtype Wanton m = Wanton { want :: StrictTVar m BL.ByteString }
+data Wanton m = Wanton {
+  want      :: !(StrictTVar m BL.ByteString),
+  wLastSent :: !(StrictTVar m Time),
+  wBucket   :: !(StrictTVar m Word32)
+  }
 
 
 -- | Process the messages from the mini protocols - there is a single
@@ -150,9 +163,18 @@ muxer egressQueue tracer Bearer { writeMany, sduSize, batchSize, egressInterval
     withTimeoutSerial $ \timeout ->
     forever $ do
       start <- getMonotonicTime
-      TLSRDemand mpc md d <- atomically $ readTBQueue egressQueue
-      sdu <- processSingleWanton egressQueue sduSize mpc md d
-      sdus <- buildBatch [sdu] (sduLength sdu)
+      (sdu, mBurst) <- atomically do
+        demand@(TLSRDemand mpc md d (ProtocolBurst pbMaxBytes _pbRefillRate)) <- readTBQueue egressQueue
+        eSdu <- processSingleWanton sduSize mpc md d
+        case eSdu of
+          Right sdu | pbMaxBytes > 0 -> do
+                          -- we do not check if the protocol has any tokens to burst,
+                          -- that is deferred to buildBatch below.
+                          (sdu, True) <$ unGetTBQueue egressQueue demand
+                    | otherwise -> (sdu, False) <$ writeTBQueue egressQueue demand
+          Left sdu -> pure (sdu, False)
+
+      sdus <- buildBatch [sdu] (sduLength sdu) mBurst start
       void $ writeMany tracer timeout sdus
       end <- getMonotonicTime
       empty <- atomically $ isEmptyTBQueue egressQueue
@@ -168,51 +190,109 @@ muxer egressQueue tracer Bearer { writeMany, sduSize, batchSize, egressInterval
     sduLength sdu = fromIntegral msHeaderLength + fromIntegral (msLength sdu)
 
     -- Build a batch of SDUs to submit in one go to the bearer.
-    -- The egress queue is still processed one SDU at the time
-    -- to ensure that we don't cause starvation.
+    -- Streams which are permitted to burst will have that many
+    -- sdu's serviced back-to-back before the scheduler moves to process the
+    -- next request on the queue. Any remaining sdu's which did not
+    -- fit in the burst allowance are placed on the back of the queue
+    -- to ensure that we don't cause starvation. In particular, a burst
+    -- of 1 will have the muxer process one sdu at a time from the queue,
+    -- and any remaining work is put on the back of the queue.
     -- The batch size is either limited by the bearer
     -- (e.g the SO_SNDBUF for Socket) or number of SDUs.
     --
-    buildBatch s sl = reverse <$> go s sl
+    buildBatch s sl mBurst0 start = reverse <$> go 1 s sl mBurst0
      where
-      go sdus _ | length sdus >= maxSDUsPerBatch   = return sdus
-      go sdus sdusLength | sdusLength >= batchSize = return sdus
-      go sdus !sdusLength = do
-        demand_m <- atomically $ tryReadTBQueue egressQueue
-        case demand_m of
-             Just (TLSRDemand mpc md d) -> do
-               sdu <- processSingleWanton egressQueue sduSize mpc md d
-               go (sdu:sdus) (sdusLength + sduLength sdu)
-             Nothing -> return sdus
+      toDouble :: DiffTime -> Double
+      toDouble = realToFrac
+
+      go !count sdus _ _ | count >= maxSDUsPerBatch      = return sdus
+      go _ sdus sdusLength _  | sdusLength >= batchSize  = return sdus
+      go count  sdus !sdusLength mBurst = do
+        mResult <- atomically $ tryReadTBQueue egressQueue
+        case mResult of
+          Nothing -> return sdus
+          Just demand@(TLSRDemand mpc md d@Wanton { wLastSent, wBucket } (ProtocolBurst pbMaxBytes pbRefillRate)) -> do
+            (count', sdusLength', sdus') <- atomically do
+              delta <- (start `diffTime`) <$> stateTVar wLastSent (, start)
+              isEmpty <- isEmptyTBQueue egressQueue
+              sduSize0 <- stateTVar wBucket \tokens ->
+                let tokens' = truncate $ min (fromIntegral pbMaxBytes)
+                                             (fromIntegral tokens + fromIntegral pbRefillRate * toDouble delta)
+                    -- we leverage burst and deduct credits only where there is contention
+                    -- between protocols
+                    sduSize0 = bool (Left sduSize) (Right $ min sduSize (fromIntegral tokens')) (mBurst && not isEmpty)
+                in (sduSize0, tokens')
+              let step (!count', !sdusLength', !sdus', !eSize) mx = do
+                    -- the first one is always free
+                    -- For Left's, we don't count the wanton bytes against the burst allowance
+                    -- to permit a full sdu in the first iteration
+                    let (size, consumedTokens) = either (, const 0) (, id) eSize
+                    x <- lift $ mx size
+                    case x of
+                      Left sdu -> do
+                        lift $ modifyTVar wBucket \tokens ->
+                                 let tokens' = tokens - consumedTokens (fromIntegral (msLength sdu))
+                                 in assert (tokens >= consumedTokens (fromIntegral $ msLength sdu))
+                                    tokens'
+                        let sdusLength'' = sdusLength' + sduLength sdu
+                        throwE (succ count', sdusLength'', sdu:sdus')
+                      Right sdu -> do
+                        nextSdu <- lift $ stateTVar wBucket \tokens ->
+                                       let tokens' = tokens - consumedTokens (fromIntegral (msLength sdu))
+                                           nextSdu = min sduSize (fromIntegral tokens')
+                                       in assert (tokens >= consumedTokens (fromIntegral $ msLength sdu))
+                                          (nextSdu, tokens')
+                        let sdusLength'' = sdusLength' + sduLength sdu
+                            count'' = succ count'
+                        if | nextSdu <= 400 -> do -- 8 bytes header / 2% burst efficiency
+                               -- there is more payload, but burst allowance has been exhausted
+                               lift $ writeTBQueue egressQueue demand
+                               throwE (count'', sdusLength'', sdu:sdus')
+                           | sdusLength'' >= batchSize || count'' >= maxSDUsPerBatch -> do
+                               lift $ unGetTBQueue egressQueue demand
+                               throwE (count'', sdusLength'', sdu:sdus')
+                           | otherwise -> pure (count'', sdusLength'', sdu:sdus', Right nextSdu)
+              either pure (\(a, b, c, _d) -> (a, b, c) <$ writeTBQueue egressQueue demand)
+                =<< runExceptT do
+                      when (either id id sduSize0 <= 400) do
+                        -- edge case where the protocol is bursty, but there aren't enough tokens
+                        -- available. The muxer forever loop does not check this
+                        -- when it calls to build a batch, so we handle it here.
+                        lift $ writeTBQueue egressQueue demand
+                        throwE (count, sdusLength, sdus)
+                      foldM step (count, sdusLength, sdus, sduSize0)
+                                 (if isEmpty
+                                    then [const $ processSingleWanton sduSize mpc md d]
+                                         -- ^ grab full sdu, save the tokens for cases with contention
+                                    else repeat (\sduSize' -> processSingleWanton sduSize' mpc md d))
+            go count' sdus' sdusLength' False
+
 
 -- | Pull a `maxSDU`s worth of data out out the `Wanton` - if there is
 -- data remaining requeue the `TranslocationServiceRequest` (this
 -- ensures that any other items on the queue will get some service
 -- first.
-processSingleWanton :: MonadSTM m
-                    => EgressQueue m
-                    -> SDUSize
+processSingleWanton :: (MonadSTM m)
+                    => SDUSize
                     -> MiniProtocolNum
                     -> MiniProtocolDir
                     -> Wanton m
-                    -> m SDU
-processSingleWanton egressQueue (SDUSize sduSize)
+                       -- Right: more sdu's remain; Left: finished
+                    -> STM m (Either SDU SDU)
+processSingleWanton sduSize
                     mpc md wanton = do
-    blob <- atomically $ do
+    (blob, wrap) <- do
       -- extract next SDU
       d <- readTVar (want wanton)
       let (frag, rest) = BL.splitAt (fromIntegral sduSize) d
       -- if more to process then enqueue remaining work
       if BL.null rest
-        then writeTVar (want wanton) BL.empty
+        then (frag, Left) <$ writeTVar (want wanton) BL.empty
         else do
           -- Note that to preserve bytestream ordering within a given
           -- miniprotocol the readTVar and writeTVar operations
           -- must be inside the same STM transaction.
-          writeTVar (want wanton) rest
-          writeTBQueue egressQueue (TLSRDemand mpc md wanton)
-      -- return data to send
-      pure frag
+          (frag, Right) <$ writeTVar (want wanton) rest
     let sdu = SDU {
                 msHeader = SDUHeader {
                     mhTimestamp = RemoteClockModel 0,
@@ -222,5 +302,5 @@ processSingleWanton egressQueue (SDUSize sduSize)
                   },
                 msBlob = blob
               }
-    return sdu
+    pure $ wrap sdu
     --paceTransmission tNow

network-mux/src/Network/Mux/Types.hs

@@ -25,6 +25,7 @@ module Network.Mux.Types
   , IngressQueue
   , MiniProtocolIx
   , MiniProtocolDir (..)
+  , ProtocolBurst (..)
   , protocolDirEnum
   , MiniProtocolState (..)
   , MiniProtocolStatus (..)
@@ -93,15 +94,23 @@ newtype MiniProtocolNum = MiniProtocolNum Word16
   deriving (Eq, Ord, Enum, Ix, Show)
 
 -- | Per Miniprotocol limits
-newtype MiniProtocolLimits =
+data MiniProtocolLimits =
      MiniProtocolLimits {
        -- | Limit on the maximum number of bytes that can be queued in the
        -- miniprotocol's ingress queue.
        --
-       maximumIngressQueue :: Int
+       maximumIngressQueue :: !Int,
+       burst               :: !(Maybe ProtocolBurst)
      }
   deriving Show
 
+
+data ProtocolBurst = ProtocolBurst {
+  pbMaxBytes   :: !Word32,
+  pbRefillRate :: !Word32
+  }
+  deriving (Eq, Show)
+
 -- $interface
 --
 -- To run a node you will also need a bearer and a way to run a server, see