chunky-bits 0.2.0

Chunky Bits is simple, unmanaged, distributed HTTP object store tool.
use std::borrow::Borrow;

use futures::stream::{
    self,
    Stream,
    StreamExt,
};
use tokio::io::{
    self,
    AsyncRead,
};
use tokio_util::io::StreamReader;

use crate::{
    error::FileReadError,
    file::{
        FilePart,
        FileReference,
        LocationContext,
    },
};

pub struct FileReadBuilder<T> {
    file: T,
    buffer: usize,
    location_context: LocationContext,
}

impl<T> FileReadBuilder<T> {
    pub(super) fn new(file: T) -> FileReadBuilder<T> {
        FileReadBuilder {
            file,
            buffer: 5,
            location_context: Default::default(),
        }
    }

    pub fn location_context(self, location_context: LocationContext) -> FileReadBuilder<T> {
        let mut new = self;
        new.location_context = location_context;
        new
    }

    fn inner_stream_reader<'a>(
        &self,
        parts: impl Iterator<Item = impl Borrow<FilePart> + Send + Sync> + Send + 'a,
        length: u64,
    ) -> impl Stream<Item = Result<Vec<u8>, FileReadError>> + Send + 'a {
        let FileReadBuilder {
            buffer,
            location_context,
            ..
        } = self;
        let mut bytes_remaining: u64 = length;
        let location_context = location_context.clone();
        stream::iter(parts.map(move |part| {
            let location_context = location_context.clone();
            async move { part.borrow().read_with_context(&location_context).await }
        }))
        .buffered(*buffer)
        .map(move |res| match res {
            Ok(mut bytes) => {
                if bytes.len() as u64 > bytes_remaining {
                    bytes.drain((bytes_remaining as usize)..);
                }
                bytes_remaining -= bytes.len() as u64;
                Ok(bytes)
            },
            Err(err) => Err(err),
        })
    }

    fn inner_reader<'a>(
        stream: impl Stream<Item = Result<Vec<u8>, FileReadError>> + Send + 'a,
    ) -> impl AsyncRead + Send + 'a {
        StreamReader::new(stream.map(|res| -> io::Result<_> {
            match res {
                Ok(bytes) => Ok(std::io::Cursor::new(bytes)),
                Err(err) => Err(io::Error::new(io::ErrorKind::Other, err)),
            }
        }))
    }
}

impl<T> FileReadBuilder<T>
where
    T: Borrow<FileReference> + Send + Sync,
{
    pub fn stream_reader(&self) -> impl Stream<Item = Result<Vec<u8>, FileReadError>> + Send + '_ {
        let FileReference { parts, length, .. } = self.file.borrow();
        self.inner_stream_reader(parts.iter(), length.unwrap())
    }

    pub fn reader(&self) -> impl AsyncRead + Send + Unpin + '_ {
        Self::inner_reader(self.stream_reader())
    }
}

impl<T> FileReadBuilder<T>
where
    T: Into<FileReference> + Send + 'static,
{
    pub fn stream_reader_owned(
        self,
    ) -> impl Stream<Item = Result<Vec<u8>, FileReadError>> + Send + 'static {
        let FileReadBuilder {
            buffer,
            location_context,
            file,
        } = self;
        let new = FileReadBuilder {
            buffer,
            location_context,
            file: (),
        };
        let file = file.into();
        let length = file.length.unwrap();
        let parts = file.parts.into_iter();
        new.inner_stream_reader(parts, length)
    }

    pub fn reader_owned(self) -> impl AsyncRead + Send + Unpin + 'static {
        FileReadBuilder::<()>::inner_reader(self.stream_reader_owned())
    }
}