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..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, } } @@ -526,6 +532,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); } @@ -541,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"); @@ -557,6 +569,64 @@ 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, + { + // 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; + } + 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..beeff5d38 100644 --- a/tests/h2-tests/tests/trailers.rs +++ b/tests/h2-tests/tests/trailers.rs @@ -262,3 +262,89 @@ 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; +} + +#[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; +}