command-stream 1.8.0

//! Bounded, composable asynchronous data streams and online analytics.
use crate::{OutputChunk, OutputStream};
use std::collections::VecDeque;
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;

struct Task(JoinHandle<()>);
impl Drop for Task {
    fn drop(&mut self) {
        self.0.abort();
    }
}

/// Each stage has at most 32 queued values; dropping a consumer stops its producer.
pub struct DataStream<T> {
    receiver: mpsc::Receiver<T>,
    task: Arc<Task>,
}

#[derive(Debug, Clone, Default)]
pub struct StreamStats {
    pub count: u64,
    pub errors: u64,
    pub error_rate: f64,
    pub throughput: f64,
    pub elapsed: f64,
    pub response_count: u64,
    pub response_mean: f64,
    pub response_min: Option<f64>,
    pub response_max: Option<f64>,
}

impl<T: Send + 'static> DataStream<T> {
    fn spawn<F, Fut>(producer: F) -> Self
    where
        F: FnOnce(mpsc::Sender<T>) -> Fut,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let (tx, receiver) = mpsc::channel(32);
        let task = Arc::new(Task(tokio::spawn(producer(tx))));
        Self { receiver, task }
    }
    pub fn from_values<I>(values: I) -> Self
    where
        I: IntoIterator<Item = T> + Send + 'static,
        I::IntoIter: Send,
    {
        Self::spawn(|tx| async move {
            for value in values {
                if tx.send(value).await.is_err() {
                    break;
                }
            }
        })
    }
    pub async fn next(&mut self) -> Option<T> {
        self.receiver.recv().await
    }
    pub fn map<U, F>(self, mut transform: F) -> DataStream<U>
    where
        U: Send + 'static,
        F: FnMut(T) -> U + Send + 'static,
    {
        self.map_async(move |value| std::future::ready(transform(value)))
    }
    pub fn map_async<U, F, Fut>(mut self, mut transform: F) -> DataStream<U>
    where
        U: Send + 'static,
        F: FnMut(T) -> Fut + Send + 'static,
        Fut: Future<Output = U> + Send,
    {
        DataStream::spawn(|tx| async move {
            while let Some(value) = self.next().await {
                if tx.send(transform(value).await).await.is_err() {
                    break;
                }
            }
        })
    }
    pub fn filter<F>(mut self, mut predicate: F) -> Self
    where
        F: FnMut(&T) -> bool + Send + 'static,
    {
        Self::spawn(|tx| async move {
            while let Some(value) = self.next().await {
                if predicate(&value) && tx.send(value).await.is_err() {
                    break;
                }
            }
        })
    }
    pub fn batch(self, size: usize) -> DataStream<Vec<T>> {
        self.batch_inner(size, None)
    }
    /// Flush at size or a deadline measured from the first value in each batch.
    /// An idle source still flushes; a slow consumer applies channel backpressure.
    pub fn batch_with_timeout(self, size: usize, time_window: Duration) -> DataStream<Vec<T>> {
        assert!(!time_window.is_zero(), "time window must be positive");
        assert!(
            tokio::time::Instant::now()
                .checked_add(time_window)
                .is_some(),
            "time window must fit the clock"
        );
        self.batch_inner(size, Some(time_window))
    }
    fn batch_inner(mut self, size: usize, time_window: Option<Duration>) -> DataStream<Vec<T>> {
        assert!(size > 0);
        DataStream::spawn(|tx| async move {
            let mut values = Vec::with_capacity(size);
            let mut deadline = None;
            loop {
                let value = if let Some(end) = deadline {
                    // mpsc::Receiver::recv is cancellation-safe: timing out a
                    // read cannot remove a value from the source channel.
                    let read = if tokio::time::Instant::now() >= end {
                        None
                    } else {
                        tokio::time::timeout_at(end, self.next()).await.ok()
                    };
                    match read {
                        Some(value) => value,
                        None => {
                            if tx.send(std::mem::take(&mut values)).await.is_err() {
                                return;
                            }
                            deadline = None;
                            continue;
                        }
                    }
                } else {
                    self.next().await
                };
                let Some(value) = value else { break };
                if values.is_empty() {
                    deadline = time_window.map(|duration| tokio::time::Instant::now() + duration);
                }
                values.push(value);
                if values.len() == size {
                    if tx.send(std::mem::take(&mut values)).await.is_err() {
                        return;
                    }
                    deadline = None;
                }
            }
            if !values.is_empty() {
                let _ = tx.send(values).await;
            }
        })
    }
    pub fn window(mut self, size: usize) -> DataStream<Vec<T>>
    where
        T: Clone,
    {
        assert!(size > 0);
        DataStream::spawn(|tx| async move {
            let mut values = VecDeque::with_capacity(size);
            while let Some(value) = self.next().await {
                values.push_back(value);
                if values.len() > size {
                    values.pop_front();
                }
                if values.len() == size && tx.send(values.iter().cloned().collect()).await.is_err()
                {
                    break;
                }
            }
        })
    }
    /// Alias for a sliding window advancing by one value.
    pub fn sliding_window(self, size: usize) -> DataStream<Vec<T>>
    where
        T: Clone,
    {
        self.window(size)
    }
    pub fn split<F>(mut self, mut predicate: F) -> (Self, Self)
    where
        F: FnMut(&T) -> bool + Send + 'static,
    {
        let (yes_tx, yes) = mpsc::channel(32);
        let (no_tx, no) = mpsc::channel(32);
        let task = Arc::new(Task(tokio::spawn(async move {
            while let Some(value) = self.next().await {
                if yes_tx.is_closed() && no_tx.is_closed() {
                    break;
                }
                let tx = if predicate(&value) { &yes_tx } else { &no_tx };
                let _ = tx.send(value).await;
            }
        })));
        (
            Self {
                receiver: yes,
                task: task.clone(),
            },
            Self { receiver: no, task },
        )
    }
    pub fn merge(mut self, mut other: Self) -> Self {
        Self::spawn(|tx| async move {
            let mut left = true;
            let mut right = true;
            while left || right {
                let value = tokio::select! { value = self.next(), if left => { if value.is_none() { left = false; } value },
                value = other.next(), if right => { if value.is_none() { right = false; } value } };
                if let Some(value) = value {
                    if tx.send(value).await.is_err() {
                        break;
                    }
                }
            }
        })
    }
    pub async fn collect(mut self) -> Vec<T> {
        let mut values = Vec::new();
        while let Some(value) = self.next().await {
            values.push(value);
        }
        values
    }
    pub async fn reduce<U, F>(mut self, mut initial: U, mut reducer: F) -> U
    where
        F: FnMut(U, T) -> U,
    {
        while let Some(value) = self.next().await {
            initial = reducer(initial, value);
        }
        initial
    }
    pub async fn for_each<F, Fut>(mut self, mut callback: F)
    where
        F: FnMut(T) -> Fut,
        Fut: Future<Output = ()>,
    {
        while let Some(value) = self.next().await {
            callback(value).await;
        }
    }
    pub async fn analyze<E, R, U>(
        mut self,
        mut error: E,
        mut response: R,
        mut update: U,
    ) -> StreamStats
    where
        E: FnMut(&T) -> bool,
        R: FnMut(&T) -> Option<f64>,
        U: FnMut(&StreamStats),
    {
        let started = Instant::now();
        let mut stats = StreamStats::default();
        let mut sum = 0.0;
        while let Some(value) = self.next().await {
            stats.count += 1;
            stats.errors += u64::from(error(&value));
            if let Some(time) = response(&value).filter(|time| time.is_finite()) {
                stats.response_count += 1;
                sum += time;
                stats.response_min = Some(stats.response_min.map_or(time, |min| min.min(time)));
                stats.response_max = Some(stats.response_max.map_or(time, |max| max.max(time)));
                stats.response_mean = sum / stats.response_count as f64;
            }
            stats.elapsed = started.elapsed().as_secs_f64();
            stats.error_rate = stats.errors as f64 / stats.count as f64;
            stats.throughput = stats.count as f64 / stats.elapsed.max(0.001);
            update(&stats);
        }
        stats
    }
}

impl<T> Drop for DataStream<T> {
    fn drop(&mut self) {
        self.receiver.close();
        let _ = &self.task;
    }
}

impl OutputStream {
    /// Frame stdout bytes into UTF-8 lines, retaining split codepoints until newline.
    pub fn lines(mut self) -> DataStream<String> {
        DataStream::spawn(|tx| async move {
            let mut pending = Vec::new();
            while let Some(chunk) = self.next().await {
                if let OutputChunk::Stdout(bytes) = chunk {
                    pending.extend(bytes);
                    while let Some(end) = pending.iter().position(|byte| *byte == b'\n') {
                        let mut line = pending.drain(..=end).collect::<Vec<_>>();
                        line.pop();
                        if line.last() == Some(&b'\r') {
                            line.pop();
                        }
                        if tx
                            .send(String::from_utf8_lossy(&line).into_owned())
                            .await
                            .is_err()
                        {
                            return;
                        }
                    }
                }
            }
            if !pending.is_empty() {
                let _ = tx
                    .send(String::from_utf8_lossy(&pending).into_owned())
                    .await;
            }
        })
    }
}