oxirs_stream/watermark/

mod.rs

//! # Stream Watermarking
//!
//! Event-time watermarks for out-of-order stream processing.
//!
//! ## Key Types
//!
//! - [`StreamWatermark`]: An immutable watermark value with source attribution
//! - [`WatermarkGenerator`]: Advances a watermark as events are observed
//! - [`WatermarkAligner`]: Computes the global watermark across multiple sources (minimum)
//! - [`LateDataHandler`]: Decides what to do with events that arrive after the watermark

use std::collections::HashMap;

// ─── StreamWatermark ──────────────────────────────────────────────────────────

/// An event-time watermark generated by a single source.
///
/// The watermark asserts that all future events from `source_id` will have
/// timestamps `>= timestamp`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamWatermark {
    /// Milliseconds since Unix epoch.
    pub timestamp: i64,
    /// Identifier of the source that produced this watermark.
    pub source_id: String,
    /// Number of times this generator has advanced the watermark.
    pub advance_count: u64,
}

impl StreamWatermark {
    pub fn new(timestamp: i64, source_id: impl Into<String>, advance_count: u64) -> Self {
        Self {
            timestamp,
            source_id: source_id.into(),
            advance_count,
        }
    }
}

// ─── WatermarkGenerator ───────────────────────────────────────────────────────

/// Generates watermarks by tracking the maximum observed event timestamp.
///
/// A watermark is emitted at `max_event_ts - max_out_of_order_ms` after every
/// `advance_threshold` events. This provides a balance between latency and
/// tolerance for out-of-order events.
pub struct WatermarkGenerator {
    max_out_of_order_ms: i64,
    current_max_ts: i64,
    advance_threshold: usize,
    event_count: usize,
    advance_count: u64,
    source_id: String,
}

impl WatermarkGenerator {
    /// Create a generator that tolerates up to `max_out_of_order_ms` of lateness.
    ///
    /// Watermarks are emitted every 100 events by default.
    pub fn new(max_out_of_order_ms: i64) -> Self {
        Self {
            max_out_of_order_ms,
            current_max_ts: i64::MIN,
            advance_threshold: 100,
            event_count: 0,
            advance_count: 0,
            source_id: "default".to_string(),
        }
    }

    /// Set the event count after which a watermark is emitted.
    pub fn with_advance_threshold(mut self, threshold: usize) -> Self {
        self.advance_threshold = threshold;
        self
    }

    /// Set the source identifier embedded in emitted watermarks.
    pub fn with_source_id(mut self, source_id: impl Into<String>) -> Self {
        self.source_id = source_id.into();
        self
    }

    /// Observe an event with `event_timestamp_ms`.
    ///
    /// Returns `Some(watermark)` when the threshold is reached, otherwise `None`.
    pub fn observe(&mut self, event_timestamp_ms: i64) -> Option<StreamWatermark> {
        if event_timestamp_ms > self.current_max_ts {
            self.current_max_ts = event_timestamp_ms;
        }
        self.event_count += 1;

        if self.event_count >= self.advance_threshold {
            self.event_count = 0;
            self.advance_count += 1;
            Some(StreamWatermark::new(
                self.current_watermark(),
                self.source_id.clone(),
                self.advance_count,
            ))
        } else {
            None
        }
    }

    /// Current watermark timestamp: `max_event_ts - max_out_of_order_ms`.
    ///
    /// Returns `i64::MIN` before any event has been observed.
    pub fn current_watermark(&self) -> i64 {
        if self.current_max_ts == i64::MIN {
            i64::MIN
        } else {
            self.current_max_ts - self.max_out_of_order_ms
        }
    }

    /// Return `true` if `event_timestamp_ms` is before the current watermark.
    pub fn is_late(&self, event_timestamp_ms: i64) -> bool {
        event_timestamp_ms < self.current_watermark()
    }
}

// ─── WatermarkAligner ─────────────────────────────────────────────────────────

/// Tracks per-source watermarks and computes a global minimum watermark.
///
/// A pipeline must wait for *all* sources to advance before it can safely
/// process events up to a given timestamp.
pub struct WatermarkAligner {
    sources: HashMap<String, i64>,
}

impl WatermarkAligner {
    /// Create an empty aligner with no tracked sources.
    pub fn new() -> Self {
        Self {
            sources: HashMap::new(),
        }
    }

    /// Update (or register) the watermark for `source_id`.
    pub fn update(&mut self, source_id: &str, watermark_ms: i64) {
        self.sources.insert(source_id.to_string(), watermark_ms);
    }

    /// The global watermark: the minimum across all tracked sources.
    ///
    /// Returns `i64::MIN` when no sources are tracked.
    pub fn global_watermark(&self) -> i64 {
        self.sources.values().copied().min().unwrap_or(i64::MIN)
    }

    /// Number of sources currently tracked.
    pub fn source_count(&self) -> usize {
        self.sources.len()
    }

    /// Return `true` iff *all* tracked sources have watermarks strictly
    /// greater than `timestamp_ms`.
    pub fn all_beyond(&self, timestamp_ms: i64) -> bool {
        if self.sources.is_empty() {
            return false;
        }
        self.sources.values().all(|&wm| wm > timestamp_ms)
    }
}

impl Default for WatermarkAligner {
    fn default() -> Self {
        Self::new()
    }
}

// ─── LateDataPolicy / LateDataDecision ────────────────────────────────────────

/// What the pipeline should do with events that arrive after the watermark.
#[derive(Debug, Clone)]
pub enum LateDataPolicy {
    /// Discard the event.
    Drop,
    /// Re-assign the event to the nearest open window that can still accept it,
    /// provided the event is at most `max_lateness_ms` behind the watermark.
    Reassign { max_lateness_ms: i64 },
    /// Route the event to a named side-output channel for separate handling.
    SideOutput { channel: String },
}

/// Result returned by [`LateDataHandler::handle`].
#[derive(Debug, Clone, PartialEq)]
pub enum LateDataDecision {
    /// Process the event in the main stream.
    Process,
    /// Drop the event.
    Drop,
    /// Re-assign the event to the given timestamp.
    Reassign(i64),
    /// Route the event to a side-output channel.
    SideOutput,
}

/// Applies the configured [`LateDataPolicy`] to late events.
pub struct LateDataHandler {
    pub policy: LateDataPolicy,
    pub late_event_count: u64,
}

impl LateDataHandler {
    /// Create a handler with the given policy.
    pub fn new(policy: LateDataPolicy) -> Self {
        Self {
            policy,
            late_event_count: 0,
        }
    }

    /// Decide what to do with an event at `event_ts_ms` given that the
    /// current watermark is `watermark_ms`.
    ///
    /// If `event_ts_ms >= watermark_ms` the event is not late and `Process` is returned.
    pub fn handle(&mut self, event_ts_ms: i64, watermark_ms: i64) -> LateDataDecision {
        if event_ts_ms >= watermark_ms {
            return LateDataDecision::Process;
        }
        self.late_event_count += 1;
        match &self.policy {
            LateDataPolicy::Drop => LateDataDecision::Drop,
            LateDataPolicy::Reassign { max_lateness_ms } => {
                let lateness = watermark_ms - event_ts_ms;
                if lateness <= *max_lateness_ms {
                    LateDataDecision::Reassign(watermark_ms)
                } else {
                    LateDataDecision::Drop
                }
            }
            LateDataPolicy::SideOutput { .. } => LateDataDecision::SideOutput,
        }
    }
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // ── WatermarkGenerator ─────────────────────────────────────────────────────

    #[test]
    fn test_generator_initial_watermark_is_min() {
        let gen = WatermarkGenerator::new(500);
        assert_eq!(gen.current_watermark(), i64::MIN);
    }

    #[test]
    fn test_generator_no_watermark_before_threshold() {
        let mut gen = WatermarkGenerator::new(100).with_advance_threshold(10);
        for ts in 0..9 {
            let wm = gen.observe(ts * 100);
            assert!(wm.is_none(), "should not emit before threshold");
        }
    }

    #[test]
    fn test_generator_emits_at_threshold() {
        let mut gen = WatermarkGenerator::new(100).with_advance_threshold(5);
        for ts in 0..5 {
            let _wm = gen.observe(ts * 1000);
        }
        let wm = gen.observe(5000);
        // threshold reached at observe #5 (indices 0–4 fill 5 slots, #5 is #6 call)
        // let's just call 5 times to trigger
        let mut gen2 = WatermarkGenerator::new(100).with_advance_threshold(5);
        let mut last_wm = None;
        for i in 0..5 {
            last_wm = gen2.observe(i * 1000);
        }
        assert!(last_wm.is_some(), "watermark must be emitted at threshold");
        let w = last_wm.unwrap();
        // max ts = 4000, watermark = 4000 - 100 = 3900
        assert_eq!(w.timestamp, 3900);
        assert_eq!(w.advance_count, 1);
        // suppress unused variable warning
        drop(wm);
    }

    #[test]
    fn test_generator_advance_count_increments() {
        let mut gen = WatermarkGenerator::new(0).with_advance_threshold(3);
        for i in 0..6 {
            gen.observe(i * 1000);
        }
        assert_eq!(gen.advance_count, 2);
    }

    #[test]
    fn test_generator_current_watermark_formula() {
        let mut gen = WatermarkGenerator::new(200).with_advance_threshold(100);
        gen.observe(5000);
        assert_eq!(gen.current_watermark(), 4800); // 5000 - 200
    }

    #[test]
    fn test_generator_is_late_true() {
        let mut gen = WatermarkGenerator::new(100).with_advance_threshold(100);
        gen.observe(10_000);
        // watermark = 9900; event at 9000 is late
        assert!(gen.is_late(9000));
    }

    #[test]
    fn test_generator_is_late_false() {
        let mut gen = WatermarkGenerator::new(100).with_advance_threshold(100);
        gen.observe(10_000);
        // watermark = 9900; event at 9900 is not late (== watermark)
        assert!(!gen.is_late(9900));
        // event at 10_000 is future
        assert!(!gen.is_late(10_000));
    }

    #[test]
    fn test_generator_source_id() {
        let mut gen = WatermarkGenerator::new(0)
            .with_advance_threshold(1)
            .with_source_id("sensor-42");
        let wm = gen.observe(1000).expect("should emit");
        assert_eq!(wm.source_id, "sensor-42");
    }

    #[test]
    fn test_generator_max_ts_tracks_maximum() {
        let mut gen = WatermarkGenerator::new(0).with_advance_threshold(100);
        // Observe out-of-order events
        gen.observe(500);
        gen.observe(1000);
        gen.observe(300);
        assert_eq!(gen.current_max_ts, 1000);
        assert_eq!(gen.current_watermark(), 1000);
    }

    // ── WatermarkAligner ───────────────────────────────────────────────────────

    #[test]
    fn test_aligner_empty_returns_min() {
        let aligner = WatermarkAligner::new();
        assert_eq!(aligner.global_watermark(), i64::MIN);
        assert_eq!(aligner.source_count(), 0);
    }

    #[test]
    fn test_aligner_single_source() {
        let mut aligner = WatermarkAligner::new();
        aligner.update("src-A", 5000);
        assert_eq!(aligner.global_watermark(), 5000);
        assert_eq!(aligner.source_count(), 1);
    }

    #[test]
    fn test_aligner_global_is_minimum() {
        let mut aligner = WatermarkAligner::new();
        aligner.update("src-A", 5000);
        aligner.update("src-B", 3000);
        aligner.update("src-C", 7000);
        assert_eq!(aligner.global_watermark(), 3000);
    }

    #[test]
    fn test_aligner_all_beyond_true() {
        let mut aligner = WatermarkAligner::new();
        aligner.update("A", 10_000);
        aligner.update("B", 12_000);
        assert!(aligner.all_beyond(9_999));
    }

    #[test]
    fn test_aligner_all_beyond_false_one_lagging() {
        let mut aligner = WatermarkAligner::new();
        aligner.update("A", 10_000);
        aligner.update("B", 5_000);
        assert!(!aligner.all_beyond(6_000));
    }

    #[test]
    fn test_aligner_all_beyond_empty_is_false() {
        let aligner = WatermarkAligner::new();
        assert!(!aligner.all_beyond(0));
    }

    #[test]
    fn test_aligner_update_overwrites() {
        let mut aligner = WatermarkAligner::new();
        aligner.update("src", 1000);
        aligner.update("src", 9000);
        assert_eq!(aligner.global_watermark(), 9000);
        assert_eq!(aligner.source_count(), 1);
    }

    // ── LateDataHandler ────────────────────────────────────────────────────────

    #[test]
    fn test_late_handler_on_time_event_is_process() {
        let mut handler = LateDataHandler::new(LateDataPolicy::Drop);
        let decision = handler.handle(5000, 4000);
        assert_eq!(decision, LateDataDecision::Process);
        assert_eq!(handler.late_event_count, 0);
    }

    #[test]
    fn test_late_handler_drop_policy() {
        let mut handler = LateDataHandler::new(LateDataPolicy::Drop);
        let decision = handler.handle(1000, 5000);
        assert_eq!(decision, LateDataDecision::Drop);
        assert_eq!(handler.late_event_count, 1);
    }

    #[test]
    fn test_late_handler_reassign_within_max_lateness() {
        let mut handler = LateDataHandler::new(LateDataPolicy::Reassign {
            max_lateness_ms: 2000,
        });
        // event at 4000, watermark at 5000 → lateness = 1000 ≤ 2000 → Reassign
        let decision = handler.handle(4000, 5000);
        assert_eq!(decision, LateDataDecision::Reassign(5000));
    }

    #[test]
    fn test_late_handler_reassign_exceeds_max_lateness_drops() {
        let mut handler = LateDataHandler::new(LateDataPolicy::Reassign {
            max_lateness_ms: 500,
        });
        // event at 3000, watermark at 5000 → lateness = 2000 > 500 → Drop
        let decision = handler.handle(3000, 5000);
        assert_eq!(decision, LateDataDecision::Drop);
        assert_eq!(handler.late_event_count, 1);
    }

    #[test]
    fn test_late_handler_side_output_policy() {
        let mut handler = LateDataHandler::new(LateDataPolicy::SideOutput {
            channel: "late-events".to_string(),
        });
        let decision = handler.handle(1000, 5000);
        assert_eq!(decision, LateDataDecision::SideOutput);
        assert_eq!(handler.late_event_count, 1);
    }

    #[test]
    fn test_late_handler_counts_accumulate() {
        let mut handler = LateDataHandler::new(LateDataPolicy::Drop);
        handler.handle(100, 1000);
        handler.handle(200, 1000);
        handler.handle(5000, 1000); // on-time: watermark = 1000, event at 5000
        assert_eq!(handler.late_event_count, 2);
    }

    #[test]
    fn test_stream_watermark_equality() {
        let w1 = StreamWatermark::new(1000, "s1", 1);
        let w2 = StreamWatermark::new(1000, "s1", 1);
        let w3 = StreamWatermark::new(2000, "s1", 1);
        assert_eq!(w1, w2);
        assert_ne!(w1, w3);
    }
}