maybe_fut/api/io/
write.rs

1use std::io::IoSlice;
2
3/// A trait for objects which are byte-oriented sinks.
4///
5/// Implementors of the [`Write`] trait are called 'writers'.
6///
7/// Writers are defined by two required methods, `write` and `flush`:
8///
9/// - The `write` method will attempt to write some data into the object, returning how many bytes were successfully written.
10/// - The `flush` method is useful for adapters and explicit buffers themselves for ensuring that all buffered data has been pushed out to the ‘true sink’.
11pub trait Write {
12    /// Writes a buffer into this writer, returning how many bytes were successfully written.
13    fn write(&mut self, buf: &[u8]) -> impl Future<Output = std::io::Result<usize>>;
14
15    /// Flushes the output streamer, ensuring that all intermediately buffered contents reach their destination.
16    fn flush(&mut self) -> impl Future<Output = std::io::Result<()>>;
17
18    /// Like `write`, except that it writes from a slice of buffers.
19    fn write_vectored(
20        &mut self,
21        bufs: &[IoSlice<'_>],
22    ) -> impl Future<Output = std::io::Result<usize>> {
23        async move {
24            let mut total = 0;
25            for buf in bufs.iter() {
26                let n = self.write(buf).await?;
27                total += n;
28            }
29            Ok(total)
30        }
31    }
32
33    /// Attempts to write an entire buffer into this writer.
34    fn write_all(&mut self, mut buf: &[u8]) -> impl Future<Output = std::io::Result<()>> {
35        async move {
36            while !buf.is_empty() {
37                let n = self.write(buf).await?;
38                if n == 0 {
39                    break;
40                } else {
41                    buf = &buf[n..];
42                }
43            }
44            Ok(())
45        }
46    }
47}
48
49#[cfg(test)]
50mod test {
51
52    use super::*;
53
54    struct MockWriter {
55        data: Vec<u8>,
56    }
57
58    impl Write for MockWriter {
59        async fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
60            let n = buf.len();
61            self.data.extend_from_slice(buf);
62            Ok(n)
63        }
64
65        async fn flush(&mut self) -> std::io::Result<()> {
66            Ok(())
67        }
68    }
69
70    #[tokio::test]
71    async fn test_write() {
72        let mut writer = MockWriter { data: Vec::new() };
73        let buf = b"Hello, world!";
74        let n = writer.write(buf).await.unwrap();
75        assert_eq!(n, buf.len());
76        assert_eq!(writer.data, buf);
77    }
78
79    #[tokio::test]
80    async fn test_write_vectored() {
81        let mut writer = MockWriter { data: Vec::new() };
82        let bufs = [b"Hello,", b"world!"];
83        let slices = bufs
84            .into_iter()
85            .map(|b| IoSlice::new(b))
86            .collect::<Vec<_>>();
87        let n = writer.write_vectored(&slices).await.unwrap();
88        assert_eq!(n, bufs.iter().map(|b| b.len()).sum::<usize>());
89        assert_eq!(writer.data, b"Hello,world!");
90    }
91
92    #[tokio::test]
93    async fn test_write_all() {
94        let mut writer = MockWriter { data: Vec::new() };
95        let buf = b"Hello, world!";
96        writer.write_all(buf).await.unwrap();
97        assert_eq!(writer.data, buf);
98    }
99}