1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
use futures_util::{
future::{select, Either},
stream::FuturesUnordered,
StreamExt,
};
use stop_token::future::FutureExt as _;
use super::*;
#[derive(Debug)]
struct DeferredStreamProcessorInner {
opt_deferred_stream_channel: Option<flume::Sender<PinBoxFutureStatic<()>>>,
opt_stopper: Option<StopSource>,
opt_join_handle: Option<MustJoinHandle<()>>,
}
/// Result of a per-item handler, controlling whether stream processing continues.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum DeferredStreamResult {
/// Stop processing the stream.
Done,
/// Keep processing the next item.
Continue,
}
/// Background processor for streams
/// Handles streams to completion, passing each item from the stream to a callback
#[derive(Debug, Clone)]
pub struct DeferredStreamProcessor {
inner: Arc<Mutex<DeferredStreamProcessorInner>>,
}
impl DeferredStreamProcessor {
/// Create a new DeferredStreamProcessor
#[must_use]
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(DeferredStreamProcessorInner {
opt_deferred_stream_channel: None,
opt_stopper: None,
opt_join_handle: None,
})),
}
}
/// Initialize the processor before use
///
/// Pairs with `terminate`. Spawns the background task; not idempotent, calling it twice panics dropping the prior un-awaited join handle.
pub fn init(&self) {
let stopper = StopSource::new();
let stop_token = stopper.token();
let mut inner = self.inner.lock();
inner.opt_stopper = Some(stopper);
let (dsc_tx, dsc_rx) = flume::unbounded::<PinBoxFutureStatic<()>>();
inner.opt_deferred_stream_channel = Some(dsc_tx);
inner.opt_join_handle = Some(spawn(
"deferred stream processor",
Self::processor(stop_token, dsc_rx),
));
}
/// Terminate the processor and ensure all streams are closed
///
/// Awaits the background task's exit. Idempotent: a no-op once already terminated or never initialized.
pub async fn terminate(&self) {
let opt_jh = {
let mut inner = self.inner.lock();
drop(inner.opt_deferred_stream_channel.take());
drop(inner.opt_stopper.take());
inner.opt_join_handle.take()
};
if let Some(jh) = opt_jh {
jh.await;
}
}
async fn processor(stop_token: StopToken, dsc_rx: flume::Receiver<PinBoxFutureStatic<()>>) {
let mut unord = FuturesUnordered::<PinBoxFutureStatic<()>>::new();
// Ensure the unord never finishes unless the stop source got dropped
unord.push(Box::pin(stop_token));
// Processor loop
let mut unord_fut = unord.next();
let mut dsc_fut = dsc_rx.recv_async();
loop {
let res = select(unord_fut, dsc_fut).await;
match res {
Either::Left((x, old_dsc_fut)) => {
// If the unord future gets empty, the stop token got dropped and all the other tasks finished
if x.is_none() {
break;
}
// Make another unord future to process
unord_fut = unord.next();
// put back the other future and keep going
dsc_fut = old_dsc_fut;
}
Either::Right((new_proc, old_unord_fut)) => {
// Immediately drop the old unord future
// because we never care about it completing
drop(old_unord_fut);
let Ok(new_proc) = new_proc else {
break;
};
// Add a new stream to process
unord.push(new_proc);
// Make a new unord future because we don't care about the
// completion of the last unord future, they never return
// anything.
unord_fut = unord.next();
// Make a new receiver future
dsc_fut = dsc_rx.recv_async();
}
}
}
}
/// Queue a stream to process in the background
///
/// * 'receiver' is the stream to process
/// * 'handler' is the callback to handle each item from the stream
///
/// Returns 'true' if the stream was added for processing, and 'false' if the stream could not be added, possibly due to not being initialized.
///
/// Non-blocking; the stream is processed by the background task until exhausted or `terminate` is called.
pub fn add_stream<
T: Send + 'static,
S: futures_util::Stream<Item = T> + Unpin + Send + 'static,
>(
&self,
mut receiver: S,
mut handler: impl FnMut(T) -> PinBoxFutureStatic<DeferredStreamResult> + Send + 'static,
) -> bool {
let (st, dsc_tx) = {
let inner = self.inner.lock();
let Some(st) = inner.opt_stopper.as_ref().map(|s| s.token()) else {
return false;
};
let Some(dsc_tx) = inner.opt_deferred_stream_channel.clone() else {
return false;
};
(st, dsc_tx)
};
let drp = Box::pin(async move {
while let Ok(Some(res)) = receiver.next().timeout_at(st.clone()).await {
if matches!(handler(res).await, DeferredStreamResult::Done) {
break;
}
}
});
if dsc_tx.send(drp).is_err() {
return false;
}
true
}
/// Queue a single future to process in the background
///
/// Non-blocking; returns false if the processor was not initialized. The future runs until done or `terminate` is called.
pub fn add_future<F>(&self, fut: F) -> bool
where
F: Future<Output = ()> + Send + 'static,
{
let (st, dsc_tx) = {
let inner = self.inner.lock();
let Some(st) = inner.opt_stopper.as_ref().map(|s| s.token()) else {
return false;
};
let Some(dsc_tx) = inner.opt_deferred_stream_channel.clone() else {
return false;
};
(st, dsc_tx)
};
if dsc_tx
.send(Box::pin(async move {
let _ = fut.timeout_at(st.clone()).await;
}))
.is_err()
{
return false;
}
true
}
}
impl Default for DeferredStreamProcessor {
fn default() -> Self {
Self::new()
}
}