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::Arc,
14    task::{Context, Poll},
15};
16
17use futures::{channel::mpsc, Future, SinkExt, Stream, StreamExt};
18
19pub struct ExternalBufferedStream<T, B, S>
20where
21    T: Send,
22    B: ExternalBuffer<T>,
23    S: Stream<Item = T>,
24{
25    buffer: Arc<B>,
26    batch_size: usize,
27    _source: PhantomData<S>,
28    notify: mpsc::UnboundedReceiver<()>,
29
30    // the pending future that be polled by the stream consumer
31    pending: Option<Pin<Box<dyn Future<Output = Result<Option<Vec<T>>, Error>> + Send>>>,
32}
33
34impl<T, B, S> ExternalBufferedStream<T, B, S>
35where
36    T: Send,
37    B: ExternalBuffer<T> + 'static,
38    S: Stream<Item = T> + Send + 'static,
39{
40    // legacy api
41    pub fn new(source: S, buffer: B) -> impl Stream<Item = T> {
42        let batched_stream = Self::new_batch(source, buffer, 1);
43        batched_stream.map(|batch| {
44            assert!(batch.len() == 1, "the size of returned batch must be 1");
45            batch.into_iter().next().unwrap()
46        })
47    }
48
49    // new batch process api
50    pub fn new_batch(source: S, buffer: B, max_batch_size: i32) -> Self {
51        assert!(max_batch_size > 0);
52
53        let source = Box::pin(source);
54
55        let buffer = Arc::new(buffer);
56        let buffer_clone = buffer.clone();
57
58        let (notify_tx, notify_rx) = mpsc::unbounded::<()>();
59
60        let handle_source = async move {
61            let mut source = source;
62            let mut notify_tx = notify_tx;
63            while let Some(item) = source.next().await {
64                match buffer_clone.push(item).await {
65                    Ok(()) => match notify_tx.send(()).await {
66                        Ok(_) => {}
67                        Err(e) => {
68                            log::error!("Failed to notify: {:?}", e);
69                            break;
70                        }
71                    },
72                    Err(e) => {
73                        log::error!("Failed to push item to buffer: {:?}", e);
74                        break;
75                    }
76                }
77            }
78            log::info!("Source of external buffer stream is ended.");
79        };
80        runtime::spawn(handle_source);
81
82        ExternalBufferedStream {
83            buffer,
84            batch_size: max_batch_size as usize,
85            _source: PhantomData,
86            notify: notify_rx,
87            pending: None,
88        }
89    }
90}
91
92impl<T, B, S> Stream for ExternalBufferedStream<T, B, S>
93where
94    T: Send,
95    B: ExternalBuffer<T> + 'static,
96    S: Stream<Item = T> + Send + 'static,
97{
98    type Item = Vec<T>;
99
100    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
101        // S is PhantomData, so here is safe to get mut
102        let this = unsafe { self.get_unchecked_mut() };
103
104        loop {
105            if this.pending.is_none() {
106                let buffer = this.buffer.clone();
107                let batch_size = this.batch_size;
108                this.pending = Some(Box::pin(async move {
109                    let mut batch = Vec::new();
110                    for _ in 0..batch_size {
111                        match buffer.shift().await {
112                            Ok(Some(item)) => batch.push(item),
113                            Ok(None) => break,
114                            Err(e) => return Err(e),
115                        }
116                    }
117                    if batch.is_empty() {
118                        Ok(None)
119                    } else {
120                        Ok(Some(batch))
121                    }
122                }));
123            }
124
125            if let Some(pending) = this.pending.as_mut() {
126                match pending.as_mut().poll(cx) {
127                    Poll::Ready(result) => {
128                        this.pending = None;
129
130                        match result {
131                            Ok(Some(batch)) => {
132                                return Poll::Ready(Some(batch));
133                            }
134                            Ok(None) => {
135                                let mut has_new = false;
136                                let is_end = loop {
137                                    // wait notify and consume all
138                                    match (&mut this.notify).poll_next_unpin(cx) {
139                                        Poll::Ready(Some(_)) => {
140                                            has_new = true;
141                                            // 消费所有通知
142                                            continue;
143                                        }
144                                        Poll::Ready(None) => break true,
145                                        Poll::Pending => break false,
146                                    }
147                                };
148                                if has_new {
149                                    continue;
150                                } else if is_end {
151                                    return Poll::Ready(None);
152                                } else {
153                                    return Poll::Pending;
154                                }
155                            }
156                            Err(err) => {
157                                log::error!("external buffer shift return error: {}", err);
158                                return Poll::Ready(None);
159                            }
160                        }
161                    }
162                    Poll::Pending => {
163                        return Poll::Pending;
164                    }
165                }
166            }
167        }
168    }
169}
170
171#[cfg(feature = "default")]
172pub fn create_external_buffered_stream<T, S, P>(
173    stream: S,
174    path: P,
175) -> Result<impl Stream<Item = T>, Error>
176where
177    T: ExternalBufferSerde + Send + 'static,
178    S: Stream<Item = T> + Send + Sync + 'static,
179    P: AsRef<std::path::Path>,
180{
181    Ok(ExternalBufferedStream::new(
182        stream,
183        ExternalBufferSled::new(path)?,
184    ))
185}
186
187#[cfg(feature = "queue")]
188pub fn create_queued_stream<T, S>(stream: S) -> Result<impl Stream<Item = T>, Error>
189where
190    T: Ord + Send + 'static,
191    S: Stream<Item = T> + Send + Sync + 'static,
192{
193    Ok(ExternalBufferedStream::new(
194        stream,
195        ExternalBufferQueue::new(),
196    ))
197}
external_buffered_stream/lib.rs

external_buffered_stream/
lib.rs
