Release a stream's read buffer after a large message is handed off

[RichieAHB]

The problem

On client-streaming and bidi requests the server's body reader keeps one BytesMut for the whole stream. Each uncompressed message is handed to the handler as buf.split_to(len).freeze(), a zero-copy slice of that buffer's allocation. (A compressed message is decompressed into a buffer of its own; the reader's buffer is retained in the same way.) The reader keeps the rest of the allocation, the unused tail, for the next message. Two properties of BytesMut then combine:

• it never shrinks;
• when the tail runs out and no slice is outstanding, reserve reclaims the whole allocation and reuses it.

So the allocation stays alive, at the size of the largest message seen so far, until the stream ends. This holds even when the handler dropped that message long ago. Growth is by doubling, so the allocation can be up to twice the message.

Example, within the default 4 MiB max_message_size: a 3 MiB first message on a long-lived stream, followed only by small messages, keeps an allocation of about 4 MiB until the stream ends. With the limit raised, a 16 MiB first message keeps about 32 MiB. A server with many long-lived streams that each begin with one large message holds that memory for every open stream.

ServerStream on the client does the same with the response body.

Nothing a handler or caller can do releases the buffer, short of ending the stream. Limits bounds how large a message may be, not how long its buffer lives.

The fix

After decoding, if the read buffer is empty and was larger than 64 KiB, replace it with BytesMut::new(). The messages already handed out become the only owners of the allocation, so it is freed when the handler (or caller) drops them. The next message starts a fresh buffer.

BodyReader::enter_drain_mode already replaces the buffer when the request stream ends, "instead of keeping them resident for the duration of the drain". This change applies the same idea between messages.

The helper release_drained_buf lives in envelope.rs. It is called from BodyReader::on_data (server) and from ServerStream::next_message_or_end (client). No public API changes.

The 64 KiB guard

Replacing the buffer every time it is empty would make a stream of small messages pay two small allocations per message (a new Vec, plus the shared header that split_to allocates), where today it can reuse one buffer. With the guard, a stream of small messages read in ordinary-sized frames behaves exactly as before. (If such a stream is read in frames larger than 64 KiB, it gets one new allocation per such frame.)

The guard compares max(bytes buffered before decoding, remaining capacity) with the threshold. capacity() on its own is not enough: after split_to it reports only the unused tail, which can be zero while the buffer still pins a large allocation. The bytes buffered before decoding are a lower bound on the allocation's size. The capacity term catches a buffer that reserve reclaimed whole, and a large tail left after a message when the next one is already partly buffered.

64 KiB is well above h2's default 16 KiB max_frame_size. It is also the largest "original capacity" that bytes itself restores when a shared BytesMut has to reallocate (MAX_ORIGINAL_CAPACITY_WIDTH in bytes_mut.rs).

Measured

A server with about 130 open streams, each carrying one 16 MiB message first and small ones after it, heap-profiled for 35 seconds, one run before and one after (the max_message_size limit was raised for the test):

• read buffer held per open stream: 32 MiB before, 0 after;
• heap allocated per open stream: 52.6 MiB before, 20.6 MiB after;
• process RSS per open stream: about 33 MiB before, about 17 MiB after. The buffer reserved twice the message and wrote half of it, so resident memory falls by about one message per stream.

After the streams ended both builds returned to the same level.

Cost

A stream whose messages are all large now starts a fresh buffer for each message and grows it by doubling. Before, it could reuse one allocation when the handler had already dropped the previous message. I have not measured this with the repository's benches, and I am happy to.

Left out on purpose

• Reserving the exact message size once the 5-byte envelope header has arrived. It would avoid the doubling, and the repeated grow-and-copy while a large message arrives. But it would let a peer make the receiver allocate up to max_message_size after sending five bytes; today memory follows the bytes actually received. That deserves its own discussion.
• Avoiding the copy in extend_from_slice when one body frame holds a whole message.

Limits of the fix

The release needs the buffer to be empty after a body frame, and at that moment either the bytes just decoded or the unused tail must be above 64 KiB. If the last frame of a large message also carries the start of the next one, and the allocation's unused tail is 64 KiB or less, neither holds. The allocation then stays until that tail is used up (at most 64 KiB of further data); a stream that goes idle first keeps it, as today. It is never worse than today.

An exact variant is possible: keep a per-stream high-water mark of buf.len(), reset on release, as one more field in BodyReader and in ServerStream. I left it out to keep the change small; say if you prefer it.

Also in this change

• A changelog entry under .changes/unreleased/.
• A note in the client-streaming guide: an item is a slice of the stream's read buffer, so a handler that keeps many items past the loop should keep to_owned_message() results.

Tests

Server (service.rs, next to the other test_body_reader_* tests):

• test_body_reader_releases_buffer_after_large_message: a 1 MiB message arrives in 16 KiB frames. Afterwards the reader's buffer has capacity 0 and the received Bytes is_unique().
• test_body_reader_releases_exactly_sized_buffer: a 128 KiB message arrives in one frame, so the buffer is sized exactly and has no tail. Only the "bytes buffered" term of the guard can release it. (With a capacity-only guard this test fails.)
• test_body_reader_releases_buffer_after_trailing_partial_message: the frame that completes a 1 MiB message also carries the start of a small one. Both arrive in order, and the buffer is released when the small one has drained it. Only the capacity term can do that. (With a buffered-only guard this test and the next one fail.)
• test_body_reader_releases_reclaimed_buffer: the limit described above. A large message fills its allocation to within 30 bytes and the next message is partly buffered, so the buffer is kept when it drains. Once the handler has dropped the messages and a later message makes the reader reclaim the allocation, the next drain releases it.
• test_body_reader_keeps_small_buffer: after a small message the buffer is still shared with the message, so it is kept for reuse.

Client (client/mod.rs): server_stream_releases_buffer_after_large_message and server_stream_releases_exactly_sized_buffer, the same two cases for ServerStream.

With the release turned into a no-op, all six "releases" tests fail and the "keeps small buffer" test passes. With the change, cargo test --all-features on the crate passes (688 unit tests and 9 doc tests on the 0.9.0 sources), and cargo clippy --all-features --all-targets and cargo fmt --check (default settings) report nothing. The conformance suite is not part of the published crate, so I did not run it locally; CI on this pull request runs it against main.

The tests rely on Bytes::is_unique. It exists from bytes 1.6.0, the crate's current minimum, but bytes 1.6.1 fixed it returning wrong values for a Bytes that came from a shared BytesMut, which is what these tests check. So in practice they need 1.6.1; you may or may not want to raise the floor.

[iainmcgin]

[claude code] Thank you for tracking this down, and for the measurements in the description, which #302 relies on.

#302 has now merged, and it supersedes this PR. It replaces the per-stream read buffer in both the server request reader and the client ServerStream with a per-message assembler. A message larger than 4 KiB gets an allocation of its own, which the reader does not keep after it hands the message off, and smaller messages share an 8 KiB slab per stream. A stream that once carried a large message therefore no longer holds that allocation until the stream ends.

Bounding the total memory held by received messages across all streams is tracked separately in #303.