use core::future::Future;
use core::ops::ControlFlow;
use core::pin::Pin;
use core::task::{Context, Poll};
use bytes::Bytes;
use crate::pipe::drain_sink::DrainSink;
use crate::pipe::fan_in::Exhausted;
use crate::pipe::primitives::UnpinPipe;
pub struct AsStream<S>(S);
impl<S> AsStream<S> {
#[must_use]
pub fn into_inner(self) -> S {
self.0
}
}
impl<S: UnpinPipe<In = (), Err = Exhausted>> futures::Stream for AsStream<S> {
type Item = S::Out;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut call = self.0.call(());
match Pin::new(&mut call).poll(cx) {
Poll::Ready(Ok(item)) => Poll::Ready(Some(item)),
Poll::Ready(Err(Exhausted)) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}
pub trait PollSourceExt: UnpinPipe<In = (), Err = Exhausted> + Sized {
fn into_stream(self) -> AsStream<Self> {
AsStream(self)
}
fn as_stream(&mut self) -> AsStream<&mut Self> {
AsStream(self)
}
}
impl<S: UnpinPipe<In = (), Err = Exhausted>> PollSourceExt for S {}
impl<S: UnpinPipe<In = (), Err = Exhausted> + ?Sized> UnpinPipe for &mut S {
type In = ();
type Out = S::Out;
type Err = Exhausted;
fn call(&self, (): ()) -> impl Future<Output = Result<S::Out, Exhausted>> + Unpin {
(**self).call(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum AsSinkError {
#[error("sink rejected the item: no capacity")]
Rejected,
}
#[cfg(feature = "io-bridge")]
impl From<AsSinkError> for std::io::Error {
fn from(error: AsSinkError) -> Self {
std::io::Error::other(error)
}
}
pub struct AsSink<S>(S);
impl<S> AsSink<S> {
#[must_use]
pub fn into_inner(self) -> S {
self.0
}
}
impl<S: DrainSink<Item = [u8]> + Unpin> futures::Sink<Bytes> for AsSink<S> {
type Error = AsSinkError;
fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if self.get_mut().0.has_capacity() {
Poll::Ready(Ok(()))
} else {
Poll::Pending
}
}
fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
match self.get_mut().0.accept(&item) {
ControlFlow::Continue(()) => Ok(()),
ControlFlow::Break(()) => Err(AsSinkError::Rejected),
}
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
pub trait DrainSinkExt: DrainSink<Item = [u8]> + Sized {
fn into_sink(self) -> AsSink<Self> {
AsSink(self)
}
fn as_sink(&mut self) -> AsSink<&mut Self> {
AsSink(self)
}
}
impl<S: DrainSink<Item = [u8]>> DrainSinkExt for S {}
impl<S: DrainSink + ?Sized> DrainSink for &mut S {
type Item = S::Item;
fn accept(&mut self, item: &Self::Item) -> ControlFlow<()> {
(**self).accept(item)
}
fn has_capacity(&self) -> bool {
(**self).has_capacity()
}
}
#[cfg(feature = "io-bridge")]
mod reader_writer {
use bytes::Bytes;
use tokio_util::compat::{Compat, TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _};
use tokio_util::io::{CopyToBytes, SinkWriter, StreamReader};
use super::{AsSink, AsStream, DrainSinkExt as _, Exhausted, PollSourceExt as _};
use crate::pipe::drain_sink::DrainSink;
use crate::pipe::primitives::UnpinPipe;
pub trait IntoReader<Error>:
UnpinPipe<In = (), Out = Result<Bytes, Error>, Err = Exhausted> + Unpin + Sized
where
Error: Into<std::io::Error>,
{
fn into_reader(self) -> Compat<StreamReader<AsStream<Self>, Bytes>> {
StreamReader::new(self.into_stream()).compat()
}
fn as_reader(&mut self) -> Compat<StreamReader<AsStream<&mut Self>, Bytes>> {
StreamReader::new(self.as_stream()).compat()
}
}
impl<S, Error> IntoReader<Error> for S
where
S: UnpinPipe<In = (), Out = Result<Bytes, Error>, Err = Exhausted> + Unpin,
Error: Into<std::io::Error>,
{
}
pub trait IntoWriter: DrainSink<Item = [u8]> + Unpin + Sized {
fn into_writer(self) -> Compat<SinkWriter<CopyToBytes<AsSink<Self>>>> {
SinkWriter::new(CopyToBytes::new(self.into_sink())).compat_write()
}
fn as_writer(&mut self) -> Compat<SinkWriter<CopyToBytes<AsSink<&mut Self>>>> {
SinkWriter::new(CopyToBytes::new(self.as_sink())).compat_write()
}
}
impl<S: DrainSink<Item = [u8]> + Unpin> IntoWriter for S {}
}
#[cfg(feature = "io-bridge")]
pub use reader_writer::{IntoReader, IntoWriter};
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use core::cell::Cell;
use core::task::Waker;
use futures::{Sink, Stream};
use super::*;
use crate::pipe::drain_sink::RingSink;
const REQUEST_LINE: &[u8] = b"GET /orders?id=42 HTTP/1.1\r\n";
const HOST_HEADER: &[u8] = b"Host: api.example.internal\r\n";
const TAIL: &[u8] = b"Accept: application/json\r\n\r\n";
struct ThreeChunkSource {
chunks: [&'static [u8]; 3],
next: Cell<usize>,
}
impl UnpinPipe for ThreeChunkSource {
type In = ();
type Out = Result<Bytes, std::io::Error>;
type Err = Exhausted;
fn call(&self, (): ()) -> impl Future<Output = Result<Self::Out, Exhausted>> + Unpin {
let index = self.next.get();
if index >= self.chunks.len() {
return core::future::ready(Err(Exhausted));
}
let chunk = self.chunks[index];
self.next.set(index + 1);
core::future::ready(Ok(Ok(Bytes::from_static(chunk))))
}
}
fn noop_context() -> Context<'static> {
Context::from_waker(Waker::noop())
}
#[test]
fn as_stream_forwards_poll_next_in_order() {
let source = ThreeChunkSource {
chunks: [REQUEST_LINE, HOST_HEADER, TAIL],
next: Cell::new(0),
};
let mut stream = source.into_stream();
let mut cx = noop_context();
let first = Pin::new(&mut stream).poll_next(&mut cx);
let Poll::Ready(Some(Ok(chunk))) = first else {
panic!("expected the first chunk, got {first:?}");
};
assert_eq!(&chunk[..], REQUEST_LINE);
let second = Pin::new(&mut stream).poll_next(&mut cx);
let Poll::Ready(Some(Ok(chunk))) = second else {
panic!("expected the second chunk, got {second:?}");
};
assert_eq!(&chunk[..], HOST_HEADER);
}
#[test]
fn as_stream_drains_to_none_when_source_is_exhausted() {
let source = ThreeChunkSource {
chunks: [REQUEST_LINE, HOST_HEADER, TAIL],
next: Cell::new(3),
};
let mut stream = source.into_stream();
let mut cx = noop_context();
let Poll::Ready(None) = Pin::new(&mut stream).poll_next(&mut cx) else {
panic!("expected the source to report exhausted");
};
}
#[test]
fn as_sink_round_trips_real_bytes_through_a_ring() {
let mut sink: AsSink<RingSink<4, 64>> = RingSink::new().into_sink();
let mut cx = noop_context();
assert!(matches!(
Pin::new(&mut sink).poll_ready(&mut cx),
Poll::Ready(Ok(()))
));
Pin::new(&mut sink)
.start_send(Bytes::from_static(REQUEST_LINE))
.expect("ring has room for one frame");
Pin::new(&mut sink)
.start_send(Bytes::from_static(HOST_HEADER))
.expect("ring has room for a second frame");
let mut ring = sink.into_inner();
assert_eq!(ring.pop(), Some(REQUEST_LINE));
assert_eq!(ring.pop(), Some(HOST_HEADER));
}
#[test]
fn as_sink_reports_rejected_when_frame_too_large_for_the_ring() {
let mut sink: AsSink<RingSink<1, 4>> = RingSink::new().into_sink();
let error = Pin::new(&mut sink)
.start_send(Bytes::from_static(b"toolong"))
.expect_err("frame exceeds the 4-byte slot capacity");
assert_eq!(error, AsSinkError::Rejected);
}
#[test]
fn as_stream_as_stream_borrow_leaves_the_source_usable_after() {
let mut source = ThreeChunkSource {
chunks: [REQUEST_LINE, HOST_HEADER, TAIL],
next: Cell::new(0),
};
let mut cx = noop_context();
{
let mut borrowed = source.as_stream();
let first = Pin::new(&mut borrowed).poll_next(&mut cx);
assert!(first.is_ready());
}
assert_eq!(source.next.get(), 1);
}
}
#[cfg(all(test, feature = "io-bridge"))]
mod io_bridge_tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use bytes::Bytes;
use core::cell::Cell;
use core::future::Future;
use futures::io::{AsyncReadExt, AsyncWriteExt};
use super::{Exhausted, IntoReader as _, IntoWriter as _};
use crate::pipe::drain_sink::RingSink;
use crate::pipe::primitives::UnpinPipe;
const REQUEST_LINE: &[u8] = b"GET /orders?id=42 HTTP/1.1\r\n";
const HOST_HEADER: &[u8] = b"Host: api.example.internal\r\n";
const TAIL: &[u8] = b"Accept: application/json\r\n\r\n";
struct ChunkedRequestSource {
chunks: [&'static [u8]; 3],
next: Cell<usize>,
}
impl UnpinPipe for ChunkedRequestSource {
type In = ();
type Out = Result<Bytes, std::io::Error>;
type Err = Exhausted;
fn call(&self, (): ()) -> impl Future<Output = Result<Self::Out, Exhausted>> + Unpin {
let index = self.next.get();
if index >= self.chunks.len() {
return core::future::ready(Err(Exhausted));
}
let chunk = self.chunks[index];
self.next.set(index + 1);
core::future::ready(Ok(Ok(Bytes::from_static(chunk))))
}
}
#[test]
fn into_reader_round_trips_a_captured_http_request() {
let source = ChunkedRequestSource {
chunks: [REQUEST_LINE, HOST_HEADER, TAIL],
next: Cell::new(0),
};
let mut reader = source.into_reader();
let mut collected = Vec::new();
futures::executor::block_on(async {
reader
.read_to_end(&mut collected)
.await
.expect("an in-memory UnpinPipe stream never errors")
});
let mut expected = Vec::new();
expected.extend_from_slice(REQUEST_LINE);
expected.extend_from_slice(HOST_HEADER);
expected.extend_from_slice(TAIL);
assert_eq!(collected, expected, "reader output equals the fed bytes");
}
#[test]
fn into_writer_round_trips_bytes_into_a_ring_sink() {
let mut writer = RingSink::<4, 64>::new().into_writer();
futures::executor::block_on(async {
writer
.write_all(REQUEST_LINE)
.await
.expect("ring has room for the request line");
writer
.write_all(HOST_HEADER)
.await
.expect("ring has room for the host header");
writer
.close()
.await
.expect("close is a no-op success for a RingSink");
});
let mut ring = writer.into_inner().into_inner().into_inner().into_inner();
assert_eq!(ring.pop(), Some(REQUEST_LINE), "first frame written");
assert_eq!(ring.pop(), Some(HOST_HEADER), "second frame written");
}
}