external-buffered-stream 0.1.7

mod buffer;
mod error;
mod runtime;
mod serde;

pub use buffer::*;
pub use error::*;
pub use serde::*;

use std::{
    marker::PhantomData,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use futures::{channel::mpsc, Future, SinkExt, Stream, StreamExt};

pub struct ExternalBufferedStream<T, B, S>
where
    T: Send,
    B: ExternalBuffer<T>,
    S: Stream<Item = T>,
{
    buffer: Arc<B>,
    batch_size: usize,
    _source: PhantomData<S>,
    notify: mpsc::UnboundedReceiver<()>,

    // the pending future that be polled by the stream consumer
    pending: Option<Pin<Box<dyn Future<Output = Result<Option<Vec<T>>, Error>> + Send>>>,
}

impl<T, B, S> ExternalBufferedStream<T, B, S>
where
    T: Send,
    B: ExternalBuffer<T> + 'static,
    S: Stream<Item = T> + Send + 'static,
{
    // legacy api
    pub fn new(source: S, buffer: B) -> impl Stream<Item = T> {
        let batched_stream = Self::new_batch(source, buffer, 1);
        batched_stream.map(|batch| {
            assert!(batch.len() == 1, "the size of returned batch must be 1");
            batch.into_iter().next().unwrap()
        })
    }

    // new batch process api
    pub fn new_batch(source: S, buffer: B, max_batch_size: i32) -> Self {
        assert!(max_batch_size > 0);

        let source = Box::pin(source);

        let buffer = Arc::new(buffer);
        let buffer_clone = buffer.clone();

        let (notify_tx, notify_rx) = mpsc::unbounded::<()>();

        let handle_source = async move {
            let mut source = source;
            let mut notify_tx = notify_tx;
            while let Some(item) = source.next().await {
                match buffer_clone.push(item).await {
                    Ok(()) => match notify_tx.send(()).await {
                        Ok(_) => {}
                        Err(e) => {
                            log::error!("Failed to notify: {:?}", e);
                            break;
                        }
                    },
                    Err(e) => {
                        log::error!("Failed to push item to buffer: {:?}", e);
                        break;
                    }
                }
            }
            log::info!("Source of external buffer stream is ended.");
        };
        runtime::spawn(handle_source);

        ExternalBufferedStream {
            buffer,
            batch_size: max_batch_size as usize,
            _source: PhantomData,
            notify: notify_rx,
            pending: None,
        }
    }
}

impl<T, B, S> Stream for ExternalBufferedStream<T, B, S>
where
    T: Send,
    B: ExternalBuffer<T> + 'static,
    S: Stream<Item = T> + Send + 'static,
{
    type Item = Vec<T>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // S is PhantomData, so here is safe to get mut
        let this = unsafe { self.get_unchecked_mut() };

        loop {
            if this.pending.is_none() {
                let buffer = this.buffer.clone();
                let batch_size = this.batch_size;
                this.pending = Some(Box::pin(async move {
                    let mut batch = Vec::new();
                    for _ in 0..batch_size {
                        match buffer.shift().await {
                            Ok(Some(item)) => batch.push(item),
                            Ok(None) => break,
                            Err(e) => return Err(e),
                        }
                    }
                    if batch.is_empty() {
                        Ok(None)
                    } else {
                        Ok(Some(batch))
                    }
                }));
            }

            if let Some(pending) = this.pending.as_mut() {
                match pending.as_mut().poll(cx) {
                    Poll::Ready(result) => {
                        this.pending = None;

                        match result {
                            Ok(Some(batch)) => {
                                return Poll::Ready(Some(batch));
                            }
                            Ok(None) => {
                                let mut has_new = false;
                                let is_end = loop {
                                    // wait notify and consume all
                                    match (&mut this.notify).poll_next_unpin(cx) {
                                        Poll::Ready(Some(_)) => {
                                            has_new = true;
                                            // 消费所有通知
                                            continue;
                                        }
                                        Poll::Ready(None) => break true,
                                        Poll::Pending => break false,
                                    }
                                };
                                if has_new {
                                    continue;
                                } else if is_end {
                                    return Poll::Ready(None);
                                } else {
                                    return Poll::Pending;
                                }
                            }
                            Err(err) => {
                                log::error!("external buffer shift return error: {}", err);
                                return Poll::Ready(None);
                            }
                        }
                    }
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                }
            }
        }
    }
}

#[cfg(feature = "default")]
pub fn create_external_buffered_stream<T, S, P>(
    stream: S,
    path: P,
) -> Result<impl Stream<Item = T>, Error>
where
    T: ExternalBufferSerde + Send + 'static,
    S: Stream<Item = T> + Send + Sync + 'static,
    P: AsRef<std::path::Path>,
{
    Ok(ExternalBufferedStream::new(
        stream,
        ExternalBufferSled::new(path)?,
    ))
}

#[cfg(feature = "queue")]
pub fn create_queued_stream<T, S>(stream: S) -> Result<impl Stream<Item = T>, Error>
where
    T: Ord + Send + 'static,
    S: Stream<Item = T> + Send + Sync + 'static,
{
    Ok(ExternalBufferedStream::new(
        stream,
        ExternalBufferQueue::new(),
    ))
}