rskit-stream 0.2.0-alpha.2

Foundational async stream toolkit: bounded fan-out broadcaster, sources, cancellable tasks, and futures::Stream extension operators
Documentation 
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

//! Leading-edge windowing and batching operators.
//!
//! These operators partition a stream into `Vec<T>` windows whose timers run
//! **leading-edge**, from the first buffered item (`tumbling_window`, `batch`),
//! or by item count (`sliding_window`). Contrast the trailing-edge rate
//! operators in [`crate::operators::rate`], whose timers are driven by quiet gaps between
//! arrivals.

use std::collections::VecDeque;
use std::time::Duration;

use futures::{Stream, StreamExt as _};

use super::buffer::{take_window, window_capacity};

/// Collect items into non-overlapping time windows.
///
/// A window is emitted after `duration` has elapsed since the first item.
pub fn tumbling_window<S, T>(
    stream: S,
    duration: Duration,
    max_items: usize,
) -> impl Stream<Item = Vec<T>> + Send + 'static
where
    S: Stream<Item = T> + Send + 'static,
    T: Send + 'static,
{
    async_stream::stream! {
        tokio::pin!(stream);
        let (max_items, initial_capacity) = window_capacity(max_items);
        let mut buf: Vec<T> = Vec::with_capacity(initial_capacity);
        let mut deadline = tokio::time::Instant::now() + duration;
        loop {
            tokio::select! {
                item = stream.next() => {
                    if let Some(v) = item {
                        buf.push(v);
                        if buf.len() >= max_items {
                            yield take_window(&mut buf, initial_capacity);
                            deadline = tokio::time::Instant::now() + duration;
                        }
                    } else {
                        if !buf.is_empty() {
                            yield take_window(&mut buf, initial_capacity);
                        }
                        break;
                    }
                }
                () = tokio::time::sleep_until(deadline) => {
                    if !buf.is_empty() {
                        yield take_window(&mut buf, initial_capacity);
                    }
                    deadline = tokio::time::Instant::now() + duration;
                }
            }
        }
    }
}

/// Collect up to `size` items into a batch.
///
/// A batch is emitted either when `size` items arrive or when `timeout`
/// elapses since the first item in the batch.
pub fn batch<S, T>(
    stream: S,
    size: usize,
    timeout: Duration,
) -> impl Stream<Item = Vec<T>> + Send + 'static
where
    S: Stream<Item = T> + Send + 'static,
    T: Send + 'static,
{
    async_stream::stream! {
        tokio::pin!(stream);
        let size = size.max(1);
        let mut buf: Vec<T> = Vec::with_capacity(size);
        let mut deadline: Option<tokio::time::Instant> = None;

        loop {
            tokio::select! {
                item = stream.next() => {
                    if let Some(v) = item {
                        if buf.is_empty() {
                            deadline = Some(tokio::time::Instant::now() + timeout);
                        }
                        buf.push(v);
                        if buf.len() >= size {
                            deadline = None;
                            yield std::mem::take(&mut buf);
                        }
                    } else {
                        if !buf.is_empty() {
                            yield std::mem::take(&mut buf);
                        }
                        break;
                    }
                }
                () = async {
                    match deadline {
                        Some(d) => tokio::time::sleep_until(d).await,
                        None => std::future::pending::<()>().await,
                    }
                } => {
                    deadline = None;
                    if !buf.is_empty() {
                        yield std::mem::take(&mut buf);
                    }
                }
            }
        }
    }
}

/// Emit sliding windows of `size` items, advancing by `step` items each time.
pub fn sliding_window<S, T>(
    stream: S,
    size: usize,
    step: usize,
) -> impl Stream<Item = Vec<T>> + Send + 'static
where
    S: Stream<Item = T> + Send + 'static,
    T: Clone + Send + 'static,
{
    async_stream::stream! {
        tokio::pin!(stream);
        let mut window = VecDeque::with_capacity(size.max(1));
        let effective_step = step.max(1);

        while let Some(item) = stream.next().await {
            window.push_back(item);
            while window.len() > size {
                window.pop_front();
            }
            if size > 0 && window.len() == size {
                yield window.iter().cloned().collect();
                for _ in 0..effective_step.min(window.len()) {
                    window.pop_front();
                }
            }
        }
    }
}