Add AsyncWriteExt::into_sink

Author: Nemo157

futures-util/src/io/into_sink.rs

@@ -0,0 +1,115 @@
+use futures_core::task::{Context, Poll};
+use futures_io::AsyncWrite;
+use futures_sink::Sink;
+use std::fmt;
+use std::io;
+use std::pin::Pin;
+use std::marker::Unpin;
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// Sink for the [`into_sink`](super::AsyncWriteExt::into_sink) method.
+#[must_use = "sinks do nothing unless polled"]
+pub struct IntoSink<W> {
+    writer: W,
+    /// An outstanding block for us to push into the underlying writer, along with an index of how
+    /// far into the block we have written already.
+    buffer: Option<(usize, Box<dyn AsRef<[u8]>>)>,
+}
+
+impl<W: Unpin> Unpin for IntoSink<W> {}
+
+impl<W: AsyncWrite> IntoSink<W> {
+    unsafe_pinned!(writer: W);
+    unsafe_unpinned!(buffer: Option<(usize, Box<dyn AsRef<[u8]>>)>);
+
+    pub(super) fn new(writer: W) -> Self {
+        IntoSink { writer, buffer: None }
+    }
+
+    fn project<'a>(self: Pin<&'a mut Self>) -> (Pin<&'a mut W>, &'a mut Option<(usize, Box<dyn AsRef<[u8]>>)>) {
+        unsafe {
+            let this = self.get_unchecked_mut();
+            (Pin::new_unchecked(&mut this.writer), &mut this.buffer)
+        }
+    }
+
+    /// If we have an outstanding block in `buffer` attempt to push it into the writer, does _not_
+    /// flush the writer after it succeeds in pushing the block into it.
+    fn poll_flush_buffer(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), io::Error>>
+    {
+        let (mut writer, buffer) = self.project();
+        if let Some((index, block)) = buffer {
+            loop {
+                let bytes = (**block).as_ref();
+                let written = ready!(writer.as_mut().poll_write(cx, &bytes[*index..]))?;
+                *index += written;
+                if *index == bytes.len() {
+                    break;
+                }
+            }
+            *buffer = None;
+        }
+        Poll::Ready(Ok(()))
+    }
+
+}
+
+impl<W: AsyncWrite, Item: AsRef<[u8]> + 'static> Sink<Item> for IntoSink<W> {
+    type SinkError = io::Error;
+
+    fn poll_ready(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), Self::SinkError>>
+    {
+        ready!(self.as_mut().poll_flush_buffer(cx))?;
+        Poll::Ready(Ok(()))
+    }
+
+    fn start_send(
+        mut self: Pin<&mut Self>,
+        item: Item,
+    ) -> Result<(), Self::SinkError>
+    {
+        debug_assert!(self.as_mut().buffer().is_none());
+        *self.as_mut().buffer() = Some((0, Box::new(item)));
+        Ok(())
+    }
+
+    fn poll_flush(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), Self::SinkError>>
+    {
+        ready!(self.as_mut().poll_flush_buffer(cx))?;
+        ready!(self.as_mut().writer().poll_flush(cx))?;
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_close(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), Self::SinkError>>
+    {
+        ready!(self.as_mut().poll_flush_buffer(cx))?;
+        ready!(self.as_mut().writer().poll_close(cx))?;
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl<W: fmt::Debug> fmt::Debug for IntoSink<W> {
+    
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let buffer = self.buffer.as_ref().map(|(size, block)| {
+            (size, format!("[... {} bytes ...]", (**block).as_ref().len())) 
+        });
+        f.debug_struct("IntoSink")
+            .field("writer", &self.writer)
+            .field("buffer", &buffer)
+            .finish()?;
+        Ok(())
+    }
+}

futures-util/src/io/mod.rs

@@ -34,6 +34,9 @@ pub use self::copy_into::CopyInto;
 mod flush;
 pub use self::flush::Flush;
 
+mod into_sink;
+pub use self::into_sink::IntoSink;
+
 mod lines;
 pub use self::lines::Lines;
 
@@ -382,6 +385,40 @@ pub trait AsyncWriteExt: AsyncWrite {
     {
         Compat::new(self)
     }
+
+
+    /// Allow using an [`AsyncWrite`] as a [`Sink`](futures_sink::Sink)`<Item: AsRef<[u8]>>`.
+    ///
+    /// This adapter produces a sink that will write each value passed to it
+    /// into the underlying writer.
+    ///
+    /// Note that this function consumes the given writer, returning a wrapped
+    /// version.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await)]
+    /// # futures::executor::block_on(async {
+    /// use futures::io::AsyncWriteExt;
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter(vec![Ok([1, 2, 3]), Ok([4, 5, 6])]);
+    ///
+    /// let mut writer = vec![];
+    ///
+    /// stream.forward((&mut writer).into_sink()).await?;
+    ///
+    /// assert_eq!(writer, vec![1, 2, 3, 4, 5, 6]);
+    /// # Ok::<(), Box<dyn std::error::Error>>(())
+    /// # })?;
+    /// # Ok::<(), Box<dyn std::error::Error>>(())
+    /// ```
+    fn into_sink(self) -> IntoSink<Self>
+        where Self: Sized,
+    {
+        IntoSink::new(self)
+    }
 }
 
 impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {}