objectstore_service/

stream.rs

//! Stream types and buffering utilities for object data.
//!
//! Data flows in streams to keep memory consumption low. Two distinct types
//! cover the two directions of data flow:
//!
//! - [`ClientStream`] — incoming data from a client PUT request body. Uses
//!   [`ClientError`] as the error type so backends can distinguish a broken
//!   client connection from a backend I/O failure (400 vs 500).
//! - [`PayloadStream`] — outgoing data returned from
//!   [`Backend::get_object`](crate::backend::common::Backend::get_object).

use std::error::Error;
use std::fmt;
use std::io;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use bytes::{Bytes, BytesMut};
use futures_util::stream::BoxStream;
use futures_util::{Stream, StreamExt, TryStreamExt};

/// Outgoing byte stream returned by [`Backend::get_object`](crate::backend::common::Backend::get_object).
///
/// Use [`single`] to construct a single-chunk `PayloadStream` from an owned value.
pub type PayloadStream = BoxStream<'static, io::Result<Bytes>>;

/// Error originating from a client-supplied input stream.
///
/// Wraps any error yielded by a [`ClientStream`] (the incoming HTTP request body).
/// Backends receive this via [`ClientStream`] and can detect it with
/// [`unpack_client_error`] to return a 4xx response rather than treating
/// it as a 5xx backend failure.
#[derive(Clone, Debug)]
pub struct ClientError(Arc<dyn Error + Send + Sync + 'static>);

impl ClientError {
    /// Creates a new [`ClientError`] wrapping `err`.
    pub fn new<E>(err: E) -> Self
    where
        E: Error + Send + Sync + 'static,
    {
        Self(Arc::new(err))
    }
}

impl fmt::Display for ClientError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Error for ClientError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        self.0.source()
    }
}

/// Required by [`tokio_util::io::StreamReader`] in the local filesystem backend.
impl From<ClientError> for io::Error {
    fn from(err: ClientError) -> Self {
        io::Error::other(err)
    }
}

/// Incoming byte stream from a client PUT request body.
///
/// Uses [`ClientError`] as the error type so that a dropped or interrupted
/// client connection is distinguishable from a backend I/O failure. Backends
/// that detect a [`ClientError`] (via [`unpack_client_error`]) can surface it
/// as [`ClientStream`](crate::error::ErrorKind::ClientStream), which the server
/// maps to HTTP 400 rather than 500.
///
/// Use [`single`] to construct a single-chunk `ClientStream` from an owned value.
pub type ClientStream = BoxStream<'static, Result<Bytes, ClientError>>;

/// Walks the source chain of `err` looking for a [`ClientError`].
///
/// At each step, two locations are checked:
///
/// - **Direct**: the error itself is a `ClientError`.
/// - **Packed in `io::Error`**: the error is an `io::Error` whose custom inner
///   value is a `ClientError`.
///
/// Use this in `put_object` implementations to reclassify body-stream errors
/// as [`ClientStream`](crate::error::ErrorKind::ClientStream) instead of an
/// opaque server error.
pub fn unpack_client_error<E>(err: &E) -> Option<ClientError>
where
    E: Error + 'static,
{
    let mut source = Some(err as &(dyn Error + 'static));

    while let Some(s) = source {
        // The client error may be wrapped as custom `io::Error`, in which case it cannot be
        // discovered by iterating `sources`.
        let target = match s.downcast_ref::<io::Error>().and_then(|e| e.get_ref()) {
            Some(inner) => inner,
            None => s,
        };

        if let Some(client_error) = target.downcast_ref::<ClientError>() {
            return Some(client_error.clone());
        }

        source = s.source();
    }
    None
}

#[derive(Debug, Default)]
enum ChunkedBytesState {
    #[default]
    Empty,
    Single(Bytes),
    Multi(BytesMut),
}

/// Lazy stream buffer that avoids copying single chunks.
///
/// Tracks three internal states:
/// - **Empty** — no allocation, no data.
/// - **Single** — first chunk stored by move (zero-copy).
/// - **Multi** — allocated on second `push`, coalesces all chunks.
#[derive(Debug)]
pub(crate) struct ChunkedBytes {
    state: ChunkedBytesState,
    capacity: usize,
}

impl ChunkedBytes {
    /// Creates a new buffer with the given capacity hint.
    ///
    /// The capacity is used when transitioning from Single to Multi state
    /// and is exposed via [`capacity()`](Self::capacity) for use as a buffer limit.
    pub fn new(capacity: usize) -> Self {
        Self {
            state: ChunkedBytesState::Empty,
            capacity,
        }
    }

    /// Appends a chunk to the buffer.
    ///
    /// Empty→Single stores by move (zero-copy). Single→Multi allocates and
    /// copies both chunks. Multi extends the existing allocation.
    pub fn push(&mut self, chunk: Bytes) {
        if chunk.is_empty() {
            return;
        }
        self.state = match std::mem::take(&mut self.state) {
            ChunkedBytesState::Empty => ChunkedBytesState::Single(chunk),
            ChunkedBytesState::Single(first) => {
                let capacity = self.capacity.max(first.len() + chunk.len());
                let mut buf = BytesMut::with_capacity(capacity);
                buf.extend_from_slice(&first);
                buf.extend_from_slice(&chunk);
                ChunkedBytesState::Multi(buf)
            }
            ChunkedBytesState::Multi(mut buf) => {
                buf.extend_from_slice(&chunk);
                ChunkedBytesState::Multi(buf)
            }
        };
    }

    /// Returns the configured capacity (used as the buffer limit).
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Returns the total number of buffered bytes.
    pub fn len(&self) -> usize {
        match &self.state {
            ChunkedBytesState::Empty => 0,
            ChunkedBytesState::Single(b) => b.len(),
            ChunkedBytesState::Multi(b) => b.len(),
        }
    }

    /// Returns `true` if no bytes have been buffered.
    #[cfg(test)]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Consumes the buffer and returns the data as a single `Bytes`.
    ///
    /// Single-chunk data is returned by move (zero-copy). Multi-chunk data
    /// is frozen from the internal `BytesMut`.
    pub fn into_bytes(self) -> Bytes {
        match self.state {
            ChunkedBytesState::Empty => Bytes::new(),
            ChunkedBytesState::Single(b) => b,
            ChunkedBytesState::Multi(b) => b.freeze(),
        }
    }
}

/// Reads up to `limit` bytes from a stream to support size-based routing decisions.
///
/// Constructed via the async [`SizedPeek::new`], which reads eagerly so that
/// [`is_exhausted()`](Self::is_exhausted) is always valid.
/// The full stream (buffered prefix plus any unconsumed remainder) is recovered
/// with [`into_stream()`](Self::into_stream).
///
/// The buffer never exceeds `limit` bytes: the chunk that would cause overflow
/// is held separately and re-emitted by `into_stream` without copying.
pub(crate) struct SizedPeek<S> {
    buffer: ChunkedBytes,
    /// The first chunk that exceeded the limit; `None` when the stream was exhausted.
    pending: Option<Bytes>,
    /// `None` when the stream was fully consumed within the limit.
    stream: Option<S>,
}

impl<S> SizedPeek<S> {
    /// Returns `true` if the stream was fully consumed within the peek limit.
    pub fn is_exhausted(&self) -> bool {
        self.pending.is_none()
    }

    /// Returns the number of bytes held in the buffer (at most `limit`).
    #[cfg(test)]
    pub fn len(&self) -> usize {
        self.buffer.len()
    }
}

impl<S, E> SizedPeek<S>
where
    S: Stream<Item = Result<Bytes, E>> + Unpin,
{
    /// Reads from `stream` into an internal buffer until `limit` bytes are accumulated
    /// or the stream ends.
    ///
    /// Uses strictly-greater-than comparison: a stream of exactly `limit` bytes is
    /// considered exhausted.
    pub async fn new(mut stream: S, limit: usize) -> Result<Self, E> {
        let mut buffer = ChunkedBytes::new(limit);

        while let Some(chunk) = stream.try_next().await? {
            if buffer.len() + chunk.len() > buffer.capacity() {
                return Ok(Self {
                    buffer,
                    pending: Some(chunk),
                    stream: Some(stream),
                });
            }
            buffer.push(chunk);
        }

        Ok(Self {
            buffer,
            pending: None,
            stream: None,
        })
    }

    /// Consumes self and returns all bytes as a single [`Bytes`].
    ///
    /// If the peek limit was exceeded, drains the remaining stream before
    /// returning. Always correct regardless of [`is_exhausted`](Self::is_exhausted).
    pub async fn into_bytes(mut self) -> Result<Bytes, E> {
        if let Some(pending) = self.pending.take() {
            self.buffer.push(pending);
        }

        if let Some(mut stream) = self.stream.take() {
            while let Some(chunk) = stream.try_next().await? {
                self.buffer.push(chunk);
            }
        }

        Ok(self.buffer.into_bytes())
    }
}

impl<S, E> SizedPeek<S>
where
    S: Stream<Item = Result<Bytes, E>>,
{
    /// Consumes self and returns a stream that yields the buffered prefix first,
    /// then any remaining data from the original stream.
    pub fn into_stream(self) -> impl Stream<Item = Result<Bytes, E>> {
        let leading = [self.buffer.into_bytes(), self.pending.unwrap_or_default()]
            .into_iter()
            .filter(|b| !b.is_empty())
            .map(Ok);

        let tail = futures_util::stream::iter(self.stream).flatten();
        futures_util::stream::iter(leading).chain(tail)
    }
}

/// Creates a single-chunk stream that yields `contents` as one item.
pub fn single<E: Send + 'static>(
    contents: impl Into<Bytes>,
) -> BoxStream<'static, Result<Bytes, E>> {
    futures_util::stream::once(std::future::ready(Ok(contents.into()))).boxed()
}

/// Wraps a stream to count the total bytes yielded by successful chunks.
///
/// Returns the shared counter and the wrapped stream. The counter is incremented
/// as the stream is consumed, so read it only after the stream is exhausted.
pub fn counting_stream<S, E>(stream: S) -> (Arc<AtomicU64>, impl Stream<Item = Result<Bytes, E>>)
where
    S: Stream<Item = Result<Bytes, E>>,
{
    let counter = Arc::new(AtomicU64::new(0));

    (
        counter.clone(),
        stream.inspect(move |res| {
            if let Ok(chunk) = res {
                counter.fetch_add(chunk.len() as u64, Ordering::Relaxed);
            }
        }),
    )
}

/// Collects a stream of `Bytes` chunks into a `Vec<u8>`.
#[cfg(test)]
pub(crate) async fn read_to_vec<S, E>(mut stream: S) -> crate::error::Result<Vec<u8>>
where
    S: Stream<Item = Result<Bytes, E>> + Unpin,
    E: Into<crate::error::Error>,
{
    let mut payload = Vec::new();
    while let Some(result) = stream.next().await {
        let chunk = result.map_err(Into::into)?;
        payload.extend(&chunk);
    }
    Ok(payload)
}

#[cfg(test)]
mod tests {
    use futures_util::stream;

    use super::*;

    // --- ChunkedBytes tests ---

    #[test]
    fn empty_into_bytes() {
        let buf = ChunkedBytes::new(1024);
        assert!(buf.is_empty());
        assert_eq!(buf.len(), 0);
        assert_eq!(buf.into_bytes().len(), 0);
    }

    #[test]
    fn single_chunk_zero_copy() {
        let original = Bytes::from_static(b"hello world");
        let ptr = original.as_ptr();

        let mut buf = ChunkedBytes::new(1024);
        buf.push(original);

        assert_eq!(buf.len(), 11);
        assert!(!buf.is_empty());

        let result = buf.into_bytes();
        assert_eq!(result.as_ref(), b"hello world");
        // Pointer equality proves zero-copy.
        assert_eq!(result.as_ptr(), ptr);
    }

    #[test]
    fn multi_chunk_coalesces() {
        let mut buf = ChunkedBytes::new(1024);
        buf.push(Bytes::from_static(b"hello "));
        buf.push(Bytes::from_static(b"world"));

        assert_eq!(buf.len(), 11);
        assert_eq!(buf.into_bytes().as_ref(), b"hello world");
    }

    #[test]
    fn push_empty_chunk_is_noop() {
        let original = Bytes::from_static(b"data");
        let ptr = original.as_ptr();

        let mut buf = ChunkedBytes::new(1024);
        buf.push(Bytes::new());
        assert!(buf.is_empty());

        buf.push(original);
        buf.push(Bytes::new());
        // Still in Single state — empty pushes don't trigger Multi.
        let result = buf.into_bytes();
        assert_eq!(result.as_ptr(), ptr);
    }

    #[test]
    fn capacity_is_preserved() {
        let buf = ChunkedBytes::new(42);
        assert_eq!(buf.capacity(), 42);
    }

    // --- SizedPeek tests ---

    #[tokio::test]
    async fn exhausted_when_stream_fits() {
        let s = stream::once(std::future::ready(Ok::<_, io::Error>(Bytes::from_static(
            b"small payload",
        ))));

        let peeked = SizedPeek::new(s, 1024).await.unwrap();

        assert!(peeked.is_exhausted());
        assert_eq!(peeked.len(), 13);
    }

    #[tokio::test]
    async fn remaining_when_stream_exceeds_limit() {
        let chunks: Vec<io::Result<Bytes>> = vec![
            Ok(Bytes::from(vec![0u8; 600])),
            Ok(Bytes::from(vec![1u8; 600])),
        ];

        let peeked = SizedPeek::new(stream::iter(chunks), 1000).await.unwrap();

        assert!(!peeked.is_exhausted());
        // Only the first 600-byte chunk fits in the buffer; the second is pending.
        assert_eq!(peeked.len(), 600);
    }

    #[tokio::test]
    async fn into_stream_reassembles_data() {
        // "aaa" fits (3 ≤ 5), "bbb" overflows (3+3=6 > 5) → pending, "ccc" stays in stream.
        let chunks: Vec<io::Result<Bytes>> = vec![
            Ok(Bytes::from_static(b"aaa")),
            Ok(Bytes::from_static(b"bbb")),
            Ok(Bytes::from_static(b"ccc")),
        ];

        let peeked = SizedPeek::new(stream::iter(chunks), 5).await.unwrap();
        assert!(!peeked.is_exhausted());

        let output = read_to_vec(peeked.into_stream()).await.unwrap();
        assert_eq!(output, b"aaabbbccc");
    }

    #[tokio::test]
    async fn into_stream_single_chunk_zero_copy() {
        let original = Bytes::from_static(b"zero-copy roundtrip");
        let ptr = original.as_ptr();
        let s = stream::once(std::future::ready(Ok::<_, io::Error>(original)));

        let peeked = SizedPeek::new(s, 1024).await.unwrap();
        let out = peeked.into_stream();
        futures_util::pin_mut!(out);
        let first = out.try_next().await.unwrap().unwrap();
        assert_eq!(first.as_ptr(), ptr);
    }

    #[tokio::test]
    async fn oversized_single_chunk_goes_to_pending() {
        // A single chunk larger than the limit never enters the buffer.
        let big = Bytes::from(vec![0u8; 2000]);
        let ptr = big.as_ptr();
        let s = stream::once(std::future::ready(Ok::<_, io::Error>(big)));

        let peeked = SizedPeek::new(s, 1000).await.unwrap();
        assert!(!peeked.is_exhausted());
        assert_eq!(peeked.len(), 0); // nothing in buffer

        let out = peeked.into_stream();
        futures_util::pin_mut!(out);
        let first = out.try_next().await.unwrap().unwrap();
        // The pending chunk is emitted by move — same pointer, no copy.
        assert_eq!(first.as_ptr(), ptr);
        assert_eq!(first.len(), 2000);
    }

    #[tokio::test]
    async fn error_propagation() {
        let chunks: Vec<io::Result<Bytes>> = vec![
            Ok(Bytes::from_static(b"ok")),
            Err(io::Error::new(io::ErrorKind::BrokenPipe, "broken")),
        ];

        let result = SizedPeek::new(stream::iter(chunks), 1024).await;
        assert!(result.is_err());
    }
}