diff --git a/h3-quinn/src/lib.rs b/h3-quinn/src/lib.rs index 4e0f9264..569a37e4 100644 --- a/h3-quinn/src/lib.rs +++ b/h3-quinn/src/lib.rs @@ -6,7 +6,8 @@ use std::{ convert::TryInto, future::Future, - pin::Pin, + pin::{pin, Pin}, + slice, sync::Arc, task::{self, ready, Poll}, }; @@ -496,7 +497,7 @@ fn convert_write_error_to_stream_error(error: quinn::WriteError) -> StreamErrorI /// Implements a [`quic::SendStream`] backed by a [`quinn::SendStream`]. pub struct SendStream { stream: quinn::SendStream, - writing: Option>, + writing: Option<(WriteBuf, Bytes)>, } impl SendStream @@ -517,12 +518,13 @@ where { #[cfg_attr(feature = "tracing", instrument(skip_all, level = "trace"))] fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll> { - if let Some(ref mut data) = self.writing { - while data.has_remaining() { - let stream = Pin::new(&mut self.stream); - let written = ready!(stream.poll_write(cx, data.chunk())) + if let Some((data, chunk)) = &mut self.writing { + while !chunk.is_empty() || data.has_remaining() { + if chunk.is_empty() { + *chunk = data.copy_to_bytes(data.chunk().len()); + } + ready!(pin!(self.stream.write_chunks(slice::from_mut(chunk))).poll(cx)) .map_err(convert_write_error_to_stream_error)?; - data.advance(written); } } // all data is written @@ -563,7 +565,7 @@ where ), }); } - self.writing = Some(data.into()); + self.writing = Some((data.into(), Bytes::new())); Ok(()) } diff --git a/h3/src/stream.rs b/h3/src/stream.rs index a3a823f6..2e5c4f9c 100644 --- a/h3/src/stream.rs +++ b/h3/src/stream.rs @@ -4,7 +4,7 @@ use std::{ task::{Context, Poll}, }; -use bytes::{Buf, BufMut, Bytes}; +use bytes::{Buf, BufMut, Bytes, BytesMut}; use futures_util::{future, ready}; use pin_project_lite::pin_project; use tokio::io::ReadBuf; @@ -253,6 +253,18 @@ where payload.advance(cnt); } } + + fn copy_to_bytes(&mut self, len: usize) -> Bytes { + if self.len == self.pos { + if let Some(payload) = self.frame.as_mut().and_then(|f| f.payload_mut()) { + return payload.copy_to_bytes(len); + } + } + assert!(len <= self.remaining(), "`len` greater than remaining"); + let mut ret = BytesMut::with_capacity(len); + ret.put(self.take(len)); + ret.freeze() + } } pub(super) enum AcceptedRecvStream @@ -791,4 +803,15 @@ mod tests { wbuf.advance(4); assert_eq!(wbuf.chunk(), b"ey"); } + + #[test] + fn write_buf_copy_to_bytes_shares_payload_after_header() { + let payload = Bytes::from("hey"); + let mut wbuf = WriteBuf::::from(Frame::Data(payload.clone())); + + assert_eq!(wbuf.copy_to_bytes(3), &b"\x00\x03h"[..]); + let rest = wbuf.copy_to_bytes(2); + assert_eq!(rest, "ey"); + assert_eq!(rest.as_ptr(), payload[1..].as_ptr()); + } }