rskit_stream/

lib.rs

//! Foundational async stream toolkit.
//!
//! `rskit-stream` owns the opinion-free, layer-zero building blocks for the
//! "observe → fan-out → consume" graph that recurs across config reloads,
//! service discovery, cache invalidation, secret rotation, and message
//! consumers. Sources, the bounded fan-out bus, cancellable consumer tasks,
//! and the `futures::Stream` extension operators that chain them all live
//! together at the foundation where any higher layer can reuse them without
//! inverting the layer order.
//!
//! Sequential, named-step workflows (run N steps, report progress, cancel)
//! are a different concern and live in `rskit-chain`.

#![warn(missing_docs)]

/// Bounded fan-out broadcaster source (`Broadcaster<T>`).
pub mod broadcaster;
/// Extension trait adding `rskit` operators to any `Stream`.
pub mod ext;
/// Higher-level stream operators (map, filter, fan-out, windowing, etc.).
pub mod operators;
/// Terminal sink combinators (`collect`, `drain`, `for_each`).
pub mod sink;
/// Stream source constructors (`from_slice`, `from_fn`, `from_channel`).
pub mod source;
/// Cancellable owned tasks (`SpawnedTask`, `TaskGroup`).
pub mod task;

pub use broadcaster::{BroadcastStream, Broadcaster, DEFAULT_BROADCAST_BUFFER};
pub use ext::RskitStreamExt;
pub use operators::combine::{concat, merge};
pub use sink::{collect, drain, for_each};
pub use source::{from_channel, from_fn, from_slice};
pub use task::{SpawnedTask, TaskGroup};

pub use tokio_util::sync::CancellationToken;

#[cfg(test)]
mod tests {
    use parking_lot::Mutex;
    use std::sync::Arc;
    use std::time::Duration;

    use futures::StreamExt as _;

    use crate::{RskitStreamExt, from_fn, from_slice, merge};

    // ── Sources ───────────────────────────────────────────────────────────

    /// `from_slice` must yield every item in the original order.
    #[tokio::test]
    async fn test_from_slice_yields_all_in_order() {
        let items = vec![1u32, 2, 3, 4, 5];
        let stream = from_slice(items.clone());
        let collected: Vec<u32> = stream.collect().await;
        assert_eq!(collected, items);
    }

    /// `from_slice` with an empty vec yields nothing.
    #[tokio::test]
    async fn test_from_slice_empty() {
        let stream = from_slice::<u32>(vec![]);
        let collected: Vec<u32> = stream.collect().await;
        assert!(collected.is_empty());
    }

    /// `from_fn` calls the function repeatedly and stops when it returns `None`.
    #[tokio::test]
    async fn test_from_fn_yields_until_none() {
        let counter = Arc::new(Mutex::new(0u32));
        let c = counter.clone();
        let stream = from_fn(move || {
            let c = c.clone();
            async move {
                let mut guard = c.lock();
                let next = if *guard < 5 {
                    let val = *guard;
                    *guard += 1;
                    Some(val)
                } else {
                    None
                };
                drop(guard);
                next
            }
        });
        let collected: Vec<u32> = stream.collect().await;
        assert_eq!(collected, vec![0, 1, 2, 3, 4]);
    }

    /// `from_fn` that immediately returns `None` yields nothing.
    #[tokio::test]
    async fn test_from_fn_immediate_none() {
        let stream = from_fn(|| async { None::<u32> });
        let collected: Vec<u32> = stream.collect().await;
        assert!(collected.is_empty());
    }

    /// `merge` interleaves two streams; the combined set of items must match.
    #[tokio::test]
    async fn test_merge_set_equality() {
        let s1 = from_slice(vec![1u32, 3, 5]);
        let s2 = from_slice(vec![2u32, 4, 6]);
        let mut combined: Vec<u32> = merge(s1, s2).collect().await;
        combined.sort_unstable();
        assert_eq!(combined, vec![1, 2, 3, 4, 5, 6]);
    }

    /// `merge` of two empty streams yields nothing.
    #[tokio::test]
    async fn test_merge_both_empty() {
        let s1 = from_slice::<u32>(vec![]);
        let s2 = from_slice::<u32>(vec![]);
        let combined: Vec<u32> = merge(s1, s2).collect().await;
        assert!(combined.is_empty());
    }

    // ── RskitStreamExt::rmap ──────────────────────────────────────────────

    /// `rmap` transforms each item via an async fallible function.
    #[tokio::test]
    async fn test_rmap_transforms_items() {
        let stream = from_slice(vec![1u32, 2, 3]);
        let results: Vec<_> = stream
            .rmap(|x| async move { Ok::<u32, rskit_errors::AppError>(x * 10) })
            .collect()
            .await;
        let values: Vec<u32> = results.into_iter().map(|r| r.unwrap()).collect();
        assert_eq!(values, vec![10, 20, 30]);
    }

    /// `rmap` propagates errors returned by the function.
    #[tokio::test]
    async fn test_rmap_propagates_error() {
        let stream = from_slice(vec![1u32, 2, 3]);
        let results: Vec<_> = stream
            .rmap(|x| async move {
                if x == 2 {
                    Err(rskit_errors::AppError::new(
                        rskit_errors::ErrorCode::Internal,
                        "bad item",
                    ))
                } else {
                    Ok(x)
                }
            })
            .collect()
            .await;
        assert!(results[0].is_ok());
        assert!(results[1].is_err());
        assert!(results[2].is_ok());
    }

    // ── RskitStreamExt::rfilter ───────────────────────────────────────────

    /// `rfilter` keeps only items satisfying the predicate.
    #[tokio::test]
    async fn test_rfilter_keeps_matching_items() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5, 6]);
        let evens: Vec<u32> = stream.rfilter(|x| x % 2 == 0).collect().await;
        assert_eq!(evens, vec![2, 4, 6]);
    }

    /// `rfilter` with a predicate that matches nothing yields an empty stream.
    #[tokio::test]
    async fn test_rfilter_no_match_yields_empty() {
        let stream = from_slice(vec![1u32, 3, 5]);
        let result: Vec<u32> = stream.rfilter(|x| x % 2 == 0).collect().await;
        assert!(result.is_empty());
    }

    // ── RskitStreamExt::rtap ──────────────────────────────────────────────

    /// `rtap` calls the side-effect for every item and passes items through unchanged.
    #[tokio::test]
    async fn test_rtap_calls_side_effect_and_passes_through() {
        let seen = Arc::new(Mutex::new(Vec::<u32>::new()));
        let seen_clone = seen.clone();

        let stream = from_slice(vec![10u32, 20, 30]);
        let output: Vec<u32> = stream
            .rtap(move |x| {
                let seen = seen_clone.clone();
                let val = *x;
                async move {
                    seen.lock().push(val);
                }
            })
            .collect()
            .await;

        assert_eq!(output, vec![10, 20, 30]);
        assert_eq!(*seen.lock(), vec![10, 20, 30]);
    }

    // ── RskitStreamExt::rreduce ───────────────────────────────────────────

    /// `rreduce` folds the entire stream into a single accumulated value.
    #[tokio::test]
    async fn test_rreduce_folds_to_single_value() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5]);
        let sum = stream.rreduce(0u32, |acc, x| acc + x).await;
        assert_eq!(sum, 15);
    }

    /// `rreduce` on an empty stream returns the initial accumulator.
    #[tokio::test]
    async fn test_rreduce_empty_stream_returns_init() {
        let stream = from_slice::<u32>(vec![]);
        let result = stream.rreduce(42u32, |acc, x| acc + x).await;
        assert_eq!(result, 42);
    }

    // ── RskitStreamExt::rparallel ─────────────────────────────────────────

    /// `rparallel` processes items concurrently and collects all results.
    #[tokio::test]
    async fn test_rparallel_collects_all_results() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5]);
        let mut results: Vec<u32> = stream
            .rparallel(
                3,
                |x| async move { Ok::<u32, rskit_errors::AppError>(x * 2) },
            )
            .collect::<Vec<_>>()
            .await
            .into_iter()
            .map(|r| r.unwrap())
            .collect();
        results.sort_unstable();
        assert_eq!(results, vec![2, 4, 6, 8, 10]);
    }

    /// `rparallel` propagates errors from the worker function.
    #[tokio::test]
    async fn test_rparallel_propagates_errors() {
        let stream = from_slice(vec![1u32, 2, 3]);
        let results: Vec<_> = stream
            .rparallel(2, |x| async move {
                if x == 2 {
                    Err(rskit_errors::AppError::new(
                        rskit_errors::ErrorCode::Internal,
                        "parallel error",
                    ))
                } else {
                    Ok(x)
                }
            })
            .collect()
            .await;
        let error_count = results.iter().filter(|r| r.is_err()).count();
        assert_eq!(error_count, 1);
    }

    // ── RskitStreamExt::rfan_out ──────────────────────────────────────────

    /// `rfan_out` applies N functions to each item and collects results in order.
    ///
    /// We use non-capturing closures (which are Copy + Clone) so the
    /// `F: Clone` bound on `rfan_out` is satisfied without unstable features.
    #[tokio::test]
    async fn test_rfan_out_applies_all_functions() {
        // Non-capturing closures are Copy, so they satisfy Clone.
        let add_one = |x: u32| std::future::ready(Ok::<u32, rskit_errors::AppError>(x + 1));
        let mul_two = |x: u32| std::future::ready(Ok::<u32, rskit_errors::AppError>(x * 2));

        // First check: single add_one function
        let stream_a = from_slice(vec![5u32, 10u32]);
        let res_a: Vec<_> = stream_a.rfan_out(1, vec![add_one]).collect().await;
        let res_a: Vec<Vec<_>> = res_a.into_iter().map(Result::unwrap).collect();
        assert_eq!(res_a[0][0], 6u32);
        assert_eq!(res_a[1][0], 11u32);

        // Second check: two homogeneous functions of the same concrete type
        let stream_b = from_slice(vec![5u32, 10u32]);
        let res_b: Vec<_> = stream_b.rfan_out(2, vec![add_one, mul_two]).collect().await;
        let res_b: Vec<Vec<_>> = res_b.into_iter().map(Result::unwrap).collect();
        // item 5  → [5+1=6, 5*2=10]
        assert_eq!(res_b[0][0], 6u32);
        assert_eq!(res_b[0][1], 10u32);
        // item 10 → [10+1=11, 10*2=20]
        assert_eq!(res_b[1][0], 11u32);
        assert_eq!(res_b[1][1], 20u32);
    }

    /// `rfan_out` with a single function behaves like rmap.
    #[tokio::test]
    async fn test_rfan_out_single_function() {
        let stream = from_slice(vec![3u32, 7u32]);
        // Non-capturing closure is Copy + Clone.
        let f = |x: u32| std::future::ready(Ok::<u32, rskit_errors::AppError>(x + 100));
        let results: Vec<_> = stream.rfan_out(1, vec![f]).collect().await;
        let results: Vec<Vec<_>> = results.into_iter().map(Result::unwrap).collect();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0][0], 103u32);
        assert_eq!(results[1][0], 107u32);
    }

    // ── Windowing: rbatch ─────────────────────────────────────────────────

    /// `rbatch` with size=3 produces batches of exactly 3 items when enough arrive.
    #[tokio::test]
    async fn test_rbatch_exact_size_batches() {
        tokio::time::pause();

        let stream = from_slice(vec![1u32, 2, 3, 4, 5, 6]);
        let handle = tokio::spawn(async move {
            stream
                .rbatch(3, Duration::from_millis(500))
                .collect::<Vec<_>>()
                .await
        });

        tokio::time::advance(Duration::from_millis(600)).await;
        let batches = handle.await.unwrap();

        assert_eq!(batches.len(), 2);
        assert_eq!(batches[0], vec![1, 2, 3]);
        assert_eq!(batches[1], vec![4, 5, 6]);
    }

    /// `rbatch` flushes a partial batch on timeout.
    #[tokio::test]
    async fn test_rbatch_partial_flush_on_timeout() {
        tokio::time::pause();

        // Channel-based stream so we can control item arrival timing.
        let (tx, rx) = tokio::sync::mpsc::channel::<u32>(16);
        let stream = crate::source::from_channel(rx);

        let handle = tokio::spawn(async move {
            stream
                .rbatch(10, Duration::from_millis(100))
                .collect::<Vec<_>>()
                .await
        });

        // Send 2 items then let the timeout fire.
        tx.send(1).await.unwrap();
        tx.send(2).await.unwrap();
        drop(tx); // close channel after items sent

        tokio::time::advance(Duration::from_millis(200)).await;
        let batches = handle.await.unwrap();

        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0], vec![1, 2]);
    }

    // ── Rate: rdebounce ───────────────────────────────────────────────────

    /// `rdebounce` only emits the last item when the quiet window expires.
    #[tokio::test]
    async fn test_rdebounce_emits_last_item() {
        tokio::time::pause();

        let (tx, rx) = tokio::sync::mpsc::channel::<u32>(16);
        let stream = crate::source::from_channel(rx);

        let handle = tokio::spawn(async move {
            stream
                .rdebounce(Duration::from_millis(100))
                .collect::<Vec<_>>()
                .await
        });

        // Three rapid items — only the last should pass through.
        tx.send(1).await.unwrap();
        tx.send(2).await.unwrap();
        tx.send(3).await.unwrap();
        drop(tx);

        tokio::time::advance(Duration::from_millis(200)).await;
        let result = handle.await.unwrap();

        // After the channel closes, the pending item must be flushed.
        assert!(!result.is_empty());
        assert_eq!(*result.last().unwrap(), 3u32);
    }

    // ── Rate: rdebounce_batch ─────────────────────────────────────────────

    /// A burst of items with no quiet gap collapses into one trailing-edge
    /// batch, in arrival order. Delays live inside the source so `start_paused`
    /// auto-advances the clock deterministically.
    #[tokio::test(start_paused = true)]
    async fn test_rdebounce_batch_collapses_a_burst() {
        let source = async_stream::stream! {
            yield 1u32;
            yield 2u32;
            // Still inside the 100ms window: only pushes the deadline out.
            tokio::time::sleep(Duration::from_millis(50)).await;
            yield 3u32;
        };
        let result: Vec<Vec<u32>> = source
            .rdebounce_batch(Duration::from_millis(100), 1024)
            .collect()
            .await;
        assert_eq!(result, vec![vec![1u32, 2, 3]]);
    }

    /// Items separated by a quiet gap wider than the window land in separate
    /// batches.
    #[tokio::test(start_paused = true)]
    async fn test_rdebounce_batch_splits_separate_windows() {
        let source = async_stream::stream! {
            yield 1u32;
            tokio::time::sleep(Duration::from_millis(150)).await;
            yield 2u32;
            tokio::time::sleep(Duration::from_millis(150)).await;
        };
        let result: Vec<Vec<u32>> = source
            .rdebounce_batch(Duration::from_millis(100), 1024)
            .collect()
            .await;
        assert_eq!(result, vec![vec![1u32], vec![2u32]]);
    }

    /// The source closing before the quiet window elapses still flushes the
    /// pending, not-yet-emitted window.
    #[tokio::test(start_paused = true)]
    async fn test_rdebounce_batch_flushes_pending_on_close() {
        let source = async_stream::stream! { yield 7u32; };
        let result: Vec<Vec<u32>> = source
            .rdebounce_batch(Duration::from_millis(100), 1024)
            .collect()
            .await;
        assert_eq!(result, vec![vec![7u32]]);
    }

    /// The `max_items` safety cap force-flushes mid-burst so a sustained input
    /// rate faster than the quiet window cannot grow the buffer without bound.
    #[tokio::test(start_paused = true)]
    async fn test_rdebounce_batch_force_flushes_at_cap() {
        // A burst of 5 items with no quiet gap: a cap of 2 must split it into
        // [1,2], [3,4], then flush the trailing [5] on close.
        let source = async_stream::stream! {
            for value in 1u32..=5 {
                yield value;
            }
        };
        let result: Vec<Vec<u32>> = source
            .rdebounce_batch(Duration::from_millis(100), 2)
            .collect()
            .await;
        assert_eq!(result, vec![vec![1u32, 2], vec![3u32, 4], vec![5u32]]);
    }

    /// Dropping the consumer (the debounced stream) tears the pipeline down: it
    /// owns the source, so the source receiver is dropped and further sends fail.
    #[tokio::test]
    async fn test_rdebounce_batch_consumer_drop_ends_pipeline() {
        let (tx, rx) = tokio::sync::mpsc::channel::<u32>(16);
        let batched =
            crate::source::from_channel(rx).rdebounce_batch(Duration::from_millis(100), 1024);
        drop(batched); // Consumer gone before any window is emitted.

        assert!(
            tx.send(0).await.is_err(),
            "dropping the debounced stream must close the source"
        );
    }

    // ── Rate: rthrottle ───────────────────────────────────────────────────

    /// `rthrottle` drops items arriving faster than the interval.
    #[tokio::test]
    async fn test_rthrottle_drops_fast_items() {
        tokio::time::pause();

        let stream = from_slice(vec![1u32, 2, 3, 4, 5]);
        let handle = tokio::spawn(async move {
            stream
                .rthrottle(Duration::from_millis(100))
                .collect::<Vec<_>>()
                .await
        });

        tokio::time::advance(Duration::from_millis(600)).await;
        let result = handle.await.unwrap();

        // The first item is always emitted; subsequent items are dropped
        // because the stream is synchronous and all items arrive "instantly"
        // before the interval can pass.
        assert!(!result.is_empty());
        assert_eq!(result[0], 1u32);
        // All items after the first should have been throttled away.
        assert!(result.len() < 5);
    }

    /// `rthrottle` with an extreme interval still emits the first item without
    /// panicking (guards against an underflowing `now - interval` seed).
    #[tokio::test]
    async fn test_rthrottle_huge_interval_emits_first_without_panic() {
        tokio::time::pause();

        let stream = from_slice(vec![1u32, 2, 3]);
        let result = stream
            .rthrottle(Duration::from_secs(u64::MAX))
            .collect::<Vec<_>>()
            .await;

        assert_eq!(result, vec![1u32]);
    }

    // ── Windowing: rtumbling_window ───────────────────────────────────────

    /// `rtumbling_window` emits a non-empty window when the timer fires.
    #[tokio::test]
    async fn test_rtumbling_window_emits_on_timer() {
        tokio::time::pause();

        let (tx, rx) = tokio::sync::mpsc::channel::<u32>(16);
        let stream = crate::source::from_channel(rx);

        let handle = tokio::spawn(async move {
            stream
                .rtumbling_window(Duration::from_millis(100), 128)
                .collect::<Vec<_>>()
                .await
        });

        // Send items that should land in the first window.
        tx.send(10).await.unwrap();
        tx.send(20).await.unwrap();
        tx.send(30).await.unwrap();
        drop(tx);

        tokio::time::advance(Duration::from_millis(200)).await;
        let windows = handle.await.unwrap();

        assert!(!windows.is_empty());
        let all_items: Vec<u32> = windows.into_iter().flatten().collect();
        let mut sorted = all_items;
        sorted.sort_unstable();
        assert_eq!(sorted, vec![10, 20, 30]);
    }

    /// `rtumbling_window` yields an empty stream when input is empty.
    #[tokio::test]
    async fn test_rtumbling_window_empty_input() {
        tokio::time::pause();

        let stream = from_slice::<u32>(vec![]);
        let handle = tokio::spawn(async move {
            stream
                .rtumbling_window(Duration::from_millis(100), 128)
                .collect::<Vec<_>>()
                .await
        });

        tokio::time::advance(Duration::from_millis(200)).await;
        let windows = handle.await.unwrap();
        assert!(windows.is_empty());
    }

    #[tokio::test]
    async fn test_rdistinct_filters_duplicates() {
        let stream = from_slice(vec![1u32, 2, 2, 3, 1, 4]);
        let values: Vec<u32> = stream.rdistinct().collect().await;
        assert_eq!(values, vec![1, 2, 3, 4]);
    }

    #[tokio::test]
    async fn test_rtake_and_rskip_compose() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5]);
        let values: Vec<u32> = stream.rskip(1).rtake(3).collect().await;
        assert_eq!(values, vec![2, 3, 4]);
    }

    #[tokio::test]
    async fn test_rpartition_splits_stream() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5, 6]);
        let (even_stream, odd_stream) = stream.rpartition(|value| value % 2 == 0);
        let (evens, odds) = tokio::join!(
            even_stream.collect::<Vec<_>>(),
            odd_stream.collect::<Vec<_>>()
        );
        assert_eq!(evens, vec![2, 4, 6]);
        assert_eq!(odds, vec![1, 3, 5]);
    }

    #[tokio::test]
    async fn test_rsliding_window_emits_overlapping_windows() {
        let stream = from_slice(vec![1u32, 2, 3, 4, 5]);
        let windows: Vec<Vec<u32>> = stream.rsliding_window(3, 1).collect().await;
        assert_eq!(windows, vec![vec![1, 2, 3], vec![2, 3, 4], vec![3, 4, 5]]);
    }
}