From 0531549b6ea66d22d301044910ddc6a77e8c7f1e Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 29 Jun 2021 21:26:47 +0300 Subject: [PATCH] io: efficient implementation of vectored writes for BufWriter (#3163) --- tokio/src/io/util/buf_reader.rs | 14 +- tokio/src/io/util/buf_stream.rs | 14 +- tokio/src/io/util/buf_writer.rs | 68 +++++++- tokio/tests/io_buf_writer.rs | 286 ++++++++++++++++++++++++++++++++ tokio/tests/support/io_vec.rs | 45 +++++ 5 files changed, 424 insertions(+), 3 deletions(-) create mode 100644 tokio/tests/support/io_vec.rs diff --git a/tokio/src/io/util/buf_reader.rs b/tokio/src/io/util/buf_reader.rs index c4d6842d4..7cfd46ce0 100644 --- a/tokio/src/io/util/buf_reader.rs +++ b/tokio/src/io/util/buf_reader.rs @@ -2,7 +2,7 @@ use crate::io::util::DEFAULT_BUF_SIZE; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; -use std::io::{self, SeekFrom}; +use std::io::{self, IoSlice, SeekFrom}; use std::pin::Pin; use std::task::{Context, Poll}; use std::{cmp, fmt, mem}; @@ -268,6 +268,18 @@ impl AsyncWrite for BufReader { self.get_pin_mut().poll_write(cx, buf) } + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + self.get_pin_mut().poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.get_ref().is_write_vectored() + } + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_pin_mut().poll_flush(cx) } diff --git a/tokio/src/io/util/buf_stream.rs b/tokio/src/io/util/buf_stream.rs index ff3d9dba8..595c142ac 100644 --- a/tokio/src/io/util/buf_stream.rs +++ b/tokio/src/io/util/buf_stream.rs @@ -2,7 +2,7 @@ use crate::io::util::{BufReader, BufWriter}; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; -use std::io::{self, SeekFrom}; +use std::io::{self, IoSlice, SeekFrom}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -127,6 +127,18 @@ impl AsyncWrite for BufStream { self.project().inner.poll_write(cx, buf) } + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + self.project().inner.poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_flush(cx) } diff --git a/tokio/src/io/util/buf_writer.rs b/tokio/src/io/util/buf_writer.rs index 4e8e493ce..8dd1bba60 100644 --- a/tokio/src/io/util/buf_writer.rs +++ b/tokio/src/io/util/buf_writer.rs @@ -3,7 +3,7 @@ use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; use std::fmt; -use std::io::{self, SeekFrom, Write}; +use std::io::{self, IoSlice, SeekFrom, Write}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -133,6 +133,72 @@ impl AsyncWrite for BufWriter { } } + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + mut bufs: &[IoSlice<'_>], + ) -> Poll> { + if self.inner.is_write_vectored() { + let total_len = bufs + .iter() + .fold(0usize, |acc, b| acc.saturating_add(b.len())); + if total_len > self.buf.capacity() - self.buf.len() { + ready!(self.as_mut().flush_buf(cx))?; + } + let me = self.as_mut().project(); + if total_len >= me.buf.capacity() { + // It's more efficient to pass the slices directly to the + // underlying writer than to buffer them. + // The case when the total_len calculation saturates at + // usize::MAX is also handled here. + me.inner.poll_write_vectored(cx, bufs) + } else { + bufs.iter().for_each(|b| me.buf.extend_from_slice(b)); + Poll::Ready(Ok(total_len)) + } + } else { + // Remove empty buffers at the beginning of bufs. + while bufs.first().map(|buf| buf.len()) == Some(0) { + bufs = &bufs[1..]; + } + if bufs.is_empty() { + return Poll::Ready(Ok(0)); + } + // Flush if the first buffer doesn't fit. + let first_len = bufs[0].len(); + if first_len > self.buf.capacity() - self.buf.len() { + ready!(self.as_mut().flush_buf(cx))?; + debug_assert!(self.buf.is_empty()); + } + let me = self.as_mut().project(); + if first_len >= me.buf.capacity() { + // The slice is at least as large as the buffering capacity, + // so it's better to write it directly, bypassing the buffer. + debug_assert!(me.buf.is_empty()); + return me.inner.poll_write(cx, &bufs[0]); + } else { + me.buf.extend_from_slice(&bufs[0]); + bufs = &bufs[1..]; + } + let mut total_written = first_len; + debug_assert!(total_written != 0); + // Append the buffers that fit in the internal buffer. + for buf in bufs { + if buf.len() > me.buf.capacity() - me.buf.len() { + break; + } else { + me.buf.extend_from_slice(buf); + total_written += buf.len(); + } + } + Poll::Ready(Ok(total_written)) + } + } + + fn is_write_vectored(&self) -> bool { + true + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(self.as_mut().flush_buf(cx))?; self.get_pin_mut().poll_flush(cx) diff --git a/tokio/tests/io_buf_writer.rs b/tokio/tests/io_buf_writer.rs index 6f4f10a8e..47a0d466f 100644 --- a/tokio/tests/io_buf_writer.rs +++ b/tokio/tests/io_buf_writer.rs @@ -8,6 +8,17 @@ use std::io::{self, Cursor}; use std::pin::Pin; use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BufWriter, SeekFrom}; +use futures::future; +use tokio_test::assert_ok; + +use std::cmp; +use std::io::IoSlice; + +mod support { + pub(crate) mod io_vec; +} +use support::io_vec::IoBufs; + struct MaybePending { inner: Vec, ready: bool, @@ -47,6 +58,14 @@ impl AsyncWrite for MaybePending { } } +async fn write_vectored(writer: &mut W, bufs: &[IoSlice<'_>]) -> io::Result +where + W: AsyncWrite + Unpin, +{ + let mut writer = Pin::new(writer); + future::poll_fn(|cx| writer.as_mut().poll_write_vectored(cx, bufs)).await +} + #[tokio::test] async fn buf_writer() { let mut writer = BufWriter::with_capacity(2, Vec::new()); @@ -249,3 +268,270 @@ async fn maybe_pending_buf_writer_seek() { &[0, 1, 8, 9, 4, 5, 6, 7] ); } + +struct MockWriter { + data: Vec, + write_len: usize, + vectored: bool, +} + +impl MockWriter { + fn new(write_len: usize) -> Self { + MockWriter { + data: Vec::new(), + write_len, + vectored: false, + } + } + + fn vectored(write_len: usize) -> Self { + MockWriter { + data: Vec::new(), + write_len, + vectored: true, + } + } + + fn write_up_to(&mut self, buf: &[u8], limit: usize) -> usize { + let len = cmp::min(buf.len(), limit); + self.data.extend_from_slice(&buf[..len]); + len + } +} + +impl AsyncWrite for MockWriter { + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let n = this.write_up_to(buf, this.write_len); + Ok(n).into() + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + _: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + let this = self.get_mut(); + let mut total_written = 0; + for buf in bufs { + let n = this.write_up_to(buf, this.write_len - total_written); + total_written += n; + if total_written == this.write_len { + break; + } + } + Ok(total_written).into() + } + + fn is_write_vectored(&self) -> bool { + self.vectored + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Ok(()).into() + } +} + +#[tokio::test] +async fn write_vectored_empty_on_non_vectored() { + let mut w = BufWriter::new(MockWriter::new(4)); + let n = assert_ok!(write_vectored(&mut w, &[]).await); + assert_eq!(n, 0); + + let io_vec = [IoSlice::new(&[]); 3]; + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 0); + + assert_ok!(w.flush().await); + assert!(w.get_ref().data.is_empty()); +} + +#[tokio::test] +async fn write_vectored_empty_on_vectored() { + let mut w = BufWriter::new(MockWriter::vectored(4)); + let n = assert_ok!(write_vectored(&mut w, &[]).await); + assert_eq!(n, 0); + + let io_vec = [IoSlice::new(&[]); 3]; + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 0); + + assert_ok!(w.flush().await); + assert!(w.get_ref().data.is_empty()); +} + +#[tokio::test] +async fn write_vectored_basic_on_non_vectored() { + let msg = b"foo bar baz"; + let bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&msg[4..8]), + IoSlice::new(&msg[8..]), + ]; + let mut w = BufWriter::new(MockWriter::new(4)); + let n = assert_ok!(write_vectored(&mut w, &bufs).await); + assert_eq!(n, msg.len()); + assert!(w.buffer() == &msg[..]); + assert_ok!(w.flush().await); + assert_eq!(w.get_ref().data, msg); +} + +#[tokio::test] +async fn write_vectored_basic_on_vectored() { + let msg = b"foo bar baz"; + let bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&msg[4..8]), + IoSlice::new(&msg[8..]), + ]; + let mut w = BufWriter::new(MockWriter::vectored(4)); + let n = assert_ok!(write_vectored(&mut w, &bufs).await); + assert_eq!(n, msg.len()); + assert!(w.buffer() == &msg[..]); + assert_ok!(w.flush().await); + assert_eq!(w.get_ref().data, msg); +} + +#[tokio::test] +async fn write_vectored_large_total_on_non_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&msg[4..8]), + IoSlice::new(&msg[8..]), + ]; + let io_vec = IoBufs::new(&mut bufs); + let mut w = BufWriter::with_capacity(8, MockWriter::new(4)); + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 8); + assert!(w.buffer() == &msg[..8]); + let io_vec = io_vec.advance(n); + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 3); + assert!(w.get_ref().data.as_slice() == &msg[..8]); + assert!(w.buffer() == &msg[8..]); +} + +#[tokio::test] +async fn write_vectored_large_total_on_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&msg[4..8]), + IoSlice::new(&msg[8..]), + ]; + let io_vec = IoBufs::new(&mut bufs); + let mut w = BufWriter::with_capacity(8, MockWriter::vectored(10)); + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 10); + assert!(w.buffer().is_empty()); + let io_vec = io_vec.advance(n); + let n = assert_ok!(write_vectored(&mut w, &io_vec).await); + assert_eq!(n, 1); + assert!(w.get_ref().data.as_slice() == &msg[..10]); + assert!(w.buffer() == &msg[10..]); +} + +struct VectoredWriteHarness { + writer: BufWriter, + buf_capacity: usize, +} + +impl VectoredWriteHarness { + fn new(buf_capacity: usize) -> Self { + VectoredWriteHarness { + writer: BufWriter::with_capacity(buf_capacity, MockWriter::new(4)), + buf_capacity, + } + } + + fn with_vectored_backend(buf_capacity: usize) -> Self { + VectoredWriteHarness { + writer: BufWriter::with_capacity(buf_capacity, MockWriter::vectored(4)), + buf_capacity, + } + } + + async fn write_all<'a, 'b>(&mut self, mut io_vec: IoBufs<'a, 'b>) -> usize { + let mut total_written = 0; + while !io_vec.is_empty() { + let n = assert_ok!(write_vectored(&mut self.writer, &io_vec).await); + assert!(n != 0); + assert!(self.writer.buffer().len() <= self.buf_capacity); + total_written += n; + io_vec = io_vec.advance(n); + } + total_written + } + + async fn flush(&mut self) -> &[u8] { + assert_ok!(self.writer.flush().await); + &self.writer.get_ref().data + } +} + +#[tokio::test] +async fn write_vectored_odd_on_non_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&[]), + IoSlice::new(&msg[4..9]), + IoSlice::new(&msg[9..]), + ]; + let mut h = VectoredWriteHarness::new(8); + let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; + assert_eq!(bytes_written, msg.len()); + assert_eq!(h.flush().await, msg); +} + +#[tokio::test] +async fn write_vectored_odd_on_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&msg[0..4]), + IoSlice::new(&[]), + IoSlice::new(&msg[4..9]), + IoSlice::new(&msg[9..]), + ]; + let mut h = VectoredWriteHarness::with_vectored_backend(8); + let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; + assert_eq!(bytes_written, msg.len()); + assert_eq!(h.flush().await, msg); +} + +#[tokio::test] +async fn write_vectored_large_slice_on_non_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&[]), + IoSlice::new(&msg[..9]), + IoSlice::new(&msg[9..]), + ]; + let mut h = VectoredWriteHarness::new(8); + let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; + assert_eq!(bytes_written, msg.len()); + assert_eq!(h.flush().await, msg); +} + +#[tokio::test] +async fn write_vectored_large_slice_on_vectored() { + let msg = b"foo bar baz"; + let mut bufs = [ + IoSlice::new(&[]), + IoSlice::new(&msg[..9]), + IoSlice::new(&msg[9..]), + ]; + let mut h = VectoredWriteHarness::with_vectored_backend(8); + let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; + assert_eq!(bytes_written, msg.len()); + assert_eq!(h.flush().await, msg); +} diff --git a/tokio/tests/support/io_vec.rs b/tokio/tests/support/io_vec.rs new file mode 100644 index 000000000..4ea47c748 --- /dev/null +++ b/tokio/tests/support/io_vec.rs @@ -0,0 +1,45 @@ +use std::io::IoSlice; +use std::ops::Deref; +use std::slice; + +pub struct IoBufs<'a, 'b>(&'b mut [IoSlice<'a>]); + +impl<'a, 'b> IoBufs<'a, 'b> { + pub fn new(slices: &'b mut [IoSlice<'a>]) -> Self { + IoBufs(slices) + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn advance(mut self, n: usize) -> IoBufs<'a, 'b> { + let mut to_remove = 0; + let mut remaining_len = n; + for slice in self.0.iter() { + if remaining_len < slice.len() { + break; + } else { + remaining_len -= slice.len(); + to_remove += 1; + } + } + self.0 = self.0.split_at_mut(to_remove).1; + if let Some(slice) = self.0.first_mut() { + let tail = &slice[remaining_len..]; + // Safety: recasts slice to the original lifetime + let tail = unsafe { slice::from_raw_parts(tail.as_ptr(), tail.len()) }; + *slice = IoSlice::new(tail); + } else if remaining_len != 0 { + panic!("advance past the end of the slice vector"); + } + self + } +} + +impl<'a, 'b> Deref for IoBufs<'a, 'b> { + type Target = [IoSlice<'a>]; + fn deref(&self) -> &[IoSlice<'a>] { + self.0 + } +}