Streaming encoder

Cargo.toml

@@ -56,14 +56,15 @@ proptest = "1.0.0"
 proptest-derive = "0.3.0"
 rand = "0.8.5"
 futures = "0.3.5"
+futures-channel = "0.3.24"
 futures-test = "0.3.5"
 ntest = "0.8.1"
 bytes-05 = { package = "bytes", version = "0.5.0" }
 bytes-06 = { package = "bytes", version = "0.6.0" }
 bytes = "1.0.0"
 tokio-02 = { package = "tokio", version = "0.2.21", default-features = false, features = ["io-util", "stream", "macros", "io-std"] }
 tokio-03 = { package = "tokio", version = "0.3.0", default-features = false, features = ["io-util", "stream"] }
-tokio = { version = "1.0.0", default-features = false, features = ["io-util"] }
+tokio = { version = "1.0.0", default-features = false, features = ["io-util", "macros", "rt", "time"] }
 tokio-util-03 = { package = "tokio-util", version = "0.3.0", default-features = false, features = ["codec"] }
 tokio-util-04 = { package = "tokio-util", version = "0.4.0", default-features = false, features = ["io"] }
 tokio-util-05 = { package = "tokio-util", version = "0.5.0", default-features = false, features = ["io"] }

src/futures/flush.rs

@@ -0,0 +1,74 @@
+//! Types related to [`AsyncFlush`](AsyncFlush) to wrap encoders
+
+use futures_core::Future;
+use futures_core::Stream;
+use futures_io::AsyncRead;
+use pin_project_lite::pin_project;
+
+use super::bufread::Encoder;
+use crate::codec::Encode;
+
+/// Flushes asynchronously
+///
+/// `AsyncRead` and `AsyncBufRead` implementations may not have enough information
+/// to know when to flush the data they have in store, so they can implement this
+/// trait and let the caller decide when data should be flushed
+pub trait AsyncFlush {
+    /// Attempts to flush in flight data from the `AsyncFlush` into `buf`.
+    ///
+    /// On success, returns `Poll::Ready(Ok(()))` and places data in the
+    /// unfilled portion of `buf`. If no data was read (`buf.filled().len()` is
+    /// unchanged), it implies that EOF has been reached.
+    ///
+    /// If no data is available for reading, the method returns `Poll::Pending`
+    /// and arranges for the current task (via `cx.waker()`) to receive a
+    /// notification when the object becomes readable or is closed.
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut [u8],
+    ) -> std::task::Poll<std::io::Result<usize>>;
+}
+
+pin_project! {
+    /// This structure wraps an `Encoder` implementing [`AsyncRead`](tokio::io::AsyncRead) to
+    /// allow the caller to flush its buffers.
+    pub struct FlushableEncoder<E: AsyncRead, Rx: Stream<Item=()>> {
+        #[pin]
+        encoder: E,
+        #[pin]
+        receiver: Rx,
+    }
+}
+
+impl<E: AsyncRead + AsyncFlush, Rx: Stream<Item = ()>> FlushableEncoder<E, Rx> {
+    /// Creates a new `FlushableEncoder` from an existing `Encoder` and a Stream
+    ///
+    /// Whenever a message is received from the stream, the encoder will flushes its buffers
+    /// and compress them.
+    pub fn new(encoder: E, receiver: Rx) -> Self {
+        Self { encoder, receiver }
+    }
+}
+
+impl<E: AsyncRead + AsyncFlush, Rx: Stream<Item = ()>> AsyncRead for FlushableEncoder<E, Rx> {
+    fn poll_read(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut [u8],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        let mut this = self.project();
+
+        match this.encoder.as_mut().poll_read(cx, buf) {
+            std::task::Poll::Ready(r) => std::task::Poll::Ready(r),
+            std::task::Poll::Pending => match this.receiver.as_mut().poll_next(cx) {
+                std::task::Poll::Pending => std::task::Poll::Pending,
+                std::task::Poll::Ready(_) => match this.encoder.poll_flush(cx, buf) {
+                    std::task::Poll::Ready(Ok(sz)) => std::task::Poll::Ready(Ok(sz)),
+                    std::task::Poll::Ready(Err(e)) => std::task::Poll::Ready(Err(e)),
+                    std::task::Poll::Pending => std::task::Poll::Pending,
+                },
+            },
+        }
+    }
+}

src/futures/mod.rs

@@ -1,4 +1,5 @@
 //! Implementations for IO traits exported by `futures`.
 
 pub mod bufread;
+pub mod flush;
 pub mod write;

src/tokio/bufread/generic/encoder.rs

@@ -4,7 +4,7 @@ use core::{
 };
 use std::io::Result;
 
-use crate::{codec::Encode, util::PartialBuffer};
+use crate::{codec::Encode, tokio::flush::AsyncFlush, util::PartialBuffer};
 use futures_core::ready;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncBufRead, AsyncRead, ReadBuf};
@@ -115,3 +115,23 @@ impl<R: AsyncBufRead, E: Encode> AsyncRead for Encoder<R, E> {
         }
     }
 }
+
+impl<R: AsyncBufRead, E: Encode> AsyncFlush for Encoder<R, E> {
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<Result<bool>> {
+        let mut output = PartialBuffer::new(buf.initialize_unfilled());
+        let mut this = self.project();
+
+        match this.encoder.flush(&mut output)? {
+            true => {
+                let len = output.written().len();
+                buf.advance(len);
+                Poll::Ready(Ok(true))
+            }
+            false => Poll::Ready(Ok(false)),
+        }
+    }
+}

src/tokio/bufread/macros/encoder.rs

@@ -62,6 +62,16 @@ macro_rules! encoder {
             }
         }
 
+        impl<$inner: tokio::io::AsyncBufRead> crate::tokio::flush::AsyncFlush for $name<$inner> {
+            fn poll_flush(
+                self: std::pin::Pin<&mut Self>,
+                cx: &mut std::task::Context<'_>,
+                buf: &mut tokio::io::ReadBuf<'_>,
+            ) -> std::task::Poll<std::io::Result<bool>> {
+                self.project().inner.poll_flush(cx, buf)
+            }
+        }
+
         const _: () = {
             fn _assert() {
                 use crate::util::{_assert_send, _assert_sync};

src/tokio/flush.rs

@@ -0,0 +1,75 @@
+//! Types related to [`AsyncFlush`](AsyncFlush) to wrap encoders
+
+use futures_core::Future;
+use futures_core::Stream;
+use pin_project_lite::pin_project;
+use tokio::io::{AsyncBufRead, AsyncRead};
+
+use super::bufread::Encoder;
+use crate::codec::Encode;
+
+/// Flushes asynchronously
+///
+/// `AsyncRead` and `AsyncBufRead` implementations may not have enough information
+/// to know when to flush the data they have in store, so they can implement this
+/// trait and let the caller decide when data should be flushed
+pub trait AsyncFlush {
+    /// Attempts to flush in flight data from the `AsyncFlush` into `buf`.
+    ///
+    /// On success, returns `Poll::Ready(Ok(()))` and places data in the
+    /// unfilled portion of `buf`. If no data was read (`buf.filled().len()` is
+    /// unchanged), it implies that EOF has been reached.
+    ///
+    /// If no data is available for reading, the method returns `Poll::Pending`
+    /// and arranges for the current task (via `cx.waker()`) to receive a
+    /// notification when the object becomes readable or is closed.
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> std::task::Poll<std::io::Result<bool>>;
+}
+
+pin_project! {
+    /// This structure wraps an `Encoder` implementing [`AsyncRead`](tokio::io::AsyncRead) to
+    /// allow the caller to flush its buffers.
+    pub struct FlushableEncoder<E: AsyncRead, Rx: Stream<Item=()>> {
+        #[pin]
+        encoder: E,
+        #[pin]
+        receiver: Rx,
+    }
+}
+
+impl<E: AsyncRead + AsyncFlush, Rx: Stream<Item = ()>> FlushableEncoder<E, Rx> {
+    /// Creates a new `FlushableEncoder` from an existing `Encoder` and a Stream
+    ///
+    /// Whenever a message is received from the stream, the encoder will flushes its buffers
+    /// and compress them.
+    pub fn new(encoder: E, receiver: Rx) -> Self {
+        Self { encoder, receiver }
+    }
+}
+
+impl<E: AsyncRead + AsyncFlush, Rx: Stream<Item = ()>> AsyncRead for FlushableEncoder<E, Rx> {
+    fn poll_read(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        let mut this = self.project();
+
+        match this.encoder.as_mut().poll_read(cx, buf) {
+            std::task::Poll::Ready(r) => std::task::Poll::Ready(r),
+            std::task::Poll::Pending => match this.receiver.as_mut().poll_next(cx) {
+                std::task::Poll::Pending => std::task::Poll::Pending,
+                std::task::Poll::Ready(_) => match this.encoder.poll_flush(cx, buf) {
+                    std::task::Poll::Ready(Ok(true)) => std::task::Poll::Ready(Ok(())),
+                    std::task::Poll::Ready(Ok(false)) => std::task::Poll::Pending,
+                    std::task::Poll::Ready(Err(e)) => std::task::Poll::Ready(Err(e)),
+                    std::task::Poll::Pending => std::task::Poll::Pending,
+                },
+            },
+        }
+    }
+}

src/tokio/mod.rs

@@ -1,4 +1,5 @@
 //! Implementations for IO traits exported by [`tokio` v1.0](::tokio).
 
 pub mod bufread;
+pub mod flush;
 pub mod write;