1 use futures_core::future::Future; 2 use futures_core::ready; 3 use futures_core::task::{Context, Poll}; 4 use futures_io::AsyncWrite; 5 use futures_io::IoSlice; 6 use std::io; 7 use std::pin::Pin; 8 9 /// Future for the 10 /// [`write_all_vectored`](super::AsyncWriteExt::write_all_vectored) method. 11 #[derive(Debug)] 12 #[must_use = "futures do nothing unless you `.await` or poll them"] 13 pub struct WriteAllVectored<'a, W: ?Sized + Unpin> { 14 writer: &'a mut W, 15 bufs: &'a mut [IoSlice<'a>], 16 } 17 18 impl<W: ?Sized + Unpin> Unpin for WriteAllVectored<'_, W> {} 19 20 impl<'a, W: AsyncWrite + ?Sized + Unpin> WriteAllVectored<'a, W> { new(writer: &'a mut W, mut bufs: &'a mut [IoSlice<'a>]) -> Self21 pub(super) fn new(writer: &'a mut W, mut bufs: &'a mut [IoSlice<'a>]) -> Self { 22 IoSlice::advance_slices(&mut bufs, 0); 23 Self { writer, bufs } 24 } 25 } 26 27 impl<W: AsyncWrite + ?Sized + Unpin> Future for WriteAllVectored<'_, W> { 28 type Output = io::Result<()>; 29 poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>30 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { 31 let this = &mut *self; 32 while !this.bufs.is_empty() { 33 let n = ready!(Pin::new(&mut this.writer).poll_write_vectored(cx, this.bufs))?; 34 if n == 0 { 35 return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); 36 } else { 37 IoSlice::advance_slices(&mut this.bufs, n); 38 } 39 } 40 41 Poll::Ready(Ok(())) 42 } 43 } 44 45 #[cfg(test)] 46 mod tests { 47 use std::cmp::min; 48 use std::future::Future; 49 use std::io; 50 use std::pin::Pin; 51 use std::task::{Context, Poll}; 52 53 use crate::io::{AsyncWrite, AsyncWriteExt, IoSlice}; 54 use crate::task::noop_waker; 55 56 /// Create a new writer that reads from at most `n_bufs` and reads 57 /// `per_call` bytes (in total) per call to write. test_writer(n_bufs: usize, per_call: usize) -> TestWriter58 fn test_writer(n_bufs: usize, per_call: usize) -> TestWriter { 59 TestWriter { n_bufs, per_call, written: Vec::new() } 60 } 61 62 // TODO: maybe move this the future-test crate? 63 struct TestWriter { 64 n_bufs: usize, 65 per_call: usize, 66 written: Vec<u8>, 67 } 68 69 impl AsyncWrite for TestWriter { poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll<io::Result<usize>>70 fn poll_write( 71 self: Pin<&mut Self>, 72 cx: &mut Context<'_>, 73 buf: &[u8], 74 ) -> Poll<io::Result<usize>> { 75 self.poll_write_vectored(cx, &[IoSlice::new(buf)]) 76 } 77 poll_write_vectored( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll<io::Result<usize>>78 fn poll_write_vectored( 79 mut self: Pin<&mut Self>, 80 _cx: &mut Context<'_>, 81 bufs: &[IoSlice<'_>], 82 ) -> Poll<io::Result<usize>> { 83 let mut left = self.per_call; 84 let mut written = 0; 85 for buf in bufs.iter().take(self.n_bufs) { 86 let n = min(left, buf.len()); 87 self.written.extend_from_slice(&buf[0..n]); 88 left -= n; 89 written += n; 90 } 91 Poll::Ready(Ok(written)) 92 } 93 poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>>94 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { 95 Poll::Ready(Ok(())) 96 } 97 poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>>98 fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { 99 Poll::Ready(Ok(())) 100 } 101 } 102 103 // TODO: maybe move this the future-test crate? 104 macro_rules! assert_poll_ok { 105 ($e:expr, $expected:expr) => { 106 let expected = $expected; 107 match $e { 108 Poll::Ready(Ok(ok)) if ok == expected => {} 109 got => { 110 panic!("unexpected result, got: {:?}, wanted: Ready(Ok({:?}))", got, expected) 111 } 112 } 113 }; 114 } 115 116 #[test] test_writer_read_from_one_buf()117 fn test_writer_read_from_one_buf() { 118 let waker = noop_waker(); 119 let mut cx = Context::from_waker(&waker); 120 121 let mut dst = test_writer(1, 2); 122 let mut dst = Pin::new(&mut dst); 123 124 assert_poll_ok!(dst.as_mut().poll_write(&mut cx, &[]), 0); 125 assert_poll_ok!(dst.as_mut().poll_write_vectored(&mut cx, &[]), 0); 126 127 // Read at most 2 bytes. 128 assert_poll_ok!(dst.as_mut().poll_write(&mut cx, &[1, 1, 1]), 2); 129 let bufs = &[IoSlice::new(&[2, 2, 2])]; 130 assert_poll_ok!(dst.as_mut().poll_write_vectored(&mut cx, bufs), 2); 131 132 // Only read from first buf. 133 let bufs = &[IoSlice::new(&[3]), IoSlice::new(&[4, 4])]; 134 assert_poll_ok!(dst.as_mut().poll_write_vectored(&mut cx, bufs), 1); 135 136 assert_eq!(dst.written, &[1, 1, 2, 2, 3]); 137 } 138 139 #[test] test_writer_read_from_multiple_bufs()140 fn test_writer_read_from_multiple_bufs() { 141 let waker = noop_waker(); 142 let mut cx = Context::from_waker(&waker); 143 144 let mut dst = test_writer(3, 3); 145 let mut dst = Pin::new(&mut dst); 146 147 // Read at most 3 bytes from two buffers. 148 let bufs = &[IoSlice::new(&[1]), IoSlice::new(&[2, 2, 2])]; 149 assert_poll_ok!(dst.as_mut().poll_write_vectored(&mut cx, bufs), 3); 150 151 // Read at most 3 bytes from three buffers. 152 let bufs = &[IoSlice::new(&[3]), IoSlice::new(&[4]), IoSlice::new(&[5, 5])]; 153 assert_poll_ok!(dst.as_mut().poll_write_vectored(&mut cx, bufs), 3); 154 155 assert_eq!(dst.written, &[1, 2, 2, 3, 4, 5]); 156 } 157 158 #[test] test_write_all_vectored()159 fn test_write_all_vectored() { 160 let waker = noop_waker(); 161 let mut cx = Context::from_waker(&waker); 162 163 #[rustfmt::skip] // Becomes unreadable otherwise. 164 let tests: Vec<(_, &'static [u8])> = vec![ 165 (vec![], &[]), 166 (vec![IoSlice::new(&[]), IoSlice::new(&[])], &[]), 167 (vec![IoSlice::new(&[1])], &[1]), 168 (vec![IoSlice::new(&[1, 2])], &[1, 2]), 169 (vec![IoSlice::new(&[1, 2, 3])], &[1, 2, 3]), 170 (vec![IoSlice::new(&[1, 2, 3, 4])], &[1, 2, 3, 4]), 171 (vec![IoSlice::new(&[1, 2, 3, 4, 5])], &[1, 2, 3, 4, 5]), 172 (vec![IoSlice::new(&[1]), IoSlice::new(&[2])], &[1, 2]), 173 (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2])], &[1, 1, 2, 2]), 174 (vec![IoSlice::new(&[1, 1, 1]), IoSlice::new(&[2, 2, 2])], &[1, 1, 1, 2, 2, 2]), 175 (vec![IoSlice::new(&[1, 1, 1, 1]), IoSlice::new(&[2, 2, 2, 2])], &[1, 1, 1, 1, 2, 2, 2, 2]), 176 (vec![IoSlice::new(&[1]), IoSlice::new(&[2]), IoSlice::new(&[3])], &[1, 2, 3]), 177 (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2]), IoSlice::new(&[3, 3])], &[1, 1, 2, 2, 3, 3]), 178 (vec![IoSlice::new(&[1, 1, 1]), IoSlice::new(&[2, 2, 2]), IoSlice::new(&[3, 3, 3])], &[1, 1, 1, 2, 2, 2, 3, 3, 3]), 179 ]; 180 181 for (mut input, wanted) in tests { 182 let mut dst = test_writer(2, 2); 183 { 184 let mut future = dst.write_all_vectored(&mut *input); 185 match Pin::new(&mut future).poll(&mut cx) { 186 Poll::Ready(Ok(())) => {} 187 other => panic!("unexpected result polling future: {:?}", other), 188 } 189 } 190 assert_eq!(&*dst.written, &*wanted); 191 } 192 } 193 } 194