command_stream/
stream_processing.rs

1//! Bounded, composable asynchronous data streams and online analytics.
2use crate::{OutputChunk, OutputStream};
3use std::collections::VecDeque;
4use std::future::Future;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7use tokio::sync::mpsc;
8use tokio::task::JoinHandle;
9
10struct Task(JoinHandle<()>);
11impl Drop for Task {
12    fn drop(&mut self) {
13        self.0.abort();
14    }
15}
16
17/// Each stage has at most 32 queued values; dropping a consumer stops its producer.
18pub struct DataStream<T> {
19    receiver: mpsc::Receiver<T>,
20    task: Arc<Task>,
21}
22
23#[derive(Debug, Clone, Default)]
24pub struct StreamStats {
25    pub count: u64,
26    pub errors: u64,
27    pub error_rate: f64,
28    pub throughput: f64,
29    pub elapsed: f64,
30    pub response_count: u64,
31    pub response_mean: f64,
32    pub response_min: Option<f64>,
33    pub response_max: Option<f64>,
34}
35
36impl<T: Send + 'static> DataStream<T> {
37    fn spawn<F, Fut>(producer: F) -> Self
38    where
39        F: FnOnce(mpsc::Sender<T>) -> Fut,
40        Fut: Future<Output = ()> + Send + 'static,
41    {
42        let (tx, receiver) = mpsc::channel(32);
43        let task = Arc::new(Task(tokio::spawn(producer(tx))));
44        Self { receiver, task }
45    }
46    pub fn from_values<I>(values: I) -> Self
47    where
48        I: IntoIterator<Item = T> + Send + 'static,
49        I::IntoIter: Send,
50    {
51        Self::spawn(|tx| async move {
52            for value in values {
53                if tx.send(value).await.is_err() {
54                    break;
55                }
56            }
57        })
58    }
59    pub async fn next(&mut self) -> Option<T> {
60        self.receiver.recv().await
61    }
62    pub fn map<U, F>(self, mut transform: F) -> DataStream<U>
63    where
64        U: Send + 'static,
65        F: FnMut(T) -> U + Send + 'static,
66    {
67        self.map_async(move |value| std::future::ready(transform(value)))
68    }
69    pub fn map_async<U, F, Fut>(mut self, mut transform: F) -> DataStream<U>
70    where
71        U: Send + 'static,
72        F: FnMut(T) -> Fut + Send + 'static,
73        Fut: Future<Output = U> + Send,
74    {
75        DataStream::spawn(|tx| async move {
76            while let Some(value) = self.next().await {
77                if tx.send(transform(value).await).await.is_err() {
78                    break;
79                }
80            }
81        })
82    }
83    pub fn filter<F>(mut self, mut predicate: F) -> Self
84    where
85        F: FnMut(&T) -> bool + Send + 'static,
86    {
87        Self::spawn(|tx| async move {
88            while let Some(value) = self.next().await {
89                if predicate(&value) && tx.send(value).await.is_err() {
90                    break;
91                }
92            }
93        })
94    }
95    pub fn batch(self, size: usize) -> DataStream<Vec<T>> {
96        self.batch_inner(size, None)
97    }
98    /// Flush at size or a deadline measured from the first value in each batch.
99    /// An idle source still flushes; a slow consumer applies channel backpressure.
100    pub fn batch_with_timeout(self, size: usize, time_window: Duration) -> DataStream<Vec<T>> {
101        assert!(!time_window.is_zero(), "time window must be positive");
102        assert!(
103            tokio::time::Instant::now()
104                .checked_add(time_window)
105                .is_some(),
106            "time window must fit the clock"
107        );
108        self.batch_inner(size, Some(time_window))
109    }
110    fn batch_inner(mut self, size: usize, time_window: Option<Duration>) -> DataStream<Vec<T>> {
111        assert!(size > 0);
112        DataStream::spawn(|tx| async move {
113            let mut values = Vec::with_capacity(size);
114            let mut deadline = None;
115            loop {
116                let value = if let Some(end) = deadline {
117                    // mpsc::Receiver::recv is cancellation-safe: timing out a
118                    // read cannot remove a value from the source channel.
119                    let read = if tokio::time::Instant::now() >= end {
120                        None
121                    } else {
122                        tokio::time::timeout_at(end, self.next()).await.ok()
123                    };
124                    match read {
125                        Some(value) => value,
126                        None => {
127                            if tx.send(std::mem::take(&mut values)).await.is_err() {
128                                return;
129                            }
130                            deadline = None;
131                            continue;
132                        }
133                    }
134                } else {
135                    self.next().await
136                };
137                let Some(value) = value else { break };
138                if values.is_empty() {
139                    deadline = time_window.map(|duration| tokio::time::Instant::now() + duration);
140                }
141                values.push(value);
142                if values.len() == size {
143                    if tx.send(std::mem::take(&mut values)).await.is_err() {
144                        return;
145                    }
146                    deadline = None;
147                }
148            }
149            if !values.is_empty() {
150                let _ = tx.send(values).await;
151            }
152        })
153    }
154    pub fn window(mut self, size: usize) -> DataStream<Vec<T>>
155    where
156        T: Clone,
157    {
158        assert!(size > 0);
159        DataStream::spawn(|tx| async move {
160            let mut values = VecDeque::with_capacity(size);
161            while let Some(value) = self.next().await {
162                values.push_back(value);
163                if values.len() > size {
164                    values.pop_front();
165                }
166                if values.len() == size && tx.send(values.iter().cloned().collect()).await.is_err()
167                {
168                    break;
169                }
170            }
171        })
172    }
173    /// Alias for a sliding window advancing by one value.
174    pub fn sliding_window(self, size: usize) -> DataStream<Vec<T>>
175    where
176        T: Clone,
177    {
178        self.window(size)
179    }
180    pub fn split<F>(mut self, mut predicate: F) -> (Self, Self)
181    where
182        F: FnMut(&T) -> bool + Send + 'static,
183    {
184        let (yes_tx, yes) = mpsc::channel(32);
185        let (no_tx, no) = mpsc::channel(32);
186        let task = Arc::new(Task(tokio::spawn(async move {
187            while let Some(value) = self.next().await {
188                if yes_tx.is_closed() && no_tx.is_closed() {
189                    break;
190                }
191                let tx = if predicate(&value) { &yes_tx } else { &no_tx };
192                let _ = tx.send(value).await;
193            }
194        })));
195        (
196            Self {
197                receiver: yes,
198                task: task.clone(),
199            },
200            Self { receiver: no, task },
201        )
202    }
203    pub fn merge(mut self, mut other: Self) -> Self {
204        Self::spawn(|tx| async move {
205            let mut left = true;
206            let mut right = true;
207            while left || right {
208                let value = tokio::select! { value = self.next(), if left => { if value.is_none() { left = false; } value },
209                value = other.next(), if right => { if value.is_none() { right = false; } value } };
210                if let Some(value) = value {
211                    if tx.send(value).await.is_err() {
212                        break;
213                    }
214                }
215            }
216        })
217    }
218    pub async fn collect(mut self) -> Vec<T> {
219        let mut values = Vec::new();
220        while let Some(value) = self.next().await {
221            values.push(value);
222        }
223        values
224    }
225    pub async fn reduce<U, F>(mut self, mut initial: U, mut reducer: F) -> U
226    where
227        F: FnMut(U, T) -> U,
228    {
229        while let Some(value) = self.next().await {
230            initial = reducer(initial, value);
231        }
232        initial
233    }
234    pub async fn for_each<F, Fut>(mut self, mut callback: F)
235    where
236        F: FnMut(T) -> Fut,
237        Fut: Future<Output = ()>,
238    {
239        while let Some(value) = self.next().await {
240            callback(value).await;
241        }
242    }
243    pub async fn analyze<E, R, U>(
244        mut self,
245        mut error: E,
246        mut response: R,
247        mut update: U,
248    ) -> StreamStats
249    where
250        E: FnMut(&T) -> bool,
251        R: FnMut(&T) -> Option<f64>,
252        U: FnMut(&StreamStats),
253    {
254        let started = Instant::now();
255        let mut stats = StreamStats::default();
256        let mut sum = 0.0;
257        while let Some(value) = self.next().await {
258            stats.count += 1;
259            stats.errors += u64::from(error(&value));
260            if let Some(time) = response(&value).filter(|time| time.is_finite()) {
261                stats.response_count += 1;
262                sum += time;
263                stats.response_min = Some(stats.response_min.map_or(time, |min| min.min(time)));
264                stats.response_max = Some(stats.response_max.map_or(time, |max| max.max(time)));
265                stats.response_mean = sum / stats.response_count as f64;
266            }
267            stats.elapsed = started.elapsed().as_secs_f64();
268            stats.error_rate = stats.errors as f64 / stats.count as f64;
269            stats.throughput = stats.count as f64 / stats.elapsed.max(0.001);
270            update(&stats);
271        }
272        stats
273    }
274}
275
276impl<T> Drop for DataStream<T> {
277    fn drop(&mut self) {
278        self.receiver.close();
279        let _ = &self.task;
280    }
281}
282
283impl OutputStream {
284    /// Frame stdout bytes into UTF-8 lines, retaining split codepoints until newline.
285    pub fn lines(mut self) -> DataStream<String> {
286        DataStream::spawn(|tx| async move {
287            let mut pending = Vec::new();
288            while let Some(chunk) = self.next().await {
289                if let OutputChunk::Stdout(bytes) = chunk {
290                    pending.extend(bytes);
291                    while let Some(end) = pending.iter().position(|byte| *byte == b'\n') {
292                        let mut line = pending.drain(..=end).collect::<Vec<_>>();
293                        line.pop();
294                        if line.last() == Some(&b'\r') {
295                            line.pop();
296                        }
297                        if tx
298                            .send(String::from_utf8_lossy(&line).into_owned())
299                            .await
300                            .is_err()
301                        {
302                            return;
303                        }
304                    }
305                }
306            }
307            if !pending.is_empty() {
308                let _ = tx
309                    .send(String::from_utf8_lossy(&pending).into_owned())
310                    .await;
311            }
312        })
313    }
314}
command_stream/stream_processing.rs

command_stream/
stream_processing.rs
