external_buffered_stream/
lib.rs

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

external_buffered_stream/
lib.rs
