diff --git a/tokio-util/src/io/write_all_vectored.rs b/tokio-util/src/io/write_all_vectored.rs index dbc59a0b7..c2b19a19c 100644 --- a/tokio-util/src/io/write_all_vectored.rs +++ b/tokio-util/src/io/write_all_vectored.rs @@ -147,7 +147,7 @@ fn advance_slices<'a>(bufs: &mut &mut [IoSlice<'a>], n: usize) { *bufs = &mut std::mem::take(bufs)[remove..]; if let Some(first) = bufs.first_mut() { - let buf = &first[..left]; + let buf = &first[left..]; // necessary due to limitating in the borrow checker, // when tokio MSRV reaches 1.81.0 this entire function // can be replaced with `IoSlice::advance_slices` diff --git a/tokio-util/tests/io_write_all_vectored.rs b/tokio-util/tests/io_write_all_vectored.rs index 5ba0e62db..300903037 100644 --- a/tokio-util/tests/io_write_all_vectored.rs +++ b/tokio-util/tests/io_write_all_vectored.rs @@ -10,6 +10,55 @@ use std::io::IoSlice; use std::pin::Pin; use std::task::{Context, Poll}; +struct PartialVectoredWriter { + buf: BytesMut, + max_write: usize, +} + +impl AsyncWrite for PartialVectoredWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + panic!("shouldn't be called") + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + let mut remaining = self.max_write; + let mut written = 0; + + for buf in bufs { + if remaining == 0 { + break; + } + + let n = remaining.min(buf.len()); + self.buf.extend_from_slice(&buf[..n]); + remaining -= n; + written += n; + } + + Ok(written).into() + } + + fn is_write_vectored(&self) -> bool { + true + } +} + #[tokio::test] async fn test_write_all_vectored() { struct Wr { @@ -140,3 +189,29 @@ async fn write_all_vectored_with_empty_slice() { write_all_vectored(&mut wr, buf).await.unwrap(); assert_eq!(&wr.buf[..], b"hello"); } + +#[tokio::test] +async fn write_all_vectored_should_continue_with_unwritten_suffix_if_write_stops_inside_buffer() { + let mut wr = PartialVectoredWriter { + buf: BytesMut::with_capacity(64), + max_write: 3, + }; + + let buf = &mut [IoSlice::new(b"hello"), IoSlice::new(b"world")]; + + write_all_vectored(&mut wr, buf).await.unwrap(); + assert_eq!(&wr.buf[..], b"helloworld"); +} + +#[tokio::test] +async fn write_all_vectored_should_continue_with_next_buffer_if_write_ends_on_boundary() { + let mut wr = PartialVectoredWriter { + buf: BytesMut::with_capacity(64), + max_write: 2, + }; + + let buf = &mut [IoSlice::new(b"ab"), IoSlice::new(b"cd")]; + + write_all_vectored(&mut wr, buf).await.unwrap(); + assert_eq!(&wr.buf[..], b"abcd"); +}