Messages are now interleaved when multiplex-mini-protocols: true

Author: nfrisby

simulation/ouroboros-leios-sim.cabal

@@ -139,6 +139,7 @@ library
     , io-classes
     , io-sim
     , kdt
+    , lens
     , leios-trace-hs
     , linear
     , mtl

simulation/src/Chan.hs

@@ -33,7 +33,7 @@ import Control.Monad.Class.MonadFork (MonadFork)
 import Control.Tracer (Contravariant (contramap), Tracer)
 import Data.Maybe (fromMaybe)
 import GHC.Generics
-import ModelTCP (kibibytes, mkTcpConnProps)
+import ModelTCP (kibibytes, lensTcpEvent, mkTcpConnProps)
 import TimeCompat (DiffTime, MonadDelay, MonadMonotonicTimeNSec, MonadTime)
 
 data ConnectionConfig = ConnectionConfig
@@ -64,14 +64,14 @@ newConnectionBundle tracer = \case
     newUnMuxedConnectionBundle $ \tofrom ->
       newConnectionSimple (traceAsBundleMsg @bundle tofrom tracer) simpleConnProps
   ConnectionConfig (TransportSimple simpleConnProps) _mux@True -> do
-    let tracer' = contramap ((fmap . fmap) fromBearerMsg) tracer
+    let tracer' = mapBearerTracer (lensLabelTcpDir . lensTcpEvent) tracer
     (mA, mB) <- newConnectionSimple tracer' simpleConnProps
     (,) <$> newMuxChan mA <*> newMuxChan mB
   ConnectionConfig (TransportTcp tcpConnProps) _mux@False ->
     newUnMuxedConnectionBundle $ \tofrom ->
       newConnectionTCP (traceAsBundleMsg @bundle tofrom tracer) tcpConnProps
   ConnectionConfig (TransportTcp tcpConnProps) _mux@True -> do
-    let tracer' = contramap ((fmap . fmap) fromBearerMsg) tracer
+    let tracer' = mapBearerTracer (lensLabelTcpDir . lensTcpEvent) tracer
     (mA, mB) <- newConnectionTCP tracer' tcpConnProps
     (,) <$> newMuxChan mA <*> newMuxChan mB
 

simulation/src/Chan/Mux.hs

@@ -5,6 +5,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTSyntax #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -14,18 +15,22 @@
 module Chan.Mux (
   ToFromBundleMsg (..),
   ConnectionBundle (..),
-  fromBearerMsg,
+  mapBearerTracer,
   newMuxChan,
 ) where
 
 import Chan.Core (Chan (..))
-import Chan.TCP (MessageSize (..))
+import Chan.TCP (Bytes, MessageSize (..))
 import qualified Control.Category as Cat
 import Control.Concurrent.Class.MonadMVar (MonadMVar (..))
-import Control.Monad (forever)
+import qualified Control.Lens as Lens
+import Control.Monad (forM_, forever)
 import Control.Monad.Class.MonadFork (MonadFork (forkIO))
+import Control.Tracer (Tracer (Tracer), traceWith)
 import Data.Array (Array, listArray, (!))
 import Data.Kind
+import Data.Foldable (traverse_)
+import Data.Functor.Const (Const (Const), getConst)
 import STMCompat
 
 class ConnectionBundle bundle where
@@ -56,66 +61,86 @@ instance Cat.Category ToFromBundleMsg where
 -- dynToFromBundleMsg :: Typeable a => ToFromBundleMsg Dynamic a
 -- dynToFromBundleMsg = ToFromBundleMsg toDyn (fromJust . fromDynamic)
 
-data BearerMsg a = BearerMsg !Int a
+data BorneMsg a = BorneMsg !Int a
 
-fromBearerMsg :: BearerMsg a -> a
-fromBearerMsg (BearerMsg _ a) = a
-
-instance MessageSize a => MessageSize (BearerMsg a) where
-  messageSizeBytes (BearerMsg _ a) = 1 + messageSizeBytes a
+-- | Each bearer message is some slices of various 'BorneMsg's
+--
+-- The mini protocols never see this, so this type is not exported. It does
+-- occur in the argument types of some exported functions, but the caller
+-- should be using parametric functions to generate those arguments.
+data BearerMsg a = BearerMsg !Bytes [BorneMsg a]
+  -- ^ the cumulative size of the slices the borne messages whose /final/ slice
+  -- is in this message
+
+instance MessageSize (BearerMsg a) where
+  messageSizeBytes (BearerMsg sz _) = 1 + sz
+
+mapBearerTracer ::
+  Applicative m =>
+  Lens.Lens s t (BearerMsg a) a ->
+  Tracer m t ->
+  Tracer m s
+mapBearerTracer lens tracer = Tracer $ \x -> do
+  let BearerMsg _ msgs = getConst $ lens Const x -- why doesn't Lens.view lens x type check?
+  flip traverse_ msgs $ \(BorneMsg _ a) -> do
+    traceWith tracer $ Lens.set lens a x
 
 newMuxChan ::
   forall bundle m.
-  (ConnectionBundle bundle, MonadMVar m, MonadSTM m, MonadFork m) =>
+  (ConnectionBundle bundle, MonadMVar m, MonadSTM m, MonadFork m, MessageSize (BundleMsg bundle)) =>
   Chan m (BearerMsg (BundleMsg bundle)) ->
   m (bundle (Chan m))
 newMuxChan bearer = do
-  sendLock <- newMVar ()
   -- Bit of a hack to use these TVars, could run the traverseConnectionBundle
   -- in a reader+state monad instead. That'd be cleaner.
   recvQueuesAccum <- newTVarIO []
   recvQueuesIx <- newTVarIO (0 :: Int)
+  sendQueue <- newTQueueIO
   chans <-
     traverseConnectionBundle
       ( newMuxChanSingle @bundle
-          bearer
-          sendLock
+          sendQueue
           recvQueuesIx
           recvQueuesAccum
       )
       toFromBundleMsgBundle
   recvQueues <- reverse <$> readTVarIO recvQueuesAccum
   let recvQueues' = listArray (0, length recvQueues - 1) recvQueues
   _ <- forkIO $ demuxer @bundle bearer recvQueues'
+  _ <- forkIO $ muxer @bundle bearer sendQueue
   return chans
 
 newMuxChanSingle ::
   forall bundle m a.
-  (MonadMVar m, MonadSTM m) =>
-  Chan m (BearerMsg (BundleMsg bundle)) ->
-  MVar m () ->
+  (MonadMVar m, MonadSTM m, MessageSize (BundleMsg bundle)) =>
+  TQueue m (MVar m (), Bytes, BorneMsg (BundleMsg bundle)) ->
   TVar m Int ->
   TVar m [RecvQueue m (BundleMsg bundle)] ->
   ToFromBundleMsg (BundleMsg bundle) a ->
   m (Chan m a)
 newMuxChanSingle
-  bearer
-  sendLock
+  sendQueue
   recvQueuesIx
   recvQueuesAccum
   ToFromBundleMsg{..} = do
-    queue <- newTQueueIO
+    recvQueue <- newTQueueIO
+    -- A mini protocol can have at most one message in the send buffer.
+    sendLock <- newMVar ()
     i <- atomically $ do
-      modifyTVar recvQueuesAccum (RecvQueue fromBundleMsg queue :)
+      modifyTVar recvQueuesAccum (RecvQueue fromBundleMsg recvQueue :)
       i <- readTVar recvQueuesIx
       writeTVar recvQueuesIx $! (i + 1)
       return i
     return
       Chan
-        { readChan = atomically (readTQueue queue)
-        , writeChan = \msg ->
-            let !muxmsg = BearerMsg i (toBundleMsg msg)
-             in withMVar sendLock $ \_ -> writeChan bearer muxmsg
+        { readChan = atomically (readTQueue recvQueue)
+        , writeChan = \msg -> do
+            let !bundleMsg = toBundleMsg msg
+                !muxmsg = BorneMsg i bundleMsg
+            takeMVar sendLock
+            atomically $
+              writeTQueue sendQueue $
+              (sendLock, messageSizeBytes bundleMsg, muxmsg)
         }
 
 data RecvQueue m mm where
@@ -129,10 +154,49 @@ demuxer ::
   m ()
 demuxer bearer queues =
   forever $ do
-    BearerMsg i msg <- readChan bearer
-    case queues ! i of
-      RecvQueue convert queue ->
-        atomically $ writeTQueue queue $! convert msg
+    BearerMsg _ msgs <- readChan bearer
+    forM_ msgs $ \(BorneMsg i msg) ->
+      case queues ! i of
+        RecvQueue convert queue ->
+          atomically $ writeTQueue queue $! convert msg
+
+muxer ::
+  forall bundle m.
+  (MonadMVar m, MonadSTM m) =>
+  Chan m (BearerMsg (BundleMsg bundle)) ->
+  TQueue m (MVar m (), Bytes, BorneMsg (BundleMsg bundle)) ->
+  m ()
+muxer bearer sendQueue =
+    forever $ do
+      x <- atomically (readTQueue sendQueue)
+      (muxmsg, locks) <- go 0 [] [] x
+      mapM_ (flip putMVar ()) locks
+      writeChan bearer muxmsg
+  where
+    --- from ouroboros-network's @Network.Mux.Bearer.makeSocketBearer'@
+    sliceBytes = 12288
+    loafBytes = 131072
+
+    go !accBytes acc locks (lock, bytes, msg) = do
+      let !accBytes' = accBytes + min sliceBytes bytes
+      (acc', locks') <-
+        if bytes <= sliceBytes
+        then do
+          -- We do not release the lock before finalizing the loaf because a
+          -- single loaf should include slices from at most one borne message
+          -- per protocol.
+          pure (msg : acc, lock : locks)
+        else do
+          -- reenqueue the rest of the message
+          let !bytes' = bytes - sliceBytes
+          atomically $ writeTQueue sendQueue (lock, bytes', msg)
+          pure (acc, locks)
+
+      let result = (BearerMsg accBytes' acc', locks')
+      if accBytes' >= loafBytes then pure result else do
+        atomically (tryReadTQueue sendQueue) >>= \case
+          Nothing -> pure result
+          Just x -> go accBytes' acc' locks' x
 
 data ExampleBundle f = ExampleBundle
   { exampleFoo :: f Int

simulation/src/Chan/TCP.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
@@ -35,6 +36,11 @@ import TimeCompat
 data LabelTcpDir e = DirClientToServer e | DirServerToClient e
   deriving (Eq, Ord, Show, Functor)
 
+lensLabelTcpDir :: Functor f => (a -> f b) -> LabelTcpDir a -> f (LabelTcpDir b)
+lensLabelTcpDir f = \case
+    DirClientToServer x -> DirClientToServer <$> f x
+    DirServerToClient x -> DirServerToClient <$> f x
+
 -- | Class for messages to be sent over a simulated TCP connection.
 -- To correctly model the timing of the messages sent over the connection we
 -- need to know a reasonable approximation of the message size. This does not

simulation/src/ModelTCP.hs

@@ -16,6 +16,7 @@ module ModelTCP (
   TcpMsgForecast (..),
   forecastTcpMsgSend,
   TcpEvent (..),
+  lensTcpEvent,
   traceTcpSend,
   mkTcpConnProps,
   kibibytes,
@@ -366,6 +367,9 @@ data TcpEvent msg
       [TcpMsgForecast] -- tcp internal activity
   deriving (Show, Functor)
 
+lensTcpEvent :: Functor f => (a -> f b) -> TcpEvent a -> f (TcpEvent b)
+lensTcpEvent f (TcpSendMsg x y z) = (\x' -> TcpSendMsg x' y z) <$> f x
+
 traceTcpSend ::
   TcpConnProps ->
   -- | message sizes to send eagerly, back-to-back