tokio-quiche: lazily allocate H3 body_recv_buf

H3Driver eagerly allocated a 64 KiB body receive buffer per connection at construction. Make it Option<Limit<BytesMut>>, allocate it lazily on the first body read, and release it once no streams or flows remain, so idle connections hold no receive buffer.

Author: LynnLi100

tokio-quiche/src/http3/driver/mod.rs

@@ -350,9 +350,11 @@ pub struct H3Driver<H: DriverHooks> {
     dgram_recv: OutboundFrameStream,
     /// Keeps the datagram channel open such that datagram flows can be created.
     dgram_send: OutboundFrameSender,
-    /// A buffer to receive H3 body data from quiche. We initialize a large
-    /// buffer and then `split()` off filled parts until we need to reallocate.
-    body_recv_buf: bytes::buf::Limit<BytesMut>,
+    /// A buffer to receive H3 body data from quiche. Lazily allocated on the
+    /// first body read and released once no streams remain, so idle
+    /// connections hold no receive buffer. We `split()` off filled parts until
+    /// we need to reallocate.
+    body_recv_buf: Option<bytes::buf::Limit<BytesMut>>,
 
     /// The buffer used to interact with the underlying IoWorker.
     io_worker_buf: Vec<u8>,
@@ -393,8 +395,7 @@ impl<H: DriverHooks> H3Driver<H> {
                 dgram_recv,
                 dgram_send: PollSender::new(dgram_send),
                 max_stream_seen: 0,
-                body_recv_buf: BytesMut::with_capacity(BufFactory::MAX_BUF_SIZE)
-                    .limit(BufFactory::MAX_BUF_SIZE),
+                body_recv_buf: None,
                 io_worker_buf: vec![0u8; BufFactory::MAX_BUF_SIZE],
 
                 waiting_streams: FuturesUnordered::new(),
@@ -521,17 +522,23 @@ impl<H: DriverHooks> H3Driver<H> {
                 };
             }
 
-            // NOTE: `self.body_recv_buf` is `Limit<BytesMut>` so
+            // Lazily allocate the receive buffer on first use; idle
+            // connections never receive body bytes and never allocate it.
+            let body_recv_buf = self.body_recv_buf.get_or_insert_with(|| {
+                BytesMut::with_capacity(BufFactory::MAX_BUF_SIZE)
+                    .limit(BufFactory::MAX_BUF_SIZE)
+            });
+            // NOTE: `body_recv_buf` is `Limit<BytesMut>` so
             // `has_remaining_mut()` will indicate if the buffer
             // has space available until the *limit* is
             //  reached. (A plain `BytesMut` can reallocate and would always
             // return true)
-            if !self.body_recv_buf.has_remaining_mut() {
-                self.body_recv_buf =
+            if !body_recv_buf.has_remaining_mut() {
+                *body_recv_buf =
                     BytesMut::with_capacity(BufFactory::MAX_BUF_SIZE)
-                        .limit(BufFactory::MAX_BUF_SIZE)
-            };
-            match conn.recv_body_buf(qconn, stream_id, &mut self.body_recv_buf) {
+                        .limit(BufFactory::MAX_BUF_SIZE);
+            }
+            match conn.recv_body_buf(qconn, stream_id, &mut *body_recv_buf) {
                 Ok(n) => {
                     ctx.audit_stats.add_downstream_bytes_recvd(n as u64);
                     let event = H3Event::BodyBytesReceived {
@@ -541,12 +548,12 @@ impl<H: DriverHooks> H3Driver<H> {
                     };
                     let _ = self.h3_event_sender.send(event.into());
                     // Take the filled part, leave the remaining capacity
-                    let filled_body = self.body_recv_buf.get_mut().split();
+                    let filled_body = body_recv_buf.get_mut().split();
                     // Sanity check: the remaining spare capacity should equal
                     // the limit.
                     debug_assert_eq!(
-                        self.body_recv_buf.get_mut().spare_capacity_mut().len(),
-                        self.body_recv_buf.remaining_mut()
+                        body_recv_buf.get_mut().spare_capacity_mut().len(),
+                        body_recv_buf.remaining_mut()
                     );
                     permit.send(InboundFrame::Body(filled_body, false));
                 },
@@ -1003,15 +1010,23 @@ impl<H: DriverHooks> H3Driver<H> {
         Ok(())
     }
 
-    /// Closes the connection with `NoError` if the H3 event receiver
-    /// has been dropped and there are no active streams or flows.
-    fn close_if_idle(&self, qconn: &mut QuicheConnection) {
-        if self.h3_event_receiver_dropped &&
-            self.stream_map.is_empty() &&
-            self.flow_map.is_empty()
-        {
-            let _ =
-                qconn.close(true, quiche::h3::WireErrorCode::NoError as u64, &[]);
+    /// Handles connection cleanup once no streams or flows remain.
+    ///
+    /// Releases the body receive buffer (it is reallocated lazily on the next
+    /// body read; body bytes only flow on active streams, so an empty stream
+    /// map means it is unused) and closes the connection with `NoError` if the
+    /// H3 event receiver has been dropped.
+    fn close_if_idle(&mut self, qconn: &mut QuicheConnection) {
+        if self.stream_map.is_empty() && self.flow_map.is_empty() {
+            self.body_recv_buf = None;
+
+            if self.h3_event_receiver_dropped {
+                let _ = qconn.close(
+                    true,
+                    quiche::h3::WireErrorCode::NoError as u64,
+                    &[],
+                );
+            }
         }
     }
 

tokio-quiche/src/http3/driver/test_utils.rs

@@ -277,7 +277,8 @@ impl<H: DriverHooks + GetConnectionForHook> DriverTestHelper<H> {
     }
 
     pub fn driver_set_body_buf_size(&mut self, limit: usize) {
-        self.driver.body_recv_buf = BytesMut::with_capacity(limit).limit(limit);
+        self.driver.body_recv_buf =
+            Some(BytesMut::with_capacity(limit).limit(limit));
     }
 }
 

tokio-quiche/src/http3/driver/tests.rs

@@ -243,6 +243,149 @@ mod client_side_driver {
         assert_eq!(helper.driver.stream_map.len(), 0);
     }
 
+    /// An idle connection (and a request/response that exchanges only
+    /// headers) must never allocate the body receive buffer.
+    #[test]
+    fn client_body_recv_buf_not_allocated_when_idle() {
+        let mut helper = DriverTestHelper::<ClientHooks>::new().unwrap();
+        helper.complete_handshake().unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        // Idle: never received body bytes, so the buffer is not allocated.
+        assert!(helper.driver.body_recv_buf.is_none());
+
+        // client sends a request with fin (no request body)
+        let stream_id = helper
+            .driver_send_request(make_request_headers("GET"), true)
+            .unwrap();
+
+        // server reads request and sends response headers (no body)
+        helper.advance_and_run_loop().unwrap();
+        assert_matches!(
+            helper.peer_server_poll().unwrap(),
+            (0, h3::Event::Headers { .. })
+        );
+        helper.peer_server_send_response(0, false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        // Client receives response headers
+        let resp = assert_matches!(
+            helper.driver_recv_core_event().unwrap(),
+            H3Event::IncomingHeaders(headers) => { headers }
+        );
+        assert_eq!(resp.stream_id, stream_id);
+
+        // Only headers have been exchanged; still no body buffer.
+        assert!(helper.driver.body_recv_buf.is_none());
+    }
+
+    /// The body receive buffer is lazily allocated on the first body read
+    /// and released once the last stream is cleaned up.
+    #[test]
+    fn client_body_recv_buf_allocated_on_body_and_released_on_close() {
+        let mut helper = DriverTestHelper::<ClientHooks>::new().unwrap();
+        helper.complete_handshake().unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        let stream_id = helper
+            .driver_send_request(make_request_headers("GET"), true)
+            .unwrap();
+
+        helper.advance_and_run_loop().unwrap();
+        assert_matches!(
+            helper.peer_server_poll().unwrap(),
+            (0, h3::Event::Headers { .. })
+        );
+        helper.peer_server_send_response(0, false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        let resp = assert_matches!(
+            helper.driver_recv_core_event().unwrap(),
+            H3Event::IncomingHeaders(headers) => { headers }
+        );
+        assert_eq!(resp.stream_id, stream_id);
+        let mut from_server = resp.recv;
+
+        // No body yet: the buffer is still unallocated.
+        assert!(helper.driver.body_recv_buf.is_none());
+
+        // Server sends a body chunk (not fin).
+        helper.peer_server_send_body(0, &[7; 10], false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        assert_eq!(helper.driver_try_recv_body(&mut from_server).0, vec![7; 10]);
+
+        // The first body read lazily allocated the buffer.
+        assert!(helper.driver.body_recv_buf.is_some());
+
+        // Server finishes the stream.
+        helper.peer_server_send_body(0, &[8; 10], true).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        let (body, fin, _) = helper.driver_try_recv_body(&mut from_server);
+        assert_eq!(body, vec![8; 10]);
+        assert!(fin);
+
+        // Stream cleaned up on both-directions-close, buffer released.
+        assert_eq!(helper.driver.stream_map.len(), 0);
+        assert!(helper.driver.body_recv_buf.is_none());
+    }
+
+    /// A body larger than the receive buffer exercises the reallocation
+    /// branch; the buffer stays allocated across reallocations and is
+    /// released once the stream closes.
+    #[test]
+    fn client_body_recv_buf_reallocates_and_releases() {
+        let mut helper = DriverTestHelper::<ClientHooks>::new().unwrap();
+        helper.complete_handshake().unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        let stream_id = helper
+            .driver_send_request(make_request_headers("GET"), true)
+            .unwrap();
+
+        helper.advance_and_run_loop().unwrap();
+        assert_matches!(
+            helper.peer_server_poll().unwrap(),
+            (0, h3::Event::Headers { .. })
+        );
+        helper.peer_server_send_response(0, false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+
+        let resp = assert_matches!(
+            helper.driver_recv_core_event().unwrap(),
+            H3Event::IncomingHeaders(headers) => { headers }
+        );
+        assert_eq!(resp.stream_id, stream_id);
+        let mut from_server = resp.recv;
+
+        // Force a small receive buffer so the body exhausts it and the
+        // reallocation branch (`*body_recv_buf = ...`) runs.
+        helper.driver_set_body_buf_size(20);
+
+        // Send 40 bytes across the 20-byte buffer, exhausting it repeatedly.
+        helper.peer_server_send_body(0, &[1; 10], false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        assert_eq!(helper.driver_try_recv_body(&mut from_server).0, vec![1; 10]);
+        helper.peer_server_send_body(0, &[2; 10], false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        assert_eq!(helper.driver_try_recv_body(&mut from_server).0, vec![2; 10]);
+        helper.peer_server_send_body(0, &[3; 10], false).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        assert_eq!(helper.driver_try_recv_body(&mut from_server).0, vec![3; 10]);
+        // Buffer remains allocated across reallocation.
+        assert!(helper.driver.body_recv_buf.is_some());
+
+        // Final chunk with fin.
+        helper.peer_server_send_body(0, &[4; 10], true).unwrap();
+        helper.advance_and_run_loop().unwrap();
+        let (body, fin, _) = helper.driver_try_recv_body(&mut from_server);
+        assert_eq!(body, vec![4; 10]);
+        assert!(fin);
+
+        // Stream cleaned up, buffer released.
+        assert_eq!(helper.driver.stream_map.len(), 0);
+        assert!(helper.driver.body_recv_buf.is_none());
+    }
+
     /// Test that dropping the OutboundFrame channel causes the driver to
     /// send a RESET_STREAM frame to the peer.
     #[test]