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).await {
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>, ctx: &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().poll_unpin(ctx) {
98                Poll::Ready(next) => match next {
99                    Ok(Some(item)) => return Poll::Ready(Some(item)),
100                    Ok(None) => {
101                        let mut wait = (&mut this.notify).next();
102                        match wait.poll_unpin(ctx) {
103                            Poll::Ready(_) => {
104                                if this.stop_flag.load(Ordering::SeqCst) {
105                                    break Poll::Ready(None);
106                                } else {
107                                    continue;
108                                }
109                            }
110                            Poll::Pending => return Poll::Pending,
111                        }
112                    }
113                    Err(err) => {
114                        log::error!("poll external buffer error: {}", err);
115                        return Poll::Ready(None);
116                    }
117                },
118                Poll::Pending => return Poll::Pending,
119            }
120        }
121    }
122}
123
124#[cfg(feature = "default")]
125pub fn create_external_buffered_stream<T, S, P>(
126    stream: S,
127    path: P,
128) -> Result<ExternalBufferedStream<T, ExternalBufferSled, S>, Error>
129where
130    T: ExternalBufferSerde + Send + 'static,
131    S: Stream<Item = T> + Send + Sync + 'static,
132    P: AsRef<std::path::Path>,
133{
134    Ok(ExternalBufferedStream::new(
135        stream,
136        ExternalBufferSled::new(path)?,
137    ))
138}
139
140#[cfg(feature = "queue")]
141pub fn create_queued_stream<T, S>(
142    stream: S,
143) -> Result<ExternalBufferedStream<T, ExternalBufferQueue<T>, S>, Error>
144where
145    T: Ord + Send + 'static,
146    S: Stream<Item = T> + Send + Sync + 'static,
147{
148    Ok(ExternalBufferedStream::new(
149        stream,
150        ExternalBufferQueue::new(),
151    ))
152}
external_buffered_stream/lib.rs

external_buffered_stream/
lib.rs
