From 339b4a555acb1476454d13c9eef050a61a549fe2 Mon Sep 17 00:00:00 2001 From: Kevin O'Shea Date: Fri, 9 Oct 2026 21:14:33 -0500 Subject: [PATCH 1/2] perf: write trailers in the same write as chained DATA A DATA payload of at least CHAIN_THRESHOLD bytes is chained rather than copied, and while it is pending the encoder reports no capacity, so a HEADERS frame already queued behind it (typically trailers) is only buffered after a flush and costs a second write. Every gRPC unary response has this shape. While a chained DATA frame is pending, let Prioritize encode one HEADERS frame from the head of the send queue into a small tail buffer that the encoder writes after the payload in the same vectored write. Frame order and stream scheduling are unchanged; CONTINUATION frames of a large HEADERS block follow once the tail is written. Codec unit tests cover the single write, byte-identical output, partial writes, and the cases where no tail is used; an integration test sends a large DATA frame followed by trailers. --- src/codec/framed_write.rs | 220 ++++++++++++++++++++++++++++++- src/codec/mod.rs | 11 ++ src/proto/streams/prioritize.rs | 57 ++++++++ tests/h2-tests/tests/trailers.rs | 41 ++++++ 4 files changed, 327 insertions(+), 2 deletions(-) diff --git a/src/codec/framed_write.rs b/src/codec/framed_write.rs index efd1966cf..6d684625a 100644 --- a/src/codec/framed_write.rs +++ b/src/codec/framed_write.rs @@ -41,6 +41,17 @@ struct Encoder { /// Next frame to encode next: Option>, + /// Frames encoded while a chained DATA frame is pending in `next`. + /// + /// Written in the same vectored write, after the DATA payload, so a + /// HEADERS frame that follows a chained DATA frame (such as trailers) + /// does not cost a separate write. Holds at most one frame. + tail: Cursor, + + /// The CONTINUATION of a HEADERS frame in `tail` that did not fit in one + /// frame. Encoded once the tail has been written. + tail_continuation: Option, + /// Last data frame last_data_frame: Option>, @@ -95,6 +106,8 @@ where hpack: hpack::Encoder::default(), buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)), next: None, + tail: Cursor::new(BytesMut::new()), + tail_continuation: None, last_data_frame: None, max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE, chain_threshold, @@ -134,6 +147,18 @@ where self.encoder.buffer(item) } + /// Returns whether a HEADERS frame can be queued behind the chained DATA + /// frame that currently prevents `has_capacity`. + pub(crate) fn can_buffer_headers_after_data(&self) -> bool { + self.encoder.can_buffer_headers_after_data() + } + + /// Queue a HEADERS frame to be written after the pending chained DATA + /// frame, in the same write. `can_buffer_headers_after_data` must be true. + pub(crate) fn buffer_headers_after_data(&mut self, item: frame::Headers) { + self.encoder.buffer_headers_after_data(item) + } + /// Flush buffered data to the wire pub fn flush(&mut self, cx: &mut Context) -> Poll> { let span = tracing::trace_span!("FramedWrite::flush"); @@ -144,7 +169,9 @@ where let n = match self.encoder.next { Some(Next::Data(ref mut frame)) => { tracing::trace!(queued_data_frame = true); - let mut buf = (&mut self.encoder.buf).chain(frame.payload_mut()); + let mut buf = (&mut self.encoder.buf) + .chain(frame.payload_mut()) + .chain(&mut self.encoder.tail); ready!(poll_write_buf(Pin::new(&mut self.inner), cx, &mut buf))? } _ => { @@ -205,6 +232,16 @@ where match self.next.take() { Some(Next::Data(frame)) => { self.last_data_frame = Some(frame); + // The tail was written with the data frame. + self.tail.set_position(0); + self.tail.get_mut().clear(); + if let Some(continuation) = self.tail_continuation.take() { + let mut buf = limited_write_buf!(self); + if let Some(continuation) = continuation.encode(&mut buf) { + self.next = Some(Next::Continuation(continuation)); + } + return ControlFlow::Continue; + } debug_assert!(self.is_empty()); ControlFlow::Break } @@ -312,9 +349,30 @@ where >= self.min_buffer_capacity) } + fn can_buffer_headers_after_data(&self) -> bool { + matches!(self.next, Some(Next::Data(_))) + && self.tail.get_ref().is_empty() + && self.tail_continuation.is_none() + } + + fn buffer_headers_after_data(&mut self, item: frame::Headers) { + assert!(self.can_buffer_headers_after_data()); + let span = tracing::trace_span!("FramedWrite::buffer_headers_after_data", frame = ?item); + let _e = span.enter(); + tracing::debug!(frame = ?item, "send"); + + // Same per-frame bound as `buf`: one frame plus its header. A larger + // header block continues in a CONTINUATION frame written afterwards. + let limit = self.max_frame_size() + frame::HEADER_LEN; + let mut tail = self.tail.get_mut().limit(limit); + self.tail_continuation = item.encode(&mut self.hpack, &mut tail); + } + fn is_empty(&self) -> bool { match self.next { - Some(Next::Data(ref frame)) => !frame.payload().has_remaining(), + Some(Next::Data(ref frame)) => { + !frame.payload().has_remaining() && !self.tail.has_remaining() + } _ => !self.buf.has_remaining(), } } @@ -376,3 +434,161 @@ mod unstable { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::frame::{Data, Headers, StreamId}; + use bytes::Bytes; + use http::HeaderMap; + use std::sync::Arc; + use std::task::{Wake, Waker}; + + struct NoopWake; + + impl Wake for NoopWake { + fn wake(self: Arc) {} + } + + fn noop_waker() -> Waker { + Waker::from(Arc::new(NoopWake)) + } + + /// Records each write call as one entry; reports vectored support. + #[derive(Default)] + struct RecordingIo { + writes: Vec>, + max_per_write: Option, + } + + impl AsyncWrite for RecordingIo { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let n = self + .max_per_write + .map_or(buf.len(), |max| buf.len().min(max)); + self.writes.push(buf[..n].to_vec()); + Poll::Ready(Ok(n)) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + let mut write = Vec::new(); + for buf in bufs { + write.extend_from_slice(buf); + } + if let Some(max) = self.max_per_write { + write.truncate(max); + } + let n = write.len(); + self.writes.push(write); + Poll::Ready(Ok(n)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + fn trailers(stream: u32) -> Headers { + let mut fields = HeaderMap::new(); + fields.insert("grpc-status", "0".parse().unwrap()); + Headers::trailers(StreamId::from(stream), fields) + } + + fn flush(write: &mut FramedWrite) { + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(write.flush(&mut cx).is_ready()); + } + + /// The bytes one frame encodes to when buffered on its own. + fn alone(frames: Vec>) -> Vec { + let mut write = FramedWrite::new(RecordingIo::default()); + for frame in frames { + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(write.poll_ready(&mut cx).is_ready()); + write.buffer(frame).unwrap(); + } + flush(&mut write); + write.inner.writes.concat() + } + + #[test] + fn headers_after_a_chained_data_frame_share_its_write() { + let payload = Bytes::from(vec![7_u8; 1_000]); + let mut write = FramedWrite::new(RecordingIo::default()); + write + .buffer(Data::new(StreamId::from(1), payload.clone()).into()) + .unwrap(); + // The chained payload blocks ordinary buffering ... + assert!(!write.has_capacity()); + // ... but a HEADERS frame may follow it in the same write. + assert!(write.can_buffer_headers_after_data()); + write.buffer_headers_after_data(trailers(1)); + assert!(!write.can_buffer_headers_after_data()); + flush(&mut write); + + assert_eq!( + write.inner.writes.len(), + 1, + "one write for data and trailers" + ); + // Byte-identical to writing the two frames in order separately. + let expected = alone(vec![ + Data::new(StreamId::from(1), payload).into(), + trailers(1).into(), + ]); + assert_eq!(write.inner.writes[0], expected); + assert!(write.has_capacity()); + assert!(write.take_last_data_frame().is_some()); + } + + #[test] + fn partial_writes_drain_data_before_the_queued_headers() { + let payload = Bytes::from((0..4_000_u32).map(|i| i as u8).collect::>()); + let mut write = FramedWrite::new(RecordingIo { + max_per_write: Some(333), + ..RecordingIo::default() + }); + write + .buffer(Data::new(StreamId::from(3), payload.clone()).into()) + .unwrap(); + write.buffer_headers_after_data(trailers(3)); + flush(&mut write); + + let expected = alone(vec![ + Data::new(StreamId::from(3), payload).into(), + trailers(3).into(), + ]); + assert_eq!(write.inner.writes.concat(), expected); + assert!(write.has_capacity()); + } + + #[test] + fn headers_cannot_follow_unchained_or_absent_data() { + let mut write = FramedWrite::<_, Bytes>::new(RecordingIo::default()); + assert!(!write.can_buffer_headers_after_data()); + // A small payload is copied, not chained, so ordinary buffering applies. + write + .buffer(Data::new(StreamId::from(1), Bytes::from_static(b"small")).into()) + .unwrap(); + assert!(write.has_capacity()); + assert!(!write.can_buffer_headers_after_data()); + } +} diff --git a/src/codec/mod.rs b/src/codec/mod.rs index d7aecee91..190d437d0 100644 --- a/src/codec/mod.rs +++ b/src/codec/mod.rs @@ -157,6 +157,17 @@ where self.framed_write().flush(cx) } + /// Returns whether a HEADERS frame can be queued behind a pending + /// chained DATA frame, to be written in the same write. + pub(crate) fn can_buffer_headers_after_data(&mut self) -> bool { + self.framed_write().can_buffer_headers_after_data() + } + + /// Queue a HEADERS frame behind the pending chained DATA frame. + pub(crate) fn buffer_headers_after_data(&mut self, item: crate::frame::Headers) { + self.framed_write().buffer_headers_after_data(item) + } + /// Shutdown the send half pub fn shutdown(&mut self, cx: &mut Context) -> Poll> { self.framed_write().shutdown(cx) diff --git a/src/proto/streams/prioritize.rs b/src/proto/streams/prioritize.rs index 910b2573d..f575f9149 100644 --- a/src/proto/streams/prioritize.rs +++ b/src/proto/streams/prioritize.rs @@ -526,6 +526,9 @@ impl Prioritize { loop { if !dst.has_send_capacity() { + if self.buffer_headers_after_data(buffer, store, counts, dst, max_frame_len) { + continue; + } return Ok(BufferStatus::CodecFull); } @@ -557,6 +560,60 @@ impl Prioritize { } } + /// While the codec holds a chained DATA frame, moves a HEADERS frame at + /// the head of the send queue into the same write, behind it. + /// + /// A unary response is HEADERS, DATA, then trailers; without this, the + /// trailers wait for the DATA payload to be flushed and cost a second + /// write. Only a HEADERS frame already at the head of the queue is taken, + /// so frame order and stream scheduling are unchanged, and no DATA frame + /// is ever added while one is in flight. + fn buffer_headers_after_data( + &mut self, + buffer: &mut Buffer>, + store: &mut Store, + counts: &mut Counts, + dst: &mut Codec>, + max_frame_len: usize, + ) -> bool + where + T: AsyncWrite + Unpin, + B: Buf, + { + if !matches!(self.in_flight_data_frame, InFlightData::DataFrame(_)) + || !dst.can_buffer_headers_after_data() + { + return false; + } + let headers_next = match self.pending_send.pop(store) { + Some(mut stream) => { + let headers_next = match stream.pending_send.pop_front(buffer) { + Some(frame) => { + let is_headers = matches!(frame, Frame::Headers(_)); + stream.pending_send.push_front(buffer, frame); + is_headers + } + None => false, + }; + self.pending_send.push_front(&mut stream); + headers_next + } + None => false, + }; + if !headers_next { + return false; + } + match self.pop_frame(buffer, store, max_frame_len, counts) { + Some(Frame::Headers(frame)) => { + tracing::trace!(?frame, "writing after in-flight data"); + dst.buffer_headers_after_data(frame); + true + } + Some(_) => unreachable!("the queue head was checked to be a HEADERS frame"), + None => false, + } + } + pub fn reclaim_written_frame( &mut self, buffer: &mut Buffer>, diff --git a/tests/h2-tests/tests/trailers.rs b/tests/h2-tests/tests/trailers.rs index b1c76707a..2aacd0c68 100644 --- a/tests/h2-tests/tests/trailers.rs +++ b/tests/h2-tests/tests/trailers.rs @@ -262,3 +262,44 @@ async fn send_trailers_rejects_connection_specific_headers() { join(srv, client).await; } + +#[tokio::test] +async fn server_sends_large_data_then_trailers_in_order() { + // A DATA payload large enough to be chained rather than copied, followed + // by trailers queued in the same turn, as a gRPC unary response does. + h2_support::trace_init!(); + let (io, mut client) = mock::new(); + let body = vec![b'x'; 4_000]; + let expected = body.clone(); + + let client = async move { + let settings = client.assert_server_handshake().await; + assert_default_settings!(settings); + client + .send_frame( + frames::headers(1) + .request("POST", "https://example.com/") + .eos(), + ) + .await; + client.recv_frame(frames::headers(1).response(200)).await; + client.recv_frame(frames::data(1, expected)).await; + client + .recv_frame(frames::headers(1).field("grpc-status", "0").eos()) + .await; + }; + + let srv = async move { + let mut srv = server::handshake(io).await.expect("handshake"); + let (_req, mut stream) = srv.next().await.unwrap().unwrap(); + let rsp = http::Response::builder().status(200).body(()).unwrap(); + let mut send = stream.send_response(rsp, false).unwrap(); + send.send_data(body.into(), false).unwrap(); + let mut trailers = HeaderMap::new(); + trailers.insert("grpc-status", "0".parse().unwrap()); + send.send_trailers(trailers).unwrap(); + assert!(srv.next().await.is_none()); + }; + + join(client, srv).await; +} From c8bfb22fedd515ca1de789474816b69262126642 Mon Sep 17 00:00:00 2001 From: Kevin O'Shea Date: Sat, 10 Oct 2026 21:13:33 -0500 Subject: [PATCH 2/2] fix: never write trailers before a DATA frame's unsent remainder When a DATA payload is larger than the frame size or the flow-control window, pop_frame sends only a prefix of the buffer; the remainder is pushed back onto the stream's queue only when the frame is reclaimed after it is written. The trailers fast path looked at the queue before that, saw the trailers at its head, and wrote them right after the first frame, ending the stream with the rest of the body unsent. Only allow HEADERS to follow an in-flight DATA frame that carries all of its buffer. The last frame of a multi-frame body still shares its write with the trailers. Adds server_sends_multi_frame_data_then_trailers_without_truncation, which fails on the previous commit (trailers arrive after the first 16,384-byte frame) and passes on v0.4.20 and with this change. --- src/proto/streams/prioritize.rs | 13 +++++++++ tests/h2-tests/tests/trailers.rs | 45 ++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/src/proto/streams/prioritize.rs b/src/proto/streams/prioritize.rs index f575f9149..d91895f87 100644 --- a/src/proto/streams/prioritize.rs +++ b/src/proto/streams/prioritize.rs @@ -54,6 +54,11 @@ pub(super) struct Prioritize { /// What `DATA` frame is currently being sent in the codec. in_flight_data_frame: InFlightData, + /// Whether the in-flight `DATA` frame carries all of its buffer. When it + /// does not, the unsent suffix is only returned to the stream's queue when + /// the frame is reclaimed, so nothing queued behind it may be written yet. + in_flight_data_is_whole: bool, + /// The maximum amount of bytes a stream should buffer. max_buffer_size: usize, } @@ -100,6 +105,7 @@ impl Prioritize { flow, last_opened_id: StreamId::ZERO, in_flight_data_frame: InFlightData::Nothing, + in_flight_data_is_whole: false, max_buffer_size: config.local_max_buffer_size, } } @@ -544,6 +550,9 @@ impl Prioritize { debug_assert_eq!(self.in_flight_data_frame, InFlightData::Nothing); if let Frame::Data(ref frame) = frame { self.in_flight_data_frame = InFlightData::DataFrame(frame.payload().stream); + let payload = &frame.payload().inner; + self.in_flight_data_is_whole = + payload.get_ref().remaining() == payload.remaining(); } dst.buffer(frame).expect("invalid frame"); @@ -580,7 +589,11 @@ impl Prioritize { T: AsyncWrite + Unpin, B: Buf, { + // A frame that carries only part of its buffer leaves the rest to be + // queued again on reclaim; HEADERS behind it would end the stream + // before that remainder is sent. if !matches!(self.in_flight_data_frame, InFlightData::DataFrame(_)) + || !self.in_flight_data_is_whole || !dst.can_buffer_headers_after_data() { return false; diff --git a/tests/h2-tests/tests/trailers.rs b/tests/h2-tests/tests/trailers.rs index 2aacd0c68..beeff5d38 100644 --- a/tests/h2-tests/tests/trailers.rs +++ b/tests/h2-tests/tests/trailers.rs @@ -303,3 +303,48 @@ async fn server_sends_large_data_then_trailers_in_order() { join(client, srv).await; } + +#[tokio::test] +async fn server_sends_multi_frame_data_then_trailers_without_truncation() { + // A DATA payload larger than the maximum frame size is split across + // several frames. Trailers queued behind it must follow the last of them, + // never the first. + h2_support::trace_init!(); + let (io, mut client) = mock::new(); + let body = (0..56_263u32).map(|i| i as u8).collect::>(); + let expected = body.clone(); + + let client = async move { + let settings = client.assert_server_handshake().await; + assert_default_settings!(settings); + client + .send_frame( + frames::headers(1) + .request("POST", "https://example.com/") + .eos(), + ) + .await; + client.recv_frame(frames::headers(1).response(200)).await; + for chunk in expected.chunks(16_384) { + client.recv_frame(frames::data(1, chunk.to_vec())).await; + } + client + .recv_frame(frames::headers(1).field("grpc-status", "0").eos()) + .await; + }; + + let srv = async move { + let mut srv = server::handshake(io).await.expect("handshake"); + let (_req, mut stream) = srv.next().await.unwrap().unwrap(); + let rsp = http::Response::builder().status(200).body(()).unwrap(); + let mut send = stream.send_response(rsp, false).unwrap(); + send.reserve_capacity(body.len()); + send.send_data(body.into(), false).unwrap(); + let mut trailers = HeaderMap::new(); + trailers.insert("grpc-status", "0".parse().unwrap()); + send.send_trailers(trailers).unwrap(); + assert!(srv.next().await.is_none()); + }; + + join(client, srv).await; +}