diff --git a/tokio-util/src/io/copy_to_bytes.rs b/tokio-util/src/io/copy_to_bytes.rs new file mode 100644 index 000000000..9509e7119 --- /dev/null +++ b/tokio-util/src/io/copy_to_bytes.rs @@ -0,0 +1,68 @@ +use bytes::Bytes; +use futures_sink::Sink; +use pin_project_lite::pin_project; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// A helper that wraps a [`Sink`]`<`[`Bytes`]`>` and converts it into a + /// [`Sink`]`<&'a [u8]>` by copying each byte slice into an owned [`Bytes`]. + /// + /// See the documentation for [`SinkWriter`] for an example. + /// + /// [`Bytes`]: bytes::Bytes + /// [`SinkWriter`]: crate::io::SinkWriter + /// [`Sink`]: futures_sink::Sink + #[derive(Debug)] + pub struct CopyToBytes { + #[pin] + inner: S, + } +} + +impl CopyToBytes { + /// Creates a new [`CopyToBytes`]. + pub fn new(inner: S) -> Self { + Self { inner } + } + + /// Gets a reference to the underlying sink. + pub fn get_ref(&self) -> &S { + &self.inner + } + + /// Gets a mutable reference to the underlying sink. + pub fn get_mut(&mut self) -> &mut S { + &mut self.inner + } + + /// Consumes this [`CopyToBytes`], returning the underlying sink. + pub fn into_inner(self) -> S { + self.inner + } +} + +impl<'a, S> Sink<&'a [u8]> for CopyToBytes +where + S: Sink, +{ + type Error = S::Error; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_ready(cx) + } + + fn start_send(self: Pin<&mut Self>, item: &'a [u8]) -> Result<(), Self::Error> { + self.project() + .inner + .start_send(Bytes::copy_from_slice(item)) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_close(cx) + } +} diff --git a/tokio-util/src/io/mod.rs b/tokio-util/src/io/mod.rs index 317d93b36..6c40d7390 100644 --- a/tokio-util/src/io/mod.rs +++ b/tokio-util/src/io/mod.rs @@ -10,9 +10,11 @@ //! [`Body`]: https://docs.rs/hyper/0.13/hyper/struct.Body.html //! [`AsyncRead`]: tokio::io::AsyncRead +mod copy_to_bytes; mod inspect; mod read_buf; mod reader_stream; +mod sink_writer; mod stream_reader; cfg_io_util! { @@ -20,8 +22,10 @@ cfg_io_util! { pub use self::sync_bridge::SyncIoBridge; } +pub use self::copy_to_bytes::CopyToBytes; pub use self::inspect::{InspectReader, InspectWriter}; pub use self::read_buf::read_buf; pub use self::reader_stream::ReaderStream; +pub use self::sink_writer::SinkWriter; pub use self::stream_reader::StreamReader; pub use crate::util::{poll_read_buf, poll_write_buf}; diff --git a/tokio-util/src/io/sink_writer.rs b/tokio-util/src/io/sink_writer.rs new file mode 100644 index 000000000..5d1acc499 --- /dev/null +++ b/tokio-util/src/io/sink_writer.rs @@ -0,0 +1,124 @@ +use futures_sink::Sink; + +use pin_project_lite::pin_project; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::AsyncWrite; + +pin_project! { + /// Convert a [`Sink`] of byte chunks into an [`AsyncWrite`]. + /// + /// Whenever you write to this [`SinkWriter`], the supplied bytes are + /// forwarded to the inner [`Sink`]. When `shutdown` is called on this + /// [`SinkWriter`], the inner sink is closed. + /// + /// This adapter takes a `Sink<&[u8]>` and provides an [`AsyncWrite`] impl + /// for it. Because of the lifetime, this trait is relatively rarely + /// implemented. The main ways to get a `Sink<&[u8]>` that you can use with + /// this type are: + /// + /// * With the codec module by implementing the [`Encoder`]`<&[u8]>` trait. + /// * By wrapping a `Sink` in a [`CopyToBytes`]. + /// * Manually implementing `Sink<&[u8]>` directly. + /// + /// The opposite conversion of implementing `Sink<_>` for an [`AsyncWrite`] + /// is done using the [`codec`] module. + /// + /// # Example + /// + /// ``` + /// use bytes::Bytes; + /// use futures_util::SinkExt; + /// use std::io::{Error, ErrorKind}; + /// use tokio::io::AsyncWriteExt; + /// use tokio_util::io::{SinkWriter, CopyToBytes}; + /// use tokio_util::sync::PollSender; + /// + /// # #[tokio::main(flavor = "current_thread")] + /// # async fn main() -> Result<(), Error> { + /// // We use an mpsc channel as an example of a `Sink`. + /// let (tx, mut rx) = tokio::sync::mpsc::channel::(1); + /// let sink = PollSender::new(tx).sink_map_err(|_| Error::from(ErrorKind::BrokenPipe)); + /// + /// // Wrap it in `CopyToBytes` to get a `Sink<&[u8]>`. + /// let mut writer = SinkWriter::new(CopyToBytes::new(sink)); + /// + /// // Write data to our interface... + /// let data: [u8; 4] = [1, 2, 3, 4]; + /// let _ = writer.write(&data).await?; + /// + /// // ... and receive it. + /// assert_eq!(data.as_slice(), &*rx.recv().await.unwrap()); + /// # Ok(()) + /// # } + /// ``` + /// + /// [`AsyncWrite`]: tokio::io::AsyncWrite + /// [`CopyToBytes`]: crate::io::CopyToBytes + /// [`Encoder`]: crate::codec::Encoder + /// [`Sink`]: futures_sink::Sink + /// [`codec`]: tokio_util::codec + #[derive(Debug)] + pub struct SinkWriter { + #[pin] + inner: S, + } +} + +impl SinkWriter { + /// Creates a new [`SinkWriter`]. + pub fn new(sink: S) -> Self { + Self { inner: sink } + } + + /// Gets a reference to the underlying sink. + pub fn get_ref(&self) -> &S { + &self.inner + } + + /// Gets a mutable reference to the underlying sink. + pub fn get_mut(&mut self) -> &mut S { + &mut self.inner + } + + /// Consumes this [`SinkWriter`], returning the underlying sink. + pub fn into_inner(self) -> S { + self.inner + } +} +impl AsyncWrite for SinkWriter +where + for<'a> S: Sink<&'a [u8], Error = E>, + E: Into, +{ + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let mut this = self.project(); + match this.inner.as_mut().poll_ready(cx) { + Poll::Ready(Ok(())) => { + if let Err(e) = this.inner.as_mut().start_send(buf) { + Poll::Ready(Err(e.into())) + } else { + Poll::Ready(Ok(buf.len())) + } + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), + Poll::Pending => { + cx.waker().wake_by_ref(); + Poll::Pending + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_flush(cx).map_err(Into::into) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_close(cx).map_err(Into::into) + } +} diff --git a/tokio-util/tests/io_sink_writer.rs b/tokio-util/tests/io_sink_writer.rs new file mode 100644 index 000000000..e76870be4 --- /dev/null +++ b/tokio-util/tests/io_sink_writer.rs @@ -0,0 +1,72 @@ +#![warn(rust_2018_idioms)] + +use bytes::Bytes; +use futures_util::SinkExt; +use std::io::{self, Error, ErrorKind}; +use tokio::io::AsyncWriteExt; +use tokio_util::codec::{Encoder, FramedWrite}; +use tokio_util::io::{CopyToBytes, SinkWriter}; +use tokio_util::sync::PollSender; + +#[tokio::test] +async fn test_copied_sink_writer() -> Result<(), Error> { + // Construct a channel pair to send data across and wrap a pollable sink. + // Note that the sink must mimic a writable object, e.g. have `std::io::Error` + // as its error type. + // As `PollSender` requires an owned copy of the buffer, we wrap it additionally + // with a `CopyToBytes` helper. + let (tx, mut rx) = tokio::sync::mpsc::channel::(1); + let mut writer = SinkWriter::new(CopyToBytes::new( + PollSender::new(tx).sink_map_err(|_| io::Error::from(ErrorKind::BrokenPipe)), + )); + + // Write data to our interface... + let data: [u8; 4] = [1, 2, 3, 4]; + let _ = writer.write(&data).await; + + // ... and receive it. + assert_eq!(data.to_vec(), rx.recv().await.unwrap().to_vec()); + + Ok(()) +} + +/// A trivial encoder. +struct SliceEncoder; + +impl SliceEncoder { + fn new() -> Self { + Self {} + } +} + +impl<'a> Encoder<&'a [u8]> for SliceEncoder { + type Error = Error; + + fn encode(&mut self, item: &'a [u8], dst: &mut bytes::BytesMut) -> Result<(), Self::Error> { + // This is where we'd write packet headers, lengths, etc. in a real encoder. + // For simplicity and demonstration purposes, we just pack a copy of + // the slice at the end of a buffer. + dst.extend_from_slice(item); + Ok(()) + } +} + +#[tokio::test] +async fn test_direct_sink_writer() -> Result<(), Error> { + // We define a framed writer which accepts byte slices + // and 'reverse' this construction immediately. + let framed_byte_lc = FramedWrite::new(Vec::new(), SliceEncoder::new()); + let mut writer = SinkWriter::new(framed_byte_lc); + + // Write multiple slices to the sink... + let _ = writer.write(&[1, 2, 3]).await; + let _ = writer.write(&[4, 5, 6]).await; + + // ... and compare it with the buffer. + assert_eq!( + writer.into_inner().write_buffer().to_vec().as_slice(), + &[1, 2, 3, 4, 5, 6] + ); + + Ok(()) +}