spiral-contract 0.1.12

//! Owned byte streams shared by service implementations and HTTP adapters.

use std::fmt::{self, Debug, Formatter};
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::Bytes;
use futures::{Stream, StreamExt, stream};

use crate::call::BoxError;

/// A fallible byte stream independent of an HTTP body implementation.
///
/// Polling drives the underlying stream. Dropping this value drops that stream;
/// it does not spawn background work. Chunks preserve their original ordering
/// and boundaries, but callers must not assign protocol meaning to boundaries.
pub struct ByteStream {
    inner: Pin<Box<dyn Stream<Item = Result<Bytes, BoxError>> + Send + 'static>>,
}

impl ByteStream {
    /// Own a stream and preserve its failures for explicit diagnostics.
    pub fn new<S, E>(stream: S) -> Self
    where
        S: Stream<Item = Result<Bytes, E>> + Send + 'static,
        E: Into<BoxError>,
    {
        Self {
            inner: Box::pin(stream.map(|chunk| chunk.map_err(Into::into))),
        }
    }

    /// Construct a stream containing a single byte chunk.
    pub fn from_bytes(bytes: Bytes) -> Self {
        Self::new(stream::iter([Ok::<_, BoxError>(bytes)]))
    }

    /// Construct a stream that finishes without yielding any chunks.
    pub fn empty() -> Self {
        Self::new(stream::empty::<Result<Bytes, BoxError>>())
    }
}

impl Stream for ByteStream {
    type Item = Result<Bytes, BoxError>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.as_mut().poll_next(context)
    }
}

impl Debug for ByteStream {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter.write_str("ByteStream([REDACTED])")
    }
}

#[cfg(test)]
mod tests {
    use std::io;
    use std::pin::Pin;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::task::{Context, Poll};

    use bytes::Bytes;
    use futures::Stream;
    use futures::task::noop_waker;

    use super::ByteStream;

    struct PendingStream(Arc<AtomicBool>);

    impl Stream for PendingStream {
        type Item = Result<Bytes, io::Error>;

        fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    impl Drop for PendingStream {
        fn drop(&mut self) {
            self.0.store(true, Ordering::SeqCst);
        }
    }

    #[test]
    fn dropping_a_pending_stream_releases_its_source() {
        let released = Arc::new(AtomicBool::new(false));
        let mut stream = ByteStream::new(PendingStream(released.clone()));
        let waker = noop_waker();
        let mut context = Context::from_waker(&waker);
        assert!(Pin::new(&mut stream).poll_next(&mut context).is_pending());
        assert!(!released.load(Ordering::SeqCst));
        drop(stream);
        assert!(released.load(Ordering::SeqCst));
    }
}