Implement an optional read buffer for mux bearers

This implements an optional read buffer for the Socket bearer.

Author: karknu

cardano-ping/src/Cardano/Network/Ping.hs

@@ -674,7 +674,7 @@ pingClient stdout stderr PingOpts{..} versions peer = bracket
     let peerStr' = TL.pack peerStr
     unless pingOptsQuiet $ TL.hPutStrLn IO.stdout $ peerStr' <> " " <> (showNetworkRtt $ toSample t0_e t0_s)
 
-    bearer <- getBearer makeSocketBearer sduTimeout nullTracer sd
+    bearer <- getBearer makeSocketBearer sduTimeout nullTracer sd Nothing
 
     !t1_s <- write bearer timeoutfn $ wrap handshakeNum InitiatorDir (handshakeReq versions pingOptsHandshakeQuery)
     (msg, !t1_e) <- nextMsg bearer timeoutfn handshakeNum

network-mux/bench/socket_read_write/Main.hs

@@ -48,10 +48,12 @@ readBenchmark sndSizeV sndSize addr = do
     (\sd -> do
       atomically $ putTMVar sndSizeV sndSize
       Socket.connect sd addr
-      bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd
+      withReadBufferIO (\buffer -> do
+        bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd buffer
 
-      let chan = bearerAsChannel bearer (MiniProtocolNum 42) InitiatorDir
-      doRead (totalPayloadLen sndSize) chan 0
+        let chan = bearerAsChannel bearer (MiniProtocolNum 42) InitiatorDir
+        doRead (totalPayloadLen sndSize) chan 0
+       )
     )
  where
    doRead :: Int64 -> ByteChannel IO -> Int64 -> IO ()
@@ -72,15 +74,17 @@ readDemuxerBenchmark sndSizeV sndSize addr = do
       atomically $ putTMVar sndSizeV sndSize
 
       Socket.connect sd addr
-      bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd
-      ms42 <- mkMiniProtocolState 42
-      ms41 <- mkMiniProtocolState 41
-      withAsync (demuxer [ms41, ms42] bearer) $ \aid -> do
-        withAsync (doRead 42 (totalPayloadLen sndSize) (miniProtocolIngressQueue ms42) 0) $ \aid42 -> do
-          withAsync (doRead 41 (totalPayloadLen 10) (miniProtocolIngressQueue ms41) 0) $ \aid41 -> do
-            _ <- waitBoth aid42 aid41
-            cancel aid
-            return ()
+      withReadBufferIO (\buffer -> do
+        bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd buffer
+        ms42 <- mkMiniProtocolState 42
+        ms41 <- mkMiniProtocolState 41
+        withAsync (demuxer [ms41, ms42] bearer) $ \aid -> do
+          withAsync (doRead 42 (totalPayloadLen sndSize) (miniProtocolIngressQueue ms42) 0) $ \aid42 -> do
+            withAsync (doRead 41 (totalPayloadLen 10) (miniProtocolIngressQueue ms41) 0) $ \aid41 -> do
+              _ <- waitBoth aid42 aid41
+              cancel aid
+              return ()
+       )
     )
  where
 
@@ -111,37 +115,39 @@ readDemuxerBenchmark sndSizeV sndSize addr = do
 startServer :: StrictTMVar IO Int64 -> Socket -> IO ()
 startServer sndSizeV ad = forever $ do
     (sd, _) <- Socket.accept ad
-    bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd
-    sndSize <- atomically $ takeTMVar sndSizeV
-
-    let chan = bearerAsChannel bearer (MiniProtocolNum 42) ResponderDir
-        payload = BL.replicate sndSize 0xa5
-        maxData = totalPayloadLen sndSize
-        numberOfSdus = fromIntegral $ maxData `div` sndSize
-    replicateM_ numberOfSdus $ do
-      send chan payload
-
+    withReadBufferIO (\buffer -> do
+      bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd buffer
+      sndSize <- atomically $ takeTMVar sndSizeV
+
+      let chan = bearerAsChannel bearer (MiniProtocolNum 42) ResponderDir
+          payload = BL.replicate sndSize 0xa5
+          maxData = totalPayloadLen sndSize
+          numberOfSdus = fromIntegral $ maxData `div` sndSize
+      replicateM_ numberOfSdus $ do
+        send chan payload
+     )
 -- | Like startServer but it uses the `writeMany` function
 -- for vector IO.
 startServerMany :: StrictTMVar IO Int64 -> Socket -> IO ()
 startServerMany sndSizeV ad = forever $ do
     (sd, _) <- Socket.accept ad
-    bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd
-    sndSize <- atomically $ takeTMVar sndSizeV
-
-    let maxData = totalPayloadLen sndSize
-        numberOfSdus = fromIntegral $ maxData `div` sndSize
-        numberOfCalls = numberOfSdus `div` 10
-        runtSdus = numberOfSdus `mod` 10
-
-    withTimeoutSerial $ \timeoutFn -> do
-      replicateM_ numberOfCalls $ do
-        let sdus = replicate 10 $ wrap $ BL.replicate sndSize 0xa5
-        void $ writeMany bearer timeoutFn sdus
-      when (runtSdus > 0) $ do
-        let sdus = replicate runtSdus $ wrap $ BL.replicate sndSize 0xa5
-        void $ writeMany bearer timeoutFn sdus
-
+    withReadBufferIO (\buffer -> do
+      bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd buffer
+      sndSize <- atomically $ takeTMVar sndSizeV
+
+      let maxData = totalPayloadLen sndSize
+          numberOfSdus = fromIntegral $ maxData `div` sndSize
+          numberOfCalls = numberOfSdus `div` 10
+          runtSdus = numberOfSdus `mod` 10
+
+      withTimeoutSerial $ \timeoutFn -> do
+        replicateM_ numberOfCalls $ do
+          let sdus = replicate 10 $ wrap $ BL.replicate sndSize 0xa5
+          void $ writeMany bearer timeoutFn sdus
+        when (runtSdus > 0) $ do
+          let sdus = replicate runtSdus $ wrap $ BL.replicate sndSize 0xa5
+          void $ writeMany bearer timeoutFn sdus
+     )
  where
   -- wrap a 'ByteString' as 'SDU'
   wrap :: BL.ByteString -> SDU
@@ -163,41 +169,43 @@ startServerMany sndSizeV ad = forever $ do
 startServerEgresss :: StrictTMVar IO Int64 -> Socket -> IO ()
 startServerEgresss sndSizeV ad = forever $ do
     (sd, _) <- Socket.accept ad
-    bearer <- getBearer makeSocketBearer sduTimeout activeTracer sd
-    sndSize <- atomically $ takeTMVar sndSizeV
-    eq <- atomically $ newTBQueue 100
-    w42 <- newTVarIO BL.empty
-    w41 <- newTVarIO BL.empty
-
-    let maxData = totalPayloadLen sndSize
-        numberOfSdus = fromIntegral $ maxData `div` sndSize
-        numberOfCalls = numberOfSdus `div` 10 :: Int
-        runtSdus = numberOfSdus `mod` 10 :: Int
-
-    withAsync (muxer eq bearer) $ \aid -> do
-
-      replicateM_ numberOfCalls $ do
-        let payload42s = replicate 10 $ BL.replicate sndSize 42
-        let payload41s = replicate 10 $ BL.replicate 10 41
-        mapM_ (sendToMux w42 eq (MiniProtocolNum 42) ResponderDir) payload42s
-        mapM_ (sendToMux w41 eq (MiniProtocolNum 41) ResponderDir) payload41s
-      when (runtSdus > 0) $ do
-        let payload42s = replicate runtSdus $ BL.replicate sndSize 42
-        let payload41s = replicate runtSdus $ BL.replicate 10 41
-        mapM_ (sendToMux w42 eq (MiniProtocolNum 42) ResponderDir) payload42s
-        mapM_ (sendToMux w41 eq (MiniProtocolNum 41) ResponderDir) payload41s
-
-      -- Wait for the egress queue to empty
-      atomically $ do
-        r42 <- readTVar w42
-        r41 <- readTVar w42
-        unless (BL.null r42 || BL.null r41) retry
-
-      -- when the client is done they will close the socket
-      -- and we will read zero bytes.
-      _ <- Socket.recv sd 128
-
-      cancel aid
+    withReadBufferIO (\buffer -> do
+      bearer <-getBearer makeSocketBearer sduTimeout activeTracer sd buffer
+      sndSize <- atomically $ takeTMVar sndSizeV
+      eq <- atomically $ newTBQueue 100
+      w42 <- newTVarIO BL.empty
+      w41 <- newTVarIO BL.empty
+
+      let maxData = totalPayloadLen sndSize
+          numberOfSdus = fromIntegral $ maxData `div` sndSize
+          numberOfCalls = numberOfSdus `div` 10 :: Int
+          runtSdus = numberOfSdus `mod` 10 :: Int
+
+      withAsync (muxer eq bearer) $ \aid -> do
+
+        replicateM_ numberOfCalls $ do
+          let payload42s = replicate 10 $ BL.replicate sndSize 42
+          let payload41s = replicate 10 $ BL.replicate 10 41
+          mapM_ (sendToMux w42 eq (MiniProtocolNum 42) ResponderDir) payload42s
+          mapM_ (sendToMux w41 eq (MiniProtocolNum 41) ResponderDir) payload41s
+        when (runtSdus > 0) $ do
+          let payload42s = replicate runtSdus $ BL.replicate sndSize 42
+          let payload41s = replicate runtSdus $ BL.replicate 10 41
+          mapM_ (sendToMux w42 eq (MiniProtocolNum 42) ResponderDir) payload42s
+          mapM_ (sendToMux w41 eq (MiniProtocolNum 41) ResponderDir) payload41s
+
+        -- Wait for the egress queue to empty
+        atomically $ do
+          r42 <- readTVar w42
+          r41 <- readTVar w42
+          unless (BL.null r42 || BL.null r41) retry
+
+        -- when the client is done they will close the socket
+        -- and we will read zero bytes.
+        _ <- Socket.recv sd 128
+
+        cancel aid
+     )
   where
     sendToMux :: StrictTVar IO BL.ByteString -> EgressQueue IO -> MiniProtocolNum -> MiniProtocolDir
               -> BL.ByteString -> IO ()

network-mux/demo/mux-demo.hs

@@ -101,7 +101,7 @@ server =
     associateWithIOManager ioManager (Left hpipe)
     Win32.Async.connectNamedPipe hpipe
     void $ forkIO $ do
-      bearer <- getBearer Mx.makeNamedPipeBearer (-1) nullTracer hpipe
+      bearer <- getBearer Mx.makeNamedPipeBearer (-1) nullTracer hpipe Nothing
       serverWorker bearer
         `finally` closeHandle hpipe
 #else
@@ -113,7 +113,7 @@ server = do
     forever $ do
       (sock', _addr) <- Socket.accept sock
       void $ forkIO $ do
-        bearer <- getBearer Mx.makeSocketBearer 1.0 nullTracer sock'
+        bearer <- getBearer Mx.makeSocketBearer 1.0 nullTracer sock' Nothing
         serverWorker bearer
           `finally` Socket.close sock'
 #endif
@@ -167,13 +167,13 @@ client n msg =
                         fILE_FLAG_OVERLAPPED
                         Nothing
     associateWithIOManager ioManager (Left hpipe)
-    bearer <- getBearer Mx.makeNamedPipeBearer (-1) nullTracer hpipe
+    bearer <- getBearer Mx.makeNamedPipeBearer (-1) nullTracer hpipe Nothing
     clientWorker bearer n msg
 #else
 client n msg = do
     sock <- Socket.socket AF_UNIX Socket.Stream Socket.defaultProtocol
     Socket.connect sock (SockAddrUnix pipeName)
-    bearer <- getBearer Mx.makeSocketBearer 1.0 nullTracer sock
+    bearer <- getBearer Mx.makeSocketBearer 1.0 nullTracer sock Nothing
     clientWorker bearer n msg
 #endif
 

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

@@ -15,6 +15,7 @@ module Network.Mux.Bearer
 #if defined(mingw32_HOST_OS)
   , makeNamedPipeBearer
 #endif
+  , withReadBufferIO
   ) where
 
 import           Control.Monad.Class.MonadSTM
@@ -28,6 +29,7 @@ import           Network.Socket (getSocketOption, SocketOption (..), Socket)
 #if defined(mingw32_HOST_OS)
 import           System.Win32 (HANDLE)
 #endif
+import           Foreign.Marshal.Alloc
 
 import           Network.Mux.Bearer.Pipe
 import           Network.Mux.Bearer.Queues
@@ -47,26 +49,34 @@ newtype MakeBearer m fd = MakeBearer {
       -- tracer
       -> fd
       -- file descriptor
+      -> Maybe (ReadBuffer m)
+      -- Optional Readbuffer
       -> m (Bearer m)
   }
 
-
 pureBearer :: Applicative m
-           => (DiffTime -> Tracer m Trace -> fd ->    Bearer m)
-           ->  DiffTime -> Tracer m Trace -> fd -> m (Bearer m)
-pureBearer f = \sduTimeout tr fd -> pure (f sduTimeout tr fd)
+           => (DiffTime -> Tracer m Trace -> fd -> Maybe (ReadBuffer m) ->   Bearer m)
+           ->  DiffTime -> Tracer m Trace -> fd -> Maybe (ReadBuffer m) -> m (Bearer m)
+pureBearer f = \sduTimeout rb tr fd -> pure (f sduTimeout rb tr fd)
+
 
 makeSocketBearer :: MakeBearer IO Socket
-makeSocketBearer = MakeBearer $ (\sduTimeout tr fd -> do
-    readBuffer <- newTVarIO BL.empty
+makeSocketBearer = MakeBearer $ (\sduTimeout tr fd rb -> do
     batch <- getSocketOption fd SendBuffer
-    return $ socketAsBearer size batch readBuffer bufSize sduTimeout tr fd)
+    return $ socketAsBearer size batch rb sduTimeout tr fd)
   where
     size = SDUSize 12_288
-    bufSize = 16_384
+
+withReadBufferIO :: (Maybe (ReadBuffer IO) -> IO b)
+                 -> IO b
+withReadBufferIO f = allocaBytesAligned size 8 $ \ptr -> do
+    v <- atomically $ newTVar BL.empty
+    f $ Just $ ReadBuffer v ptr size
+  where
+    size = 131_072
 
 makePipeChannelBearer :: MakeBearer IO PipeChannel
-makePipeChannelBearer = MakeBearer $ pureBearer (\_ -> pipeAsBearer size)
+makePipeChannelBearer = MakeBearer $ pureBearer (\_ tr fd _ -> pipeAsBearer size tr fd)
   where
     size = SDUSize 32_768
 
@@ -75,13 +85,13 @@ makeQueueChannelBearer :: ( MonadSTM   m
                           , MonadThrow m
                           )
                        => MakeBearer m (QueueChannel m)
-makeQueueChannelBearer = MakeBearer $ pureBearer (\_ -> queueChannelAsBearer size)
+makeQueueChannelBearer = MakeBearer $ pureBearer (\_ tr q _-> queueChannelAsBearer size tr q)
   where
     size = SDUSize 1_280
 
 #if defined(mingw32_HOST_OS)
 makeNamedPipeBearer :: MakeBearer IO HANDLE
-makeNamedPipeBearer = MakeBearer $ pureBearer (\_ -> namedPipeAsBearer size)
+MakeBearer $ pureBearer (\_ _ -> namedPipeAsBearer size)
   where
     size = SDUSize 24_576
 #endif