Skip to main content

asupersync/io/
write.rs

1//! AsyncWrite trait and adapters.
2
3use std::io::{self, IoSlice};
4use std::ops::DerefMut;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8/// Async non-blocking write.
9pub trait AsyncWrite {
10    /// Attempt to write data from `buf`.
11    fn poll_write(
12        self: Pin<&mut Self>,
13        cx: &mut Context<'_>,
14        buf: &[u8],
15    ) -> Poll<io::Result<usize>>;
16
17    /// Attempt to write data from multiple buffers (vectored I/O).
18    fn poll_write_vectored(
19        self: Pin<&mut Self>,
20        cx: &mut Context<'_>,
21        bufs: &[IoSlice<'_>],
22    ) -> Poll<io::Result<usize>> {
23        // Default implementation: write first non-empty buffer
24        for buf in bufs {
25            if !buf.is_empty() {
26                return self.poll_write(cx, buf);
27            }
28        }
29        Poll::Ready(Ok(0))
30    }
31
32    /// Returns whether this writer has efficient vectored writes.
33    fn is_write_vectored(&self) -> bool {
34        false
35    }
36
37    /// Attempt to flush buffered data.
38    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>;
39
40    /// Attempt to shutdown the writer.
41    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>;
42}
43
44/// Async non-blocking write from multiple buffers (vectored I/O).
45pub trait AsyncWriteVectored: AsyncWrite {
46    /// Attempt to write data from multiple buffers (vectored I/O).
47    fn poll_write_vectored(
48        self: Pin<&mut Self>,
49        cx: &mut Context<'_>,
50        bufs: &[IoSlice<'_>],
51    ) -> Poll<io::Result<usize>> {
52        AsyncWrite::poll_write_vectored(self, cx, bufs)
53    }
54
55    /// Returns whether this writer has efficient vectored writes.
56    fn is_write_vectored(&self) -> bool {
57        AsyncWrite::is_write_vectored(self)
58    }
59}
60
61impl<W> AsyncWriteVectored for W where W: AsyncWrite + ?Sized {}
62
63impl AsyncWrite for Vec<u8> {
64    fn poll_write(
65        self: Pin<&mut Self>,
66        _cx: &mut Context<'_>,
67        buf: &[u8],
68    ) -> Poll<io::Result<usize>> {
69        let this = self.get_mut();
70        this.extend_from_slice(buf);
71        Poll::Ready(Ok(buf.len()))
72    }
73
74    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
75        Poll::Ready(Ok(()))
76    }
77
78    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
79        Poll::Ready(Ok(()))
80    }
81}
82
83impl AsyncWrite for std::io::Cursor<&mut [u8]> {
84    fn poll_write(
85        self: Pin<&mut Self>,
86        _cx: &mut Context<'_>,
87        buf: &[u8],
88    ) -> Poll<io::Result<usize>> {
89        use std::io::Write as _;
90
91        let this = self.get_mut();
92        let n = this.write(buf)?;
93        Poll::Ready(Ok(n))
94    }
95
96    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
97        Poll::Ready(Ok(()))
98    }
99
100    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
101        Poll::Ready(Ok(()))
102    }
103}
104
105impl AsyncWrite for std::io::Cursor<Vec<u8>> {
106    fn poll_write(
107        self: Pin<&mut Self>,
108        _cx: &mut Context<'_>,
109        buf: &[u8],
110    ) -> Poll<io::Result<usize>> {
111        use std::io::Write as _;
112
113        let this = self.get_mut();
114        let n = this.write(buf)?;
115        Poll::Ready(Ok(n))
116    }
117
118    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
119        Poll::Ready(Ok(()))
120    }
121
122    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
123        Poll::Ready(Ok(()))
124    }
125}
126
127impl AsyncWrite for std::io::Cursor<Box<[u8]>> {
128    fn poll_write(
129        self: Pin<&mut Self>,
130        _cx: &mut Context<'_>,
131        buf: &[u8],
132    ) -> Poll<io::Result<usize>> {
133        use std::io::Write as _;
134
135        let this = self.get_mut();
136        let n = this.write(buf)?;
137        Poll::Ready(Ok(n))
138    }
139
140    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
141        Poll::Ready(Ok(()))
142    }
143
144    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
145        Poll::Ready(Ok(()))
146    }
147}
148
149impl<W> AsyncWrite for &mut W
150where
151    W: AsyncWrite + Unpin + ?Sized,
152{
153    fn poll_write(
154        self: Pin<&mut Self>,
155        cx: &mut Context<'_>,
156        buf: &[u8],
157    ) -> Poll<io::Result<usize>> {
158        let this = self.get_mut();
159        Pin::new(&mut **this).poll_write(cx, buf)
160    }
161
162    fn poll_write_vectored(
163        self: Pin<&mut Self>,
164        cx: &mut Context<'_>,
165        bufs: &[IoSlice<'_>],
166    ) -> Poll<io::Result<usize>> {
167        let this = self.get_mut();
168        Pin::new(&mut **this).poll_write_vectored(cx, bufs)
169    }
170
171    fn is_write_vectored(&self) -> bool {
172        (**self).is_write_vectored()
173    }
174
175    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
176        let this = self.get_mut();
177        Pin::new(&mut **this).poll_flush(cx)
178    }
179
180    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
181        let this = self.get_mut();
182        Pin::new(&mut **this).poll_shutdown(cx)
183    }
184}
185
186impl<W> AsyncWrite for Box<W>
187where
188    W: AsyncWrite + Unpin + ?Sized,
189{
190    fn poll_write(
191        self: Pin<&mut Self>,
192        cx: &mut Context<'_>,
193        buf: &[u8],
194    ) -> Poll<io::Result<usize>> {
195        let this = self.get_mut();
196        Pin::new(&mut **this).poll_write(cx, buf)
197    }
198
199    fn poll_write_vectored(
200        self: Pin<&mut Self>,
201        cx: &mut Context<'_>,
202        bufs: &[IoSlice<'_>],
203    ) -> Poll<io::Result<usize>> {
204        let this = self.get_mut();
205        Pin::new(&mut **this).poll_write_vectored(cx, bufs)
206    }
207
208    fn is_write_vectored(&self) -> bool {
209        (**self).is_write_vectored()
210    }
211
212    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
213        let this = self.get_mut();
214        Pin::new(&mut **this).poll_flush(cx)
215    }
216
217    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
218        let this = self.get_mut();
219        Pin::new(&mut **this).poll_shutdown(cx)
220    }
221}
222
223impl<W, P> AsyncWrite for Pin<P>
224where
225    P: DerefMut<Target = W> + Unpin,
226    W: AsyncWrite + Unpin + ?Sized,
227{
228    fn poll_write(
229        self: Pin<&mut Self>,
230        cx: &mut Context<'_>,
231        buf: &[u8],
232    ) -> Poll<io::Result<usize>> {
233        let this = self.get_mut();
234        Pin::new(&mut **this).poll_write(cx, buf)
235    }
236
237    fn poll_write_vectored(
238        self: Pin<&mut Self>,
239        cx: &mut Context<'_>,
240        bufs: &[IoSlice<'_>],
241    ) -> Poll<io::Result<usize>> {
242        let this = self.get_mut();
243        Pin::new(&mut **this).poll_write_vectored(cx, bufs)
244    }
245
246    fn is_write_vectored(&self) -> bool {
247        (**self).is_write_vectored()
248    }
249
250    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
251        let this = self.get_mut();
252        Pin::new(&mut **this).poll_flush(cx)
253    }
254
255    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
256        let this = self.get_mut();
257        Pin::new(&mut **this).poll_shutdown(cx)
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use std::sync::Arc;
265    use std::task::{Context, Wake, Waker};
266
267    struct NoopWaker;
268
269    impl Wake for NoopWaker {
270        fn wake(self: Arc<Self>) {}
271    }
272
273    fn noop_waker() -> Waker {
274        Waker::from(Arc::new(NoopWaker))
275    }
276
277    fn init_test(name: &str) {
278        crate::test_utils::init_test_logging();
279        crate::test_phase!(name);
280    }
281
282    #[test]
283    fn write_to_vec() {
284        init_test("write_to_vec");
285        let mut output = Vec::new();
286        let waker = noop_waker();
287        let mut cx = Context::from_waker(&waker);
288
289        let poll = Pin::new(&mut output).poll_write(&mut cx, b"hello");
290        let ready = matches!(poll, Poll::Ready(Ok(5)));
291        crate::assert_with_log!(ready, "write 5", true, ready);
292        crate::assert_with_log!(output == b"hello", "output", b"hello", output);
293        crate::test_complete!("write_to_vec");
294    }
295
296    #[test]
297    fn write_to_cursor() {
298        init_test("write_to_cursor");
299        let mut buf = [0u8; 8];
300        let mut cursor = std::io::Cursor::new(&mut buf[..]);
301        let waker = noop_waker();
302        let mut cx = Context::from_waker(&waker);
303
304        let poll = Pin::new(&mut cursor).poll_write(&mut cx, b"test");
305        let ready = matches!(poll, Poll::Ready(Ok(4)));
306        crate::assert_with_log!(ready, "write 4", true, ready);
307        crate::assert_with_log!(&buf[..4] == b"test", "buf", b"test", &buf[..4]);
308        crate::test_complete!("write_to_cursor");
309    }
310
311    #[test]
312    fn flush_and_shutdown_vec() {
313        init_test("flush_and_shutdown_vec");
314        let mut output = Vec::new();
315        let waker = noop_waker();
316        let mut cx = Context::from_waker(&waker);
317
318        let poll = Pin::new(&mut output).poll_flush(&mut cx);
319        let ready = matches!(poll, Poll::Ready(Ok(())));
320        crate::assert_with_log!(ready, "flush ready", true, ready);
321
322        let poll = Pin::new(&mut output).poll_shutdown(&mut cx);
323        let ready = matches!(poll, Poll::Ready(Ok(())));
324        crate::assert_with_log!(ready, "shutdown ready", true, ready);
325        crate::test_complete!("flush_and_shutdown_vec");
326    }
327
328    #[test]
329    fn write_via_ref() {
330        init_test("write_via_ref");
331        let mut output = Vec::new();
332        let mut writer = &mut output;
333        let waker = noop_waker();
334        let mut cx = Context::from_waker(&waker);
335
336        let poll = Pin::new(&mut writer).poll_write(&mut cx, b"via ref");
337        let ready = matches!(poll, Poll::Ready(Ok(7)));
338        crate::assert_with_log!(ready, "write 7", true, ready);
339        crate::assert_with_log!(output == b"via ref", "output", b"via ref", output);
340        crate::test_complete!("write_via_ref");
341    }
342
343    #[test]
344    fn write_via_box() {
345        init_test("write_via_box");
346        let mut output: Box<Vec<u8>> = Box::default();
347        let waker = noop_waker();
348        let mut cx = Context::from_waker(&waker);
349
350        let poll = Pin::new(&mut output).poll_write(&mut cx, b"boxed");
351        let ready = matches!(poll, Poll::Ready(Ok(5)));
352        crate::assert_with_log!(ready, "write 5", true, ready);
353        crate::assert_with_log!(*output == b"boxed", "output", b"boxed", *output);
354        crate::test_complete!("write_via_box");
355    }
356}