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