use alloc::collections::VecDeque;
use alloc::vec::Vec;
use core::time::Duration;
use std::time::Instant;
use crate::emit::Emit;
use crate::source::Infallible;
use crate::stage::Stage;
pub trait Clock: Send {
fn now(&self) -> Instant;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct SystemClock;
impl Clock for SystemClock {
fn now(&self) -> Instant {
Instant::now()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WindowPolicy {
Tumbling {
size: Duration,
},
Sliding {
size: Duration,
slide: Duration,
},
Session {
idle: Duration,
},
}
#[derive(Debug, Clone)]
pub struct Window<T> {
items: Vec<T>,
start: Instant,
end: Instant,
}
impl<T> Window<T> {
#[must_use]
pub fn new(items: Vec<T>, start: Instant, end: Instant) -> Self {
Self { items, start, end }
}
#[must_use]
pub fn items(&self) -> &[T] {
&self.items
}
#[must_use]
pub fn len(&self) -> usize {
self.items.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
#[must_use]
pub fn start(&self) -> Instant {
self.start
}
#[must_use]
pub fn end(&self) -> Instant {
self.end
}
#[must_use]
pub fn into_inner(self) -> Vec<T> {
self.items
}
}
impl<T> IntoIterator for Window<T> {
type Item = T;
type IntoIter = alloc::vec::IntoIter<T>;
fn into_iter(self) -> Self::IntoIter {
self.items.into_iter()
}
}
impl<'a, T> IntoIterator for &'a Window<T> {
type Item = &'a T;
type IntoIter = core::slice::Iter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.items.iter()
}
}
struct PendingWindow<T> {
start: Instant,
end: Instant,
items: Vec<T>,
}
pub(crate) struct WindowStage<T: Clone, C: Clock> {
policy: WindowPolicy,
clock: C,
state: WindowState<T>,
}
enum WindowState<T> {
Tumbling {
start: Option<Instant>,
items: Vec<T>,
},
Sliding {
windows: VecDeque<PendingWindow<T>>,
next_window_start: Option<Instant>,
},
Session {
start: Option<Instant>,
last_seen: Option<Instant>,
items: Vec<T>,
},
}
impl<T: Clone, C: Clock> WindowStage<T, C> {
pub(crate) fn new(policy: WindowPolicy, clock: C) -> Self {
let state = match policy {
WindowPolicy::Tumbling { .. } => WindowState::Tumbling {
start: None,
items: Vec::new(),
},
WindowPolicy::Sliding { .. } => WindowState::Sliding {
windows: VecDeque::new(),
next_window_start: None,
},
WindowPolicy::Session { .. } => WindowState::Session {
start: None,
last_seen: None,
items: Vec::new(),
},
};
Self {
policy,
clock,
state,
}
}
}
impl<T, C> Stage for WindowStage<T, C>
where
T: Clone + Send + 'static,
C: Clock + 'static,
{
type Input = T;
type Output = Window<T>;
type Error = Infallible;
fn process(
&mut self,
item: Self::Input,
out: &mut dyn Emit<Item = Self::Output>,
) -> Result<(), Self::Error> {
let now = self.clock.now();
match (&self.policy, &mut self.state) {
(WindowPolicy::Tumbling { size }, WindowState::Tumbling { start, items }) => {
if start.is_none() {
*start = Some(now);
}
while let Some(s) = *start {
if now.saturating_duration_since(s) >= *size {
let window_items = core::mem::take(items);
let _ = out.emit(Window::new(window_items, s, s + *size));
*start = Some(s + *size);
} else {
break;
}
}
items.push(item);
}
(
WindowPolicy::Session { idle },
WindowState::Session {
start,
last_seen,
items,
},
) => {
if let Some(ls) = *last_seen {
if now.saturating_duration_since(ls) > *idle {
if let Some(s) = *start {
let window_items = core::mem::take(items);
let _ = out.emit(Window::new(window_items, s, ls));
}
*start = Some(now);
}
} else {
*start = Some(now);
}
*last_seen = Some(now);
items.push(item);
}
(
WindowPolicy::Sliding { size, slide },
WindowState::Sliding {
windows,
next_window_start,
},
) => {
if next_window_start.is_none() {
*next_window_start = Some(now);
}
while let Some(s) = *next_window_start {
if s <= now {
windows.push_back(PendingWindow {
start: s,
end: s + *size,
items: Vec::new(),
});
*next_window_start = Some(s + *slide);
} else {
break;
}
}
while let Some(w) = windows.front() {
if w.end <= now {
let w = windows.pop_front().expect("front exists");
let _ = out.emit(Window::new(w.items, w.start, w.end));
} else {
break;
}
}
for w in windows.iter_mut() {
if w.start <= now && now < w.end {
w.items.push(item.clone());
}
}
}
_ => unreachable!(
"policy/state mismatch is impossible by construction; \
WindowStage::new enforces alignment"
),
}
Ok(())
}
fn flush(&mut self, out: &mut dyn Emit<Item = Self::Output>) -> Result<(), Self::Error> {
let now = self.clock.now();
match &mut self.state {
WindowState::Tumbling { start, items } => {
if !items.is_empty() {
let s = start.unwrap_or(now);
let window_items = core::mem::take(items);
let _ = out.emit(Window::new(window_items, s, now));
}
}
WindowState::Session {
start,
last_seen,
items,
} => {
if !items.is_empty() {
let s = start.unwrap_or(now);
let e = last_seen.unwrap_or(now);
let window_items = core::mem::take(items);
let _ = out.emit(Window::new(window_items, s, e));
}
}
WindowState::Sliding { windows, .. } => {
while let Some(w) = windows.pop_front() {
if !w.items.is_empty() {
let _ = out.emit(Window::new(w.items, w.start, w.end));
}
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::emit::EmitError;
use std::sync::{Arc, Mutex};
#[derive(Clone)]
struct FakeClock {
inner: Arc<Mutex<Instant>>,
}
impl FakeClock {
fn new(start: Instant) -> Self {
Self {
inner: Arc::new(Mutex::new(start)),
}
}
fn advance(&self, by: Duration) {
let mut g = self.inner.lock().unwrap();
*g += by;
}
}
impl Clock for FakeClock {
fn now(&self) -> Instant {
*self.inner.lock().unwrap()
}
}
struct Collect<T> {
out: Vec<Window<T>>,
}
impl<T> Emit for Collect<T> {
type Item = Window<T>;
fn emit(&mut self, w: Window<T>) -> Result<(), EmitError> {
self.out.push(w);
Ok(())
}
}
#[test]
fn tumbling_emits_on_boundary() {
let t0 = Instant::now();
let clock = FakeClock::new(t0);
let mut stage = WindowStage::<u32, _>::new(
WindowPolicy::Tumbling {
size: Duration::from_secs(10),
},
clock.clone(),
);
let mut emit = Collect::<u32> { out: Vec::new() };
stage.process(1, &mut emit).unwrap();
assert!(emit.out.is_empty());
clock.advance(Duration::from_secs(5));
stage.process(2, &mut emit).unwrap();
assert!(emit.out.is_empty());
clock.advance(Duration::from_secs(5));
stage.process(3, &mut emit).unwrap();
assert_eq!(emit.out.len(), 1);
assert_eq!(emit.out[0].items(), &[1, 2]);
clock.advance(Duration::from_secs(10));
stage.process(4, &mut emit).unwrap();
assert_eq!(emit.out.len(), 2);
assert_eq!(emit.out[1].items(), &[3]);
stage.flush(&mut emit).unwrap();
assert_eq!(emit.out.len(), 3);
assert_eq!(emit.out[2].items(), &[4]);
}
#[test]
fn session_closes_after_idle() {
let t0 = Instant::now();
let clock = FakeClock::new(t0);
let mut stage = WindowStage::<u32, _>::new(
WindowPolicy::Session {
idle: Duration::from_secs(5),
},
clock.clone(),
);
let mut emit = Collect::<u32> { out: Vec::new() };
stage.process(1, &mut emit).unwrap();
clock.advance(Duration::from_secs(2));
stage.process(2, &mut emit).unwrap();
clock.advance(Duration::from_secs(2));
stage.process(3, &mut emit).unwrap();
assert!(emit.out.is_empty());
clock.advance(Duration::from_secs(16));
stage.process(4, &mut emit).unwrap();
assert_eq!(emit.out.len(), 1);
assert_eq!(emit.out[0].items(), &[1, 2, 3]);
stage.flush(&mut emit).unwrap();
assert_eq!(emit.out.len(), 2);
assert_eq!(emit.out[1].items(), &[4]);
}
#[test]
fn sliding_overlapping_windows() {
let t0 = Instant::now();
let clock = FakeClock::new(t0);
let mut stage = WindowStage::<u32, _>::new(
WindowPolicy::Sliding {
size: Duration::from_secs(10),
slide: Duration::from_secs(5),
},
clock.clone(),
);
let mut emit = Collect::<u32> { out: Vec::new() };
stage.process(1, &mut emit).unwrap();
clock.advance(Duration::from_secs(3));
stage.process(2, &mut emit).unwrap();
clock.advance(Duration::from_secs(2));
stage.process(3, &mut emit).unwrap();
clock.advance(Duration::from_secs(5));
stage.process(4, &mut emit).unwrap();
assert_eq!(emit.out.len(), 1);
assert_eq!(emit.out[0].items(), &[1, 2, 3]);
stage.flush(&mut emit).unwrap();
assert_eq!(emit.out.len(), 3);
assert_eq!(emit.out[1].items(), &[3, 4]);
assert_eq!(emit.out[2].items(), &[4]);
}
#[test]
fn tumbling_flush_emits_partial() {
let t0 = Instant::now();
let clock = FakeClock::new(t0);
let mut stage = WindowStage::<u32, _>::new(
WindowPolicy::Tumbling {
size: Duration::from_secs(10),
},
clock.clone(),
);
let mut emit = Collect::<u32> { out: Vec::new() };
stage.process(1, &mut emit).unwrap();
stage.process(2, &mut emit).unwrap();
stage.flush(&mut emit).unwrap();
assert_eq!(emit.out.len(), 1);
assert_eq!(emit.out[0].items(), &[1, 2]);
}
#[test]
fn window_into_inner_returns_items() {
let t = Instant::now();
let w = Window::new(alloc::vec![10u32, 20, 30], t, t);
assert_eq!(w.len(), 3);
assert!(!w.is_empty());
assert_eq!(w.into_inner(), alloc::vec![10, 20, 30]);
}
}