use super::types::Window;
use chrono::{DateTime, Duration, Utc};
use std::collections::HashMap;
use std::fmt;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TriggerResult {
Continue,
Fire,
FireAndPurge,
Purge,
}
#[derive(Debug, Clone)]
pub struct TriggerContext {
pub window: Window,
pub event_time: Option<DateTime<Utc>>,
pub processing_time: DateTime<Utc>,
pub watermark: Option<DateTime<Utc>>,
pub element_count: usize,
}
pub trait WindowTrigger: Send + Sync + fmt::Debug {
fn on_element(&mut self, ctx: &TriggerContext) -> TriggerResult {
let _ = ctx;
TriggerResult::Continue
}
fn on_watermark(&mut self, ctx: &TriggerContext) -> TriggerResult {
let _ = ctx;
TriggerResult::Continue
}
fn on_processing_time(&mut self, ctx: &TriggerContext) -> TriggerResult {
let _ = ctx;
TriggerResult::Continue
}
fn clear(&mut self, window: &Window) {
let _ = window;
}
}
#[derive(Debug, Clone)]
pub struct OnWatermarkTrigger;
impl OnWatermarkTrigger {
pub fn new() -> Self {
Self
}
}
impl Default for OnWatermarkTrigger {
fn default() -> Self {
Self::new()
}
}
impl WindowTrigger for OnWatermarkTrigger {
fn on_watermark(&mut self, ctx: &TriggerContext) -> TriggerResult {
if let Some(watermark) = ctx.watermark {
if watermark >= ctx.window.bounds.end {
return TriggerResult::Fire;
}
}
TriggerResult::Continue
}
}
#[derive(Debug, Clone)]
pub struct ProcessingTimeTrigger {
delay: Duration,
}
impl ProcessingTimeTrigger {
pub fn new() -> Self {
Self {
delay: Duration::zero(),
}
}
pub fn with_delay(delay: Duration) -> Self {
Self { delay }
}
}
impl Default for ProcessingTimeTrigger {
fn default() -> Self {
Self::new()
}
}
impl WindowTrigger for ProcessingTimeTrigger {
fn on_element(&mut self, ctx: &TriggerContext) -> TriggerResult {
let fire_time = ctx.window.bounds.end + self.delay;
if ctx.processing_time >= fire_time {
TriggerResult::Fire
} else {
TriggerResult::Continue
}
}
fn on_processing_time(&mut self, ctx: &TriggerContext) -> TriggerResult {
let fire_time = ctx.window.bounds.end + self.delay;
if ctx.processing_time >= fire_time {
TriggerResult::Fire
} else {
TriggerResult::Continue
}
}
}
#[derive(Debug)]
pub struct CountTrigger {
threshold: usize,
counts: HashMap<String, usize>,
}
impl CountTrigger {
pub fn new(threshold: usize) -> Self {
assert!(threshold > 0, "Count threshold must be positive");
Self {
threshold,
counts: HashMap::new(),
}
}
fn increment_count(&mut self, window_id: &str) -> usize {
let count = self.counts.entry(window_id.to_string()).or_insert(0);
*count += 1;
*count
}
}
impl WindowTrigger for CountTrigger {
fn on_element(&mut self, ctx: &TriggerContext) -> TriggerResult {
let count = self.increment_count(&ctx.window.id);
if count >= self.threshold {
TriggerResult::Fire
} else {
TriggerResult::Continue
}
}
fn clear(&mut self, window: &Window) {
self.counts.remove(&window.id);
}
}
#[derive(Debug)]
pub struct CompositeTrigger {
triggers: Vec<Box<dyn WindowTrigger>>,
mode: CompositeMode,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompositeMode {
Any,
All,
}
impl CompositeTrigger {
pub fn any(triggers: Vec<Box<dyn WindowTrigger>>) -> Self {
assert!(!triggers.is_empty(), "Composite trigger must have at least one trigger");
Self {
triggers,
mode: CompositeMode::Any,
}
}
pub fn all(triggers: Vec<Box<dyn WindowTrigger>>) -> Self {
assert!(!triggers.is_empty(), "Composite trigger must have at least one trigger");
Self {
triggers,
mode: CompositeMode::All,
}
}
fn evaluate<F>(&mut self, ctx: &TriggerContext, f: F) -> TriggerResult
where
F: Fn(&mut Box<dyn WindowTrigger>, &TriggerContext) -> TriggerResult,
{
let mut results = Vec::new();
for trigger in &mut self.triggers {
let result = f(trigger, ctx);
results.push(result);
}
match self.mode {
CompositeMode::Any => {
if results.iter().any(|r| matches!(r, TriggerResult::Fire | TriggerResult::FireAndPurge)) {
TriggerResult::Fire
} else {
TriggerResult::Continue
}
}
CompositeMode::All => {
if results.iter().all(|r| matches!(r, TriggerResult::Fire | TriggerResult::FireAndPurge)) {
TriggerResult::Fire
} else {
TriggerResult::Continue
}
}
}
}
}
impl WindowTrigger for CompositeTrigger {
fn on_element(&mut self, ctx: &TriggerContext) -> TriggerResult {
self.evaluate(ctx, |trigger, ctx| trigger.on_element(ctx))
}
fn on_watermark(&mut self, ctx: &TriggerContext) -> TriggerResult {
self.evaluate(ctx, |trigger, ctx| trigger.on_watermark(ctx))
}
fn on_processing_time(&mut self, ctx: &TriggerContext) -> TriggerResult {
self.evaluate(ctx, |trigger, ctx| trigger.on_processing_time(ctx))
}
fn clear(&mut self, window: &Window) {
for trigger in &mut self.triggers {
trigger.clear(window);
}
}
}
#[derive(Debug, Clone)]
pub struct ImmediateTrigger;
impl ImmediateTrigger {
pub fn new() -> Self {
Self
}
}
impl Default for ImmediateTrigger {
fn default() -> Self {
Self::new()
}
}
impl WindowTrigger for ImmediateTrigger {
fn on_element(&mut self, _ctx: &TriggerContext) -> TriggerResult {
TriggerResult::Fire
}
}
#[derive(Debug, Clone)]
pub struct NeverTrigger;
impl NeverTrigger {
pub fn new() -> Self {
Self
}
}
impl Default for NeverTrigger {
fn default() -> Self {
Self::new()
}
}
impl WindowTrigger for NeverTrigger {
}
#[cfg(test)]
mod tests {
use super::*;
use crate::window::types::WindowBounds;
use chrono::TimeZone;
fn create_timestamp(millis: i64) -> DateTime<Utc> {
Utc.timestamp_millis_opt(millis).unwrap()
}
fn create_window(start_millis: i64, end_millis: i64) -> Window {
Window::new(WindowBounds::new(
create_timestamp(start_millis),
create_timestamp(end_millis),
))
}
fn create_context(
window: Window,
watermark: Option<DateTime<Utc>>,
processing_time: DateTime<Utc>,
element_count: usize,
) -> TriggerContext {
TriggerContext {
window,
event_time: None,
processing_time,
watermark,
element_count,
}
}
#[test]
fn test_watermark_trigger() {
let mut trigger = OnWatermarkTrigger::new();
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
Some(create_timestamp(1500)),
create_timestamp(3000),
5,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
Some(create_timestamp(2000)),
create_timestamp(3000),
5,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Fire);
let ctx = create_context(
window.clone(),
Some(create_timestamp(2500)),
create_timestamp(3000),
5,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Fire);
let ctx = create_context(
window,
None,
create_timestamp(3000),
5,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Continue);
}
#[test]
fn test_processing_time_trigger() {
let mut trigger = ProcessingTimeTrigger::new();
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1500),
5,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
None,
create_timestamp(2000),
5,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
let ctx = create_context(
window,
None,
create_timestamp(2500),
5,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
}
#[test]
fn test_processing_time_trigger_with_delay() {
let mut trigger = ProcessingTimeTrigger::with_delay(Duration::milliseconds(1000));
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
None,
create_timestamp(2500),
5,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
None,
create_timestamp(3000),
5,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
}
#[test]
fn test_count_trigger() {
let mut trigger = CountTrigger::new(3);
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1500),
1,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1600),
2,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1700),
3,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1800),
4,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
trigger.clear(&window);
let ctx = create_context(
window.clone(),
None,
create_timestamp(1900),
1,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
}
#[test]
fn test_composite_trigger_any() {
let count_trigger: Box<dyn WindowTrigger> = Box::new(CountTrigger::new(5));
let watermark_trigger: Box<dyn WindowTrigger> = Box::new(OnWatermarkTrigger::new());
let mut trigger = CompositeTrigger::any(vec![count_trigger, watermark_trigger]);
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
Some(create_timestamp(1500)),
create_timestamp(1500),
1,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
Some(create_timestamp(2500)),
create_timestamp(1500),
2,
);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Fire);
}
#[test]
fn test_composite_trigger_all() {
let trigger1: Box<dyn WindowTrigger> = Box::new(ProcessingTimeTrigger::new());
let trigger2: Box<dyn WindowTrigger> = Box::new(ProcessingTimeTrigger::with_delay(Duration::milliseconds(500)));
let mut trigger = CompositeTrigger::all(vec![trigger1, trigger2]);
let window = create_window(1000, 2000);
let ctx = create_context(
window.clone(),
None,
create_timestamp(2000),
1,
);
assert_eq!(trigger.on_processing_time(&ctx), TriggerResult::Continue);
let ctx = create_context(
window.clone(),
None,
create_timestamp(2500),
1,
);
assert_eq!(trigger.on_processing_time(&ctx), TriggerResult::Fire);
}
#[test]
fn test_immediate_trigger() {
let mut trigger = ImmediateTrigger::new();
let window = create_window(1000, 2000);
let ctx = create_context(
window,
None,
create_timestamp(1500),
1,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Fire);
}
#[test]
fn test_never_trigger() {
let mut trigger = NeverTrigger::new();
let window = create_window(1000, 2000);
let ctx = create_context(
window,
Some(create_timestamp(5000)),
create_timestamp(5000),
100,
);
assert_eq!(trigger.on_element(&ctx), TriggerResult::Continue);
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Continue);
assert_eq!(trigger.on_processing_time(&ctx), TriggerResult::Continue);
}
#[test]
#[should_panic(expected = "Count threshold must be positive")]
fn test_count_trigger_invalid_threshold() {
CountTrigger::new(0);
}
#[test]
#[should_panic(expected = "Composite trigger must have at least one trigger")]
fn test_composite_trigger_empty() {
CompositeTrigger::any(vec![]);
}
}