external_buffered_stream/
lib.rs

1mod buffer;
2mod error;
3mod runtime;
4mod serde;
5
6pub use buffer::*;
7pub use error::*;
8pub use serde::*;
9
10use std::{
11    marker::PhantomData,
12    pin::Pin,
13    sync::{
14        atomic::{AtomicBool, Ordering},
15        Arc,
16    },
17    task::{Context, Poll},
18};
19
20use futures::{channel::mpsc, Future, SinkExt, Stream, StreamExt};
21
22pub struct ExternalBufferedStream<T, B, S>
23where
24    T: Send,
25    B: ExternalBuffer<T>,
26    S: Stream<Item = T>,
27{
28    buffer: Arc<B>,
29    _source: PhantomData<S>,
30    notify: mpsc::UnboundedReceiver<()>,
31    stop_flag: Arc<AtomicBool>,
32
33    // the pending future that be polled by the stream consumer
34    pending: Option<Pin<Box<dyn Future<Output = Result<Option<T>, Error>> + Send>>>,
35}
36
37impl<T, B, S> ExternalBufferedStream<T, B, S>
38where
39    T: Send,
40    B: ExternalBuffer<T> + 'static,
41    S: Stream<Item = T> + Send + 'static,
42{
43    pub fn new(source: S, buffer: B) -> Self {
44        let source = Box::pin(source);
45
46        let buffer = Arc::new(buffer);
47        let buffer_clone = buffer.clone();
48
49        let (notify_tx, notify_rx) = mpsc::unbounded::<()>();
50
51        let stop_flag = Arc::new(AtomicBool::new(false));
52        let stop_flag_clone = stop_flag.clone();
53
54        let handle_source = async move {
55            let mut source = source;
56            let mut notify_tx = notify_tx;
57            while let Some(item) = source.next().await {
58                match buffer_clone.push(item).await {
59                    Ok(()) => match notify_tx.send(()).await {
60                        Ok(_) => {}
61                        Err(e) => {
62                            log::error!("Failed to notify: {:?}", e);
63                            break;
64                        }
65                    },
66                    Err(e) => {
67                        log::error!("Failed to push item to buffer: {:?}", e);
68                        break;
69                    }
70                }
71            }
72            log::info!("Source of external buffer stream is ended.");
73            stop_flag_clone.store(true, Ordering::SeqCst);
74            _ = notify_tx.send(())
75        };
76        runtime::spawn(handle_source);
77
78        ExternalBufferedStream {
79            buffer,
80            _source: PhantomData,
81            notify: notify_rx,
82            stop_flag,
83            pending: None,
84        }
85    }
86}
87
88impl<T, B, S> Stream for ExternalBufferedStream<T, B, S>
89where
90    T: Send,
91    B: ExternalBuffer<T> + 'static,
92    S: Stream<Item = T> + Send + 'static,
93{
94    type Item = T;
95
96    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
97        // S is PhantomData, so here is safe to get mut
98        let this = unsafe { self.get_unchecked_mut() };
99
100        loop {
101            if this.stop_flag.load(Ordering::SeqCst) {
102                return Poll::Ready(None);
103            }
104
105            if let Some(pending) = this.pending.as_mut() {
106                match pending.as_mut().poll(cx) {
107                    Poll::Ready(result) => {
108                        this.pending = None;
109
110                        match result {
111                            Ok(Some(item)) => {
112                                return Poll::Ready(Some(item));
113                            }
114                            Ok(None) => {
115                                // fall through to wait notify
116                            }
117                            Err(err) => {
118                                log::error!("external buffer shift return error: {}", err);
119                                return Poll::Ready(None);
120                            }
121                        }
122                    }
123                    Poll::Pending => {
124                        return Poll::Pending;
125                    }
126                }
127            }
128
129            match (&mut this.notify).poll_next_unpin(cx) {
130                Poll::Ready(Some(_)) => {
131                    let buffer = this.buffer.clone();
132                    this.pending = Some(Box::pin(async move { buffer.shift().await }));
133                    continue;
134                }
135                Poll::Ready(None) => return Poll::Ready(None),
136                Poll::Pending => return Poll::Pending,
137            }
138        }
139    }
140}
141
142#[cfg(feature = "default")]
143pub fn create_external_buffered_stream<T, S, P>(
144    stream: S,
145    path: P,
146) -> Result<ExternalBufferedStream<T, ExternalBufferSled, S>, Error>
147where
148    T: ExternalBufferSerde + Send + 'static,
149    S: Stream<Item = T> + Send + Sync + 'static,
150    P: AsRef<std::path::Path>,
151{
152    Ok(ExternalBufferedStream::new(
153        stream,
154        ExternalBufferSled::new(path)?,
155    ))
156}
157
158#[cfg(feature = "queue")]
159pub fn create_queued_stream<T, S>(
160    stream: S,
161) -> Result<ExternalBufferedStream<T, ExternalBufferQueue<T>, S>, Error>
162where
163    T: Ord + Send + 'static,
164    S: Stream<Item = T> + Send + Sync + 'static,
165{
166    Ok(ExternalBufferedStream::new(
167        stream,
168        ExternalBufferQueue::new(),
169    ))
170}
external_buffered_stream/lib.rs

external_buffered_stream/
lib.rs
