io: advance partially written buffers correctly in write_all_vectored (#8159)

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`

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<io::Result<usize>> {
+        panic!("shouldn't be called")
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Ok(()).into()
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Ok(()).into()
+    }
+
+    fn poll_write_vectored(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<Result<usize, io::Error>> {
+        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");
+}