feat: add streamReadChunk to StreamIO to unblock perf download throughput

src/LibP2P/MultistreamSelect/Negotiation.hs

@@ -8,14 +8,14 @@ module LibP2P.MultistreamSelect.Negotiation
   , StreamIO (..)
   , negotiateInitiator
   , negotiateResponder
+  , mkByteStreamIO
   , mkMemoryStreamPair
   , readExactBounded
   , closeQuietly
   ) where
 
 import Control.Concurrent.STM
 import Control.Exception (IOException, SomeException, catch)
-import Control.Monad (replicateM)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as BS
 import Data.Text (Text)
@@ -42,9 +42,28 @@ data NegotiationResult
 
 -- | Abstraction for stream I/O to enable testing with in-memory buffers.
 data StreamIO = StreamIO
-  { streamWrite    :: ByteString -> IO ()
-  , streamReadByte :: IO Word8   -- ^ Read exactly one byte (blocks until available)
-  , streamClose    :: IO ()      -- ^ Close/half-close the stream (signals EOF to remote)
+  { streamWrite     :: ByteString -> IO ()
+  , streamReadByte  :: IO Word8   -- ^ Read exactly one byte (blocks until available)
+  , streamReadChunk :: Int -> IO ByteString
+    -- ^ Read between 1 and @n@ bytes (@n >= 1@): whatever is already
+    -- buffered or arrives next, without waiting for the full @n@.
+    -- Blocks until at least one byte is available and never returns an
+    -- empty ByteString; EOF and failures surface as 'IOException',
+    -- exactly like 'streamReadByte'. Bulk readers use this to move
+    -- data at chunk granularity instead of byte-at-a-time (#276).
+  , streamClose     :: IO ()      -- ^ Close/half-close the stream (signals EOF to remote)
+  }
+
+-- | Build a 'StreamIO' from byte-level primitives: 'streamReadChunk'
+-- falls back to one byte per call. Correct for any consumer (chunk
+-- reads promise at least one byte, not @n@), just not fast — intended
+-- for tests and mocks built on byte queues.
+mkByteStreamIO :: (ByteString -> IO ()) -> IO Word8 -> IO () -> StreamIO
+mkByteStreamIO write readByte close = StreamIO
+  { streamWrite     = write
+  , streamReadByte  = readByte
+  , streamReadChunk = \_ -> BS.singleton <$> readByte
+  , streamClose     = close
   }
 
 -- | Create an in-memory stream pair for testing using STM TQueue.
@@ -55,13 +74,28 @@ mkMemoryStreamPair = do
   queueBtoA <- newTQueueIO :: IO (TQueue Word8)
   let writeToQueue q bs = mapM_ (atomically . writeTQueue q) (BS.unpack bs)
       readFromQueue q = atomically (readTQueue q)
+      -- Chunk read: block for the first byte, then drain whatever else
+      -- is already queued (up to the requested length) in the same
+      -- transaction.
+      drainUpTo q k
+        | k <= (0 :: Int) = pure []
+        | otherwise = do
+            mb <- tryReadTQueue q
+            case mb of
+              Nothing -> pure []
+              Just b  -> (b :) <$> drainUpTo q (k - 1)
+      readChunkFromQueue q n = atomically $ do
+        b <- readTQueue q
+        rest <- drainUpTo q (n - 1)
+        pure (BS.pack (b : rest))
   pure
-    ( StreamIO (writeToQueue queueAtoB) (readFromQueue queueBtoA) (pure ())
-    , StreamIO (writeToQueue queueBtoA) (readFromQueue queueAtoB) (pure ())
+    ( StreamIO (writeToQueue queueAtoB) (readFromQueue queueBtoA) (readChunkFromQueue queueBtoA) (pure ())
+    , StreamIO (writeToQueue queueBtoA) (readFromQueue queueAtoB) (readChunkFromQueue queueAtoB) (pure ())
     )
 
--- | Chunk size for 'readExactBounded'. Bounds the transient boxed-list
--- allocation per read step regardless of the requested length.
+-- | Maximum bytes requested per 'streamReadChunk' call in
+-- 'readExactBounded'. Bounds transient allocation per read step
+-- regardless of the requested length.
 readChunkSize :: Int
 readChunkSize = 32768
 
@@ -71,8 +105,10 @@ readChunkSize = 32768
 -- #169): the declared length is validated against the caller's
 -- protocol-defined cap before a single byte is read or allocated, so a
 -- hostile length prefix cannot trigger an unbounded allocation. Bytes
--- are accumulated in chunks of at most 'readChunkSize', keeping
--- transient memory use proportional to the chunk size, not to @n@.
+-- are read via 'streamReadChunk' in requests of at most
+-- 'readChunkSize', keeping transient memory use proportional to the
+-- chunk size, not to @n@. A chunk request never exceeds the bytes
+-- still owed, so no byte beyond @n@ is consumed from the stream.
 --
 -- I/O failures during the read (stream reset, EOF) are returned as
 -- 'Left' instead of propagating as 'IOException's.
@@ -93,11 +129,13 @@ readExactBounded stream maxLen n
         pure (Left ("readExactBounded: read failed: " <> show e))
   where
     go :: Int -> IO [ByteString]
-    go 0 = pure []
-    go remaining = do
-      let m = min readChunkSize remaining
-      chunk <- BS.pack <$> replicateM m (streamReadByte stream)
-      (chunk :) <$> go (remaining - m)
+    go remaining
+      | remaining <= 0 = pure []
+      | otherwise = do
+          chunk <- streamReadChunk stream (min readChunkSize remaining)
+          if BS.null chunk
+            then fail "readExactBounded: streamReadChunk returned no bytes"
+            else (chunk :) <$> go (remaining - BS.length chunk)
 
 -- | Close a stream, swallowing any exception (best-effort EOF signal).
 -- Shared by protocol handlers that must release a stream on every exit

src/LibP2P/NAT/Relay.hs

@@ -312,20 +312,22 @@ bridgeStreams mLimit streamA streamB = do
       streamClose streamB
 
 -- | Forward bytes from source to destination with a byte limit.
--- The limit is checked before each read, so the circuit terminates as soon
--- as exactly @limit@ bytes have been forwarded — no byte beyond the limit
--- is consumed from the source.
+-- Data moves at chunk granularity ('streamReadChunk'), but a chunk
+-- request never exceeds the bytes still allowed, so the circuit
+-- terminates as soon as exactly @limit@ bytes have been forwarded —
+-- no byte beyond the limit is consumed from the source.
 forwardWithLimit :: StreamIO -> StreamIO -> IORef Int -> Int -> IO ()
 forwardWithLimit src dst countRef limit = go
   where
+    forwardChunkSize = 32768
     go = do
       count <- readIORef countRef
       if count >= limit
         then pure ()  -- limit reached, stop forwarding
         else do
-          b <- streamReadByte src
-          modifyIORef' countRef (+ 1)
-          streamWrite dst (BS.singleton b)
+          chunk <- streamReadChunk src (min forwardChunkSize (limit - count))
+          modifyIORef' countRef (+ BS.length chunk)
+          streamWrite dst chunk
           go
 
 -- | Build a relay multiaddr in binary format.

src/LibP2P/Protocol/Perf.hs

@@ -84,19 +84,25 @@ writeZeros stream = go
       go (n - chunk)
 
 -- | Read and discard bytes until EOF (the initiator's half-close).
+-- Chunk-level reads (#276) keep the drain off the byte-at-a-time path
+-- that bounded download throughput.
 drainUntilEof :: StreamIO -> IO ()
 drainUntilEof stream = loop `catch` \(_ :: SomeException) -> pure ()
   where
-    loop = streamReadByte stream >> loop
+    loop = streamReadChunk stream perfBlockSize >> loop
 
--- | Read and discard exactly @n@ bytes. The payload carries no meaning,
--- so no ByteString is built; premature EOF throws.
+-- | Read and discard exactly @n@ bytes at chunk granularity (#276). The
+-- payload carries no meaning, so the chunks are dropped; premature EOF
+-- throws. A chunk request never exceeds the bytes still owed, so no
+-- byte beyond @n@ is consumed from the stream.
 discardExactly :: StreamIO -> Word64 -> IO ()
 discardExactly stream = go
   where
     go :: Word64 -> IO ()
     go 0 = pure ()
-    go !n = streamReadByte stream >> go (n - 1)
+    go !n = do
+      chunk <- streamReadChunk stream (fromIntegral (min n (fromIntegral perfBlockSize)))
+      go (n - fromIntegral (BS.length chunk))
 
 -- | Handle an inbound perf request (responder).
 --

src/LibP2P/Switch/Upgrade.hs

@@ -214,6 +214,7 @@ noiseSessionToStreamIO
 noiseSessionToStreamIO sendRef recvRef bufRef rawIO = StreamIO
   { streamWrite = encryptAndWrite sendRef rawIO
   , streamReadByte = decryptAndReadByte recvRef bufRef rawIO
+  , streamReadChunk = decryptAndReadChunk recvRef bufRef rawIO
   , streamClose = pure ()  -- Encryption layer does not own the connection
   }
 
@@ -233,37 +234,47 @@ encryptAndWrite sendRef rawIO plaintext =
           writeIORef sendRef sess'
           writeFramedMessage rawIO ct
 
--- | Read and decrypt a byte from the Noise channel.
--- If the buffer has bytes, return the first. Otherwise, read Noise
--- frames from the raw stream until one decrypts to a non-empty
--- plaintext, and buffer the result. A transport message with an empty
--- plaintext (a frame carrying only the AEAD tag) is legal — some
--- implementations send it as a keepalive — and a zero-length frame
--- carries no Noise message at all; both yield zero application bytes,
--- so reading continues at the next frame.
+-- | Read Noise frames from the raw stream until one decrypts to a
+-- non-empty plaintext, and return that plaintext. A transport message
+-- with an empty plaintext (a frame carrying only the AEAD tag) is
+-- legal — some implementations send it as a keepalive — and a
+-- zero-length frame carries no Noise message at all; both yield zero
+-- application bytes, so reading continues at the next frame.
+nextPlaintext :: IORef NoiseSession -> StreamIO -> IO ByteString
+nextPlaintext recvRef rawIO = do
+  ct <- readFramedMessage rawIO
+  if BS.null ct
+    then nextPlaintext recvRef rawIO -- zero-length frame: no message to decrypt
+    else do
+      sess <- readIORef recvRef
+      case decryptMessage sess ct of
+        Left err -> fail $ "nextPlaintext: decrypt failed: " <> err
+        Right (pt, sess') -> do
+          writeIORef recvRef sess'
+          if BS.null pt
+            then nextPlaintext recvRef rawIO -- empty transport message (keepalive)
+            else pure pt
+
+-- | Read and decrypt a byte from the Noise channel: pop the buffer if
+-- it has bytes, otherwise decrypt the next frame and buffer the rest.
 decryptAndReadByte :: IORef NoiseSession -> IORef ByteString -> StreamIO -> IO Word8
 decryptAndReadByte recvRef bufRef rawIO = do
   buf <- readIORef bufRef
-  if BS.null buf
-    then fillFromNextFrame
-    else popByte buf
-  where
-    popByte bs = do
-      writeIORef bufRef (BS.tail bs)
-      pure (BS.head bs)
-    fillFromNextFrame = do
-      ct <- readFramedMessage rawIO
-      if BS.null ct
-        then fillFromNextFrame -- zero-length frame: no message to decrypt
-        else do
-          sess <- readIORef recvRef
-          case decryptMessage sess ct of
-            Left err -> fail $ "decryptAndReadByte: " <> err
-            Right (pt, sess') -> do
-              writeIORef recvRef sess'
-              if BS.null pt
-                then fillFromNextFrame -- empty transport message (keepalive)
-                else popByte pt
+  bs <- if BS.null buf then nextPlaintext recvRef rawIO else pure buf
+  writeIORef bufRef (BS.tail bs)
+  pure (BS.head bs)
+
+-- | Chunk-level read from the Noise channel: hand back up to @n@ bytes
+-- of the buffered plaintext (a decrypted frame is already a chunk),
+-- decrypting the next frame only when the buffer is empty. Bytes
+-- beyond @n@ stay buffered for the next read.
+decryptAndReadChunk :: IORef NoiseSession -> IORef ByteString -> StreamIO -> Int -> IO ByteString
+decryptAndReadChunk recvRef bufRef rawIO n = do
+  buf <- readIORef bufRef
+  bs <- if BS.null buf then nextPlaintext recvRef rawIO else pure buf
+  let (front, rest) = BS.splitAt n bs
+  writeIORef bufRef rest
+  pure front
 
 -- | Bounded window given to the send loop to flush the GoAway frame
 -- before the transport is closed underneath it.
@@ -312,18 +323,13 @@ yamuxToMuxerSession yamuxSess closeTransport = do
     }
 
 -- | Convert a YamuxStream to StreamIO with a read buffer.
--- Yamux delivers data in chunks via streamRead, but StreamIO requires
--- byte-by-byte reads. An IORef buffer bridges this gap.
+-- Yamux delivers data in chunks via streamRead; an IORef buffer holds
+-- the bytes a byte- or chunk-level read did not consume.
 yamuxStreamToStreamIO :: YamuxStream -> IO StreamIO
 yamuxStreamToStreamIO yamuxStream = do
   readBuf <- newIORef BS.empty
-  pure StreamIO
-    { streamWrite = \bs -> do
-        result <- YS.streamWrite yamuxStream bs
-        case result of
-          Right () -> pure ()
-          Left err -> fail $ "yamuxStreamWrite: " <> show err
-    , streamReadByte = do
+  let -- Buffered bytes if any, otherwise the next yamux chunk.
+      nextChunk = do
         buf <- readIORef readBuf
         if BS.null buf
           then do
@@ -332,13 +338,23 @@ yamuxStreamToStreamIO yamuxStream = do
               Left err -> fail $ "yamuxStreamRead: " <> show err
               Right chunk
                 | BS.null chunk -> fail "yamuxStreamRead: empty chunk"
-                | BS.length chunk == 1 -> pure (BS.head chunk)
-                | otherwise -> do
-                    writeIORef readBuf (BS.tail chunk)
-                    pure (BS.head chunk)
-          else do
-            writeIORef readBuf (BS.tail buf)
-            pure (BS.head buf)
+                | otherwise -> pure chunk
+          else pure buf
+  pure StreamIO
+    { streamWrite = \bs -> do
+        result <- YS.streamWrite yamuxStream bs
+        case result of
+          Right () -> pure ()
+          Left err -> fail $ "yamuxStreamWrite: " <> show err
+    , streamReadByte = do
+        chunk <- nextChunk
+        writeIORef readBuf (BS.tail chunk)
+        pure (BS.head chunk)
+    , streamReadChunk = \n -> do
+        chunk <- nextChunk
+        let (front, rest) = BS.splitAt n chunk
+        writeIORef readBuf rest
+        pure front
     , streamClose = do
         _ <- YS.streamClose yamuxStream  -- Sends FIN flag
         pure ()

src/LibP2P/Transport/TCP.hs

@@ -134,13 +134,16 @@ tcpListen _ = fail "tcpListen: unsupported multiaddr"
 socketToStreamIO :: NS.Socket -> StreamIO
 socketToStreamIO sock = StreamIO
   { streamWrite = NSB.sendAll sock
-  , streamReadByte = do
-      bs <- NSB.recv sock 1
-      if BS.null bs
-        then fail "socketToStreamIO: connection closed"
-        else pure (BS.head bs)
+  , streamReadByte = BS.head <$> recvChunk 1
+  , streamReadChunk = recvChunk
   , streamClose = NS.close sock
   }
+  where
+    recvChunk n = do
+      bs <- NSB.recv sock n
+      if BS.null bs
+        then fail "socketToStreamIO: connection closed"
+        else pure bs
 
 -- | Convert a SockAddr to a Multiaddr.
 sockAddrToMultiaddr :: NS.SockAddr -> IO Multiaddr

test/LibP2P/ConformanceSpec.hs

@@ -24,6 +24,7 @@ import LibP2P.Crypto.PeerId (PeerId, fromPublicKey)
 import LibP2P.MultistreamSelect.Negotiation
   ( NegotiationResult (..)
   , StreamIO (..)
+  , mkByteStreamIO
   , mkMemoryStreamPair
   , negotiateInitiator
   , negotiateResponder
@@ -55,15 +56,15 @@ mkScriptedStream :: ByteString -> IO (StreamIO, IO ByteString)
 mkScriptedStream canned = do
   writtenRef <- newIORef BS.empty
   readRef <- newIORef canned
-  let stream = StreamIO
-        { streamWrite = \bs -> modifyIORef' writtenRef (`BS.append` bs)
-        , streamReadByte = do
-            buf <- readIORef readRef
-            case BS.uncons buf of
-              Nothing -> ioError (userError "scripted stream: EOF")
-              Just (b, rest) -> writeIORef readRef rest >> pure b
-        , streamClose = pure ()
-        }
+  let readB = do
+        buf <- readIORef readRef
+        case BS.uncons buf of
+          Nothing -> ioError (userError "scripted stream: EOF")
+          Just (b, rest) -> writeIORef readRef rest >> pure b
+      stream = mkByteStreamIO
+        (\bs -> modifyIORef' writtenRef (`BS.append` bs))
+        readB
+        (pure ())
   pure (stream, readIORef writtenRef)
 
 -- multistream-select messages, hand-derived from the spec:

test/LibP2P/DHT/DHTSpec.hs

@@ -27,7 +27,7 @@
 import LibP2P.DHT.RoutingTable (allPeers, bucketForPeer, insertPeer, newRoutingTable)
 import LibP2P.DHT.Types
 import LibP2P.Multiaddr (Multiaddr, fromText, toBytes)
-import LibP2P.MultistreamSelect.Negotiation (StreamIO (..), negotiateResponder)
+import LibP2P.MultistreamSelect.Negotiation (StreamIO (..), mkByteStreamIO, negotiateResponder)
 import LibP2P.Switch.ConnPool (addConn)
 import LibP2P.Switch.Types
   ( ConnState (..)
@@ -151,16 +151,14 @@
           Nothing -> do
             closed <- readTVar closedVar
             if closed then throwSTM (userError "stream closed") else retry
-      streamA = StreamIO
-        { streamWrite = writeAll q1
-        , streamReadByte = readOrEOF q2 closedBtoA
-        , streamClose = atomically (writeTVar closedAtoB True)
-        }
-      streamB = StreamIO
-        { streamWrite = writeAll q2
-        , streamReadByte = readOrEOF q1 closedAtoB
-        , streamClose = atomically (writeTVar closedBtoA True)
-        }
+      streamA = mkByteStreamIO
+        (writeAll q1)
+        (readOrEOF q2 closedBtoA)
+        (atomically (writeTVar closedAtoB True))
+      streamB = mkByteStreamIO
+        (writeAll q2)
+        (readOrEOF q1 closedAtoB)
+        (atomically (writeTVar closedBtoA True))
   pure (streamA, streamB)
 
 -- | A mock Connection that hands out the given stream on the first
@@ -427,7 +425,7 @@
       -- Provider should now be persisted
       stored <- getProviders node key
       length stored `shouldBe` 1
       peProvider (head stored) `shouldBe` remotePid
 
     it "ADD_PROVIDER round-trip via GET_PROVIDERS" $ do
       node <- mkTestNode localPid
@@ -461,7 +459,7 @@
         Right resp -> do
           msgType resp `shouldBe` GetProviders
           length (msgProviderPeers resp) `shouldBe` 1
           dhtPeerId (head (msgProviderPeers resp)) `shouldBe` peerIdBytes remotePid
         Left err -> expectationFailure $ "Failed: " ++ err
 
     -- Issue #147: specs/kad-dht requires handling additional RPC request
@@ -755,7 +753,7 @@
           bucketPeers = sameBucketPeers (kValue + 1)
           initial = take kValue bucketPeers
           newcomer = last bucketPeers
           lrs = head initial
       mapM_ (\pid -> do
                e <- mkEntry pid
                atomically $ modifyTVar' (dhtRoutingTable node) (fst . insertPeer e))
@@ -776,7 +774,7 @@
           bucketPeers = sameBucketPeers (kValue + 1)
           initial = take kValue bucketPeers
           newcomer = last bucketPeers
           lrs = head initial
       mapM_ (\pid -> do
                e <- mkEntry pid
                atomically $ modifyTVar' (dhtRoutingTable node) (fst . insertPeer e))