tokio_util/io/
std_adapter.rs

1use pin_project_lite::pin_project;
2use std::future::Future;
3use std::marker::PhantomData;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use tokio::io::AsyncWrite;
7
8/// Adapts an implementor of [`AsyncWrite`] to be usable in some context requiring [`std::io::Write`]
9///
10/// **Important**: this adapter is limited to a specific scenario where the calling code supports
11/// restarting. In particular, the calling code MUST NOT call the `write_all` method on the writer.
12#[derive(Debug)]
13pub struct StdWriteAdapter<'a, 'b, W> {
14    writer: W,
15    cx: &'a mut Context<'b>,
16}
17
18impl<'a, 'b, W: AsyncWrite> StdWriteAdapter<'a, 'b, W> {
19    /// Constructs the adapter.
20    ///
21    /// This method produces a restricted implementation of [`std::io::Write`] and may only be
22    /// passed to functions that properly implement restarting and don't block. This property is
23    /// neither checked by the type system nor required by the `Write` trait and thus it has to be
24    /// explicitly guaranteed by each function that takes the returned value as argument.
25    pub fn new_assuming_correct_usage(writer: W, cx: &'a mut Context<'b>) -> Self {
26        Self {
27            writer,
28            cx,
29        }
30    }
31}
32
33impl<'a, 'b, W: AsyncWrite + Unpin> std::io::Write for StdWriteAdapter<'a, 'b, W> {
34    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
35        match Pin::new(&mut self.writer).poll_write(self.cx, buf) {
36            Poll::Ready(result) => result,
37            Poll::Pending => Err(std::io::ErrorKind::WouldBlock.into()),
38        }
39    }
40
41    /// Calling this method would be wrong therefore it always panics.
42    #[track_caller]
43    fn write_all(&mut self, _: &[u8]) -> std::io::Result<()> {
44        panic!("operation not supported by the adapter")
45    }
46
47    fn flush(&mut self) -> std::io::Result<()> {
48        match Pin::new(&mut self.writer).poll_flush(self.cx) {
49            Poll::Ready(result) => result,
50            Poll::Pending => Err(std::io::ErrorKind::WouldBlock.into()),
51        }
52    }
53}
54
55pin_project! {
56    /// An asynchronous operation implemented using a restricted function taking `std::io::Write`.
57    ///
58    /// This is a [`Future`] that calls `F` each time it's polled, passing in a restricted writer.
59    /// The closure `F` may pass this writer to other functions which require `std::io::Write` as
60    /// long as those functions guarantee correct restarting in the presence of `WouldBlock`
61    /// errors.
62    ///
63    /// E.g. given `fn perform_single_write(&mut self, writer: impl Write) -> io::Result<()>` which
64    /// guarantees to keep consistent state in `self` even if `writer` returns error, one might
65    /// write code like this:
66    ///
67    /// ```no-compile
68    /// StdWriteAdapterFuture::new_assuming_correct_usage(writer, |writer| {
69    ///     if self.needs_write() {
70    ///         match self.perform_single_write() {
71    ///             Ok(()) => (),
72    ///             Err(error) if error.kind() == ErrorKind::WouldBlock => {
73    ///                 Poll::Pending
74    ///             }
75    ///             Err(error) => Poll::Ready(Err(error)),
76    ///         }
77    ///     } else {
78    ///         Poll::Ready(Ok(()))
79    ///     }
80    /// })
81    /// ```
82    pub struct StdWriteAdapterFuture<W, T, F> {
83        #[pin]
84        writer: W,
85        f: F,
86        _phantom: PhantomData<fn() -> T>,
87    }
88}
89
90impl<W: AsyncWrite, T, F: FnMut(StdWriteAdapter<'_, '_, Pin<&mut W>>) -> Poll<T>> StdWriteAdapterFuture<W, T, F> {
91    /// Constructs the future.
92    ///
93    /// It is required that the closure `f` only passes the writer given to it to functions that
94    /// guarantee correct restarting in case of `WouldBlock` errors. In particular, the `write_all`
95    /// method on the writer MUST NOT be called (it will panic) but also any other code that
96    /// behaves similarly by keeping the state on the stack and losing it if error happens (easy to
97    /// do by accident unless specific care was taken to avoid it).
98    pub fn new_assuming_correct_usage(writer: W, f: F) -> Self {
99        Self {
100            writer,
101            f,
102            _phantom: PhantomData,
103        }
104    }
105}
106
107impl<W: AsyncWrite, T, F: FnMut(StdWriteAdapter<'_, '_, Pin<&mut W>>) -> Poll<T>> Future for StdWriteAdapterFuture<W, T, F> {
108    type Output = T;
109
110    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
111        let this = self.project();
112        (this.f)(StdWriteAdapter::new_assuming_correct_usage(this.writer, cx))
113    }
114}
tokio_util/io/std_adapter.rs

tokio_util/io/
std_adapter.rs
