pulse_core/

lib.rs

//! pulse-core: Fundamental types, traits and a basic executor.
//!
//! Goal: provide the essential abstractions for a streaming pipeline:
//! - Record, EventTime, Watermark
//! - Traits: Source, Operator, Sink, Context, KvState, Timers
//! - Simple tokio-based executor
//!
//! Quick example:
//! ```no_run
//! use async_trait::async_trait;
//! use pulse_core::prelude::*;
//!
//! struct MySource;
//! #[async_trait]
//! impl Source for MySource {
//!     async fn run(&mut self, ctx: &mut dyn Context) -> Result<()> {
//!         ctx.collect(Record::from_value("hello"));
//!         Ok(())
//!     }
//! }
//!
//! struct MyOp;
//! #[async_trait]
//! impl Operator for MyOp {
//!     async fn on_element(&mut self, ctx: &mut dyn Context, rec: Record) -> Result<()> {
//!         // pass-through
//!         ctx.collect(rec);
//!         Ok(())
//!     }
//! }
//!
//! struct MySink;
//! #[async_trait]
//! impl Sink for MySink {
//!     async fn on_element(&mut self, rec: Record) -> Result<()> {
//!         println!("{}", rec.value);
//!         Ok(())
//!     }
//! }
//!
//! #[tokio::main]
//! async fn main() -> Result<()> {
//!     let mut exec = Executor::new();
//!     exec.source(MySource).operator(MyOp).sink(MySink);
//!     exec.run().await?;
//!     Ok(())
//! }
//! ```

use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;

use chrono::{DateTime, Utc};
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};

pub mod checkpoint;
pub use checkpoint::{CheckpointMeta, SnapshotId};

pub mod record;
pub use record::Record;
pub mod config;
pub mod metrics;

/// Logical event-time as wall-clock timestamp in UTC.
/// Use [`EventTime::now`] for current time.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
pub struct EventTime(pub DateTime<Utc>);

impl EventTime {
    pub fn now() -> Self {
        EventTime(Utc::now())
    }
}

// Record is defined in `record` module and re-exported above.

/// A low-watermark indicating no future records <= this event-time are expected.
#[derive(Debug, Clone, Copy)]
pub struct Watermark(pub EventTime);

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error(transparent)]
    Anyhow(#[from] anyhow::Error),
    #[error(transparent)]
    Io(#[from] std::io::Error),
    #[error(transparent)]
    Json(#[from] serde_json::Error),
    #[error(transparent)]
    Csv(#[from] csv::Error),
}

pub type Result<T> = std::result::Result<T, Error>;

/// Key-Value state abstraction for stateful operators.
#[async_trait::async_trait]
pub trait KvState: Send + Sync {
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
    async fn put(&self, key: &[u8], value: Vec<u8>) -> Result<()>;
    async fn delete(&self, key: &[u8]) -> Result<()>;
    /// Returns all key/value pairs, optionally filtered by a prefix.
    async fn iter_prefix(&self, prefix: Option<&[u8]>) -> Result<Vec<(Vec<u8>, Vec<u8>)>>;
    /// Creates a point-in-time snapshot of the current state and returns its id.
    async fn snapshot(&self) -> Result<SnapshotId>;
    /// Restores the state to a previously created snapshot id.
    async fn restore(&self, snapshot: SnapshotId) -> Result<()>;
}

/// Timer service for event-time callbacks requested by operators.
#[async_trait::async_trait]
pub trait Timers: Send + Sync {
    async fn register_event_time_timer(&self, when: EventTime, key: Option<Vec<u8>>) -> Result<()>;
}

/// Execution context visible to Sources and Operators.
#[async_trait::async_trait]
pub trait Context: Send {
    fn collect(&mut self, record: Record);
    fn watermark(&mut self, wm: Watermark);
    fn kv(&self) -> Arc<dyn KvState>;
    fn timers(&self) -> Arc<dyn Timers>;
}

/// A data source that pushes records into the pipeline.
#[async_trait::async_trait]
pub trait Source: Send {
    async fn run(&mut self, ctx: &mut dyn Context) -> Result<()>;
}

// Allow using boxed trait objects as sources seamlessly.
#[async_trait::async_trait]
impl<T> Source for Box<T>
where
    T: Source + ?Sized,
{
    async fn run(&mut self, ctx: &mut dyn Context) -> Result<()> {
        (**self).run(ctx).await
    }
}

/// Core operator interface. Override `on_watermark`/`on_timer` if needed.
#[async_trait::async_trait]
pub trait Operator: Send {
    async fn on_element(&mut self, ctx: &mut dyn Context, record: Record) -> Result<()>;
    async fn on_watermark(&mut self, _ctx: &mut dyn Context, _wm: Watermark) -> Result<()> {
        Ok(())
    }
    async fn on_timer(
        &mut self,
        _ctx: &mut dyn Context,
        _when: EventTime,
        _key: Option<Vec<u8>>,
    ) -> Result<()> {
        Ok(())
    }
}

/// A terminal sink that receives records (and optional watermarks).
#[async_trait::async_trait]
pub trait Sink: Send {
    async fn on_element(&mut self, record: Record) -> Result<()>;
    async fn on_watermark(&mut self, _wm: Watermark) -> Result<()> {
        Ok(())
    }
}

#[derive(Default)]
struct SimpleStateInner {
    map: std::collections::HashMap<Vec<u8>, Vec<u8>>,
    snapshots: std::collections::HashMap<SnapshotId, std::collections::HashMap<Vec<u8>, Vec<u8>>>,
}

#[derive(Clone)]
pub struct SimpleInMemoryState(Arc<Mutex<SimpleStateInner>>);

impl Default for SimpleInMemoryState {
    fn default() -> Self {
        Self(Arc::new(Mutex::new(SimpleStateInner {
            map: Default::default(),
            snapshots: Default::default(),
        })))
    }
}

#[async_trait::async_trait]
impl KvState for SimpleInMemoryState {
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        Ok(self.0.lock().map.get(key).cloned())
    }
    async fn put(&self, key: &[u8], value: Vec<u8>) -> Result<()> {
        self.0.lock().map.insert(key.to_vec(), value);
        let sz = self.0.lock().map.len() as i64;
        metrics::STATE_SIZE
            .with_label_values(&["SimpleInMemoryState"])
            .set(sz);
        Ok(())
    }
    async fn delete(&self, key: &[u8]) -> Result<()> {
        self.0.lock().map.remove(key);
        let sz = self.0.lock().map.len() as i64;
        metrics::STATE_SIZE
            .with_label_values(&["SimpleInMemoryState"])
            .set(sz);
        Ok(())
    }
    async fn iter_prefix(&self, prefix: Option<&[u8]>) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let guard = self.0.lock();
        let mut v: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        if let Some(p) = prefix {
            for (k, val) in guard.map.iter() {
                if k.starts_with(p) {
                    v.push((k.clone(), val.clone()));
                }
            }
        } else {
            v.extend(guard.map.iter().map(|(k, val)| (k.clone(), val.clone())));
        }
        Ok(v)
    }
    async fn snapshot(&self) -> Result<SnapshotId> {
        use std::time::{SystemTime, UNIX_EPOCH};
        let mut guard = self.0.lock();
        let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis();
        let id: SnapshotId = format!("mem-{}", ts);
        let current = guard.map.clone();
        guard.snapshots.insert(id.clone(), current);
        Ok(id)
    }
    async fn restore(&self, snapshot: SnapshotId) -> Result<()> {
        let mut guard = self.0.lock();
        if let Some(m) = guard.snapshots.get(&snapshot) {
            guard.map = m.clone();
            Ok(())
        } else {
            // No-op if snapshot not found
            Ok(())
        }
    }
}

/// Minimal in-memory timer service used by the demo executor.
#[derive(Clone, Default)]
pub struct SimpleTimers;

#[async_trait::async_trait]
impl Timers for SimpleTimers {
    async fn register_event_time_timer(&self, _when: EventTime, _key: Option<Vec<u8>>) -> Result<()> {
        // No-op basic impl; real executors would drive timer callbacks.
        Ok(())
    }
}

/// A tiny, single-pipeline executor.
/// Wires: Source -> Operators -> Sink. Drives watermarks and event-time timers.
pub struct Executor {
    source: Option<Box<dyn Source>>,
    operators: Vec<Box<dyn Operator>>,
    sink: Option<Box<dyn Sink>>,
    kv: Arc<dyn KvState>,
    timers: Arc<dyn Timers>,
}

impl Executor {
    /// Create a new empty executor.
    pub fn new() -> Self {
        Self {
            source: None,
            operators: Vec::new(),
            sink: None,
            kv: Arc::new(SimpleInMemoryState::default()),
            timers: Arc::new(SimpleTimers::default()),
        }
    }

    /// Set the pipeline source.
    pub fn source<S: Source + 'static>(&mut self, s: S) -> &mut Self {
        self.source = Some(Box::new(s));
        self
    }

    /// Append an operator to the pipeline.
    pub fn operator<O: Operator + 'static>(&mut self, o: O) -> &mut Self {
        self.operators.push(Box::new(o));
        self
    }

    /// Set the pipeline sink.
    pub fn sink<K: Sink + 'static>(&mut self, s: K) -> &mut Self {
        self.sink = Some(Box::new(s));
        self
    }

    /// Run the pipeline to completion. The loop exits when the source finishes
    /// and the internal channel closes. Watermarks are propagated and due timers fired.
    pub async fn run(&mut self) -> Result<()> {
        let kv = self.kv.clone();
        let timers = self.timers.clone();

        // Shared timer queue used to schedule per-operator event-time timers
        #[derive(Clone)]
        struct TimerEntry {
            op_idx: usize,
            when: EventTime,
            key: Option<Vec<u8>>,
        }
        #[derive(Clone, Default)]
        struct SharedTimers(Arc<Mutex<Vec<TimerEntry>>>);
        impl SharedTimers {
            fn add(&self, op_idx: usize, when: EventTime, key: Option<Vec<u8>>) {
                self.0.lock().push(TimerEntry { op_idx, when, key });
            }
            fn drain_due(&self, wm: EventTime) -> Vec<TimerEntry> {
                let mut guard = self.0.lock();
                let mut fired = Vec::new();
                let mut i = 0;
                while i < guard.len() {
                    if guard[i].when.0 <= wm.0 {
                        fired.push(guard.remove(i));
                    } else {
                        i += 1;
                    }
                }
                fired
            }
        }

        enum EventMsg {
            Data(Record),
            Wm(Watermark),
        }
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<EventMsg>();
        // Soft-bound drop policy controlled by env PULSE_CHANNEL_BOUND (>0):
        // If the in-flight depth reaches the bound, new Data records are dropped at source Collect.
        let bound = std::env::var("PULSE_CHANNEL_BOUND")
            .ok()
            .and_then(|s| s.parse::<i64>().ok())
            .unwrap_or(0);
        let depth = Arc::new(AtomicI64::new(0));

        struct ExecCtx {
            tx: tokio::sync::mpsc::UnboundedSender<EventMsg>,
            kv: Arc<dyn KvState>,
            timers: Arc<dyn Timers>,
            bound: i64,
            depth: Arc<AtomicI64>,
        }

        #[async_trait::async_trait]
        impl Context for ExecCtx {
            fn collect(&mut self, record: Record) {
                // Soft-bound: drop data if depth reached bound
                if self.bound > 0 && self.depth.load(Ordering::Relaxed) >= self.bound {
                    metrics::DROPPED_RECORDS
                        .with_label_values(&["channel_full"])
                        .inc();
                    return;
                }
                if self.tx.send(EventMsg::Data(record)).is_ok() {
                    self.depth.fetch_add(1, Ordering::Relaxed);
                    metrics::QUEUE_DEPTH.inc();
                } else {
                    metrics::DROPPED_RECORDS.with_label_values(&["send_failed"]).inc();
                }
            }
            fn watermark(&mut self, wm: Watermark) {
                // Always forward watermarks; never drop
                let _ = self.tx.send(EventMsg::Wm(wm));
            }
            fn kv(&self) -> Arc<dyn KvState> {
                self.kv.clone()
            }
            fn timers(&self) -> Arc<dyn Timers> {
                self.timers.clone()
            }
        }

        let mut source = self.source.take().ok_or_else(|| anyhow::anyhow!("no source"))?;
        let mut ops = std::mem::take(&mut self.operators);
        let mut sink = self.sink.take().ok_or_else(|| anyhow::anyhow!("no sink"))?;

        // Shared timers queue
        let shared_timers = SharedTimers::default();

        // Source task
        let mut sctx = ExecCtx {
            tx: tx.clone(),
            kv: kv.clone(),
            timers: timers.clone(),
            bound,
            depth: depth.clone(),
        };
        let src_handle = tokio::spawn(async move { source.run(&mut sctx).await });
        // Drop our local sender so the channel closes once the source finishes (its clone will drop then)
        drop(tx);

        // Operator chain processing task
        let op_handle = tokio::spawn(async move {
            // Local Timers wrapper capturing operator index
            struct LocalTimers {
                op_idx: usize,
                shared: SharedTimers,
            }
            #[async_trait::async_trait]
            impl Timers for LocalTimers {
                async fn register_event_time_timer(
                    &self,
                    when: EventTime,
                    key: Option<Vec<u8>>,
                ) -> Result<()> {
                    self.shared.add(self.op_idx, when, key);
                    Ok(())
                }
            }

            // Local Context used for operators; collects into a Vec to be forwarded
            struct LocalCtx<'a> {
                out: &'a mut Vec<Record>,
                kv: Arc<dyn KvState>,
                timers: Arc<dyn Timers>,
            }
            #[async_trait::async_trait]
            impl<'a> Context for LocalCtx<'a> {
                fn collect(&mut self, record: Record) {
                    self.out.push(record);
                }
                fn watermark(&mut self, _wm: Watermark) {}
                fn kv(&self) -> Arc<dyn KvState> {
                    self.kv.clone()
                }
                fn timers(&self) -> Arc<dyn Timers> {
                    self.timers.clone()
                }
            }

            while let Some(msg) = rx.recv().await {
                // Adjust queue depth when we pull an item
                depth.fetch_sub(1, Ordering::Relaxed);
                metrics::QUEUE_DEPTH.dec();
                match msg {
                    EventMsg::Data(rec) => {
                        // Pipe record through the operator chain collecting outputs at each step
                        let mut batch = vec![rec];
                        for (i, op) in ops.iter_mut().enumerate() {
                            let mut next = Vec::new();
                            let timers = Arc::new(LocalTimers {
                                op_idx: i,
                                shared: shared_timers.clone(),
                            });
                            for item in batch.drain(..) {
                                let mut lctx = LocalCtx {
                                    out: &mut next,
                                    kv: kv.clone(),
                                    timers: timers.clone(),
                                };
                                let t0 = std::time::Instant::now();
                                op.on_element(&mut lctx, item).await?;
                                let dt = t0.elapsed().as_secs_f64() * 1000.0;
                                metrics::OP_PROC_LATENCY_MS.observe(dt);
                            }
                            batch = next;
                            if batch.is_empty() {
                                break;
                            }
                        }
                        for out in batch.into_iter() {
                            let t0 = std::time::Instant::now();
                            sink.on_element(out).await?;
                            let dt = t0.elapsed().as_secs_f64() * 1000.0;
                            metrics::SINK_PROC_LATENCY_MS.observe(dt);
                        }
                    }
                    EventMsg::Wm(wm) => {
                        // Propagate watermark to operators in order, allowing them to emit
                        let mut emitted = Vec::new();
                        // Update lag metric: now - watermark
                        let now = chrono::Utc::now();
                        let lag = (now - wm.0 .0).num_milliseconds();
                        metrics::LAG_WATERMARK_MS.set(lag as i64);
                        for (i, op) in ops.iter_mut().enumerate() {
                            let timers = Arc::new(LocalTimers {
                                op_idx: i,
                                shared: shared_timers.clone(),
                            });
                            let mut lctx = LocalCtx {
                                out: &mut emitted,
                                kv: kv.clone(),
                                timers: timers.clone(),
                            };
                            op.on_watermark(&mut lctx, wm).await?;
                        }
                        // Fire any timers due at this watermark
                        let due = shared_timers.drain_due(wm.0);
                        for t in due.into_iter() {
                            if let Some(op) = ops.get_mut(t.op_idx) {
                                let timers = Arc::new(LocalTimers {
                                    op_idx: t.op_idx,
                                    shared: shared_timers.clone(),
                                });
                                let mut lctx = LocalCtx {
                                    out: &mut emitted,
                                    kv: kv.clone(),
                                    timers: timers.clone(),
                                };
                                op.on_timer(&mut lctx, t.when, t.key.clone()).await?;
                            }
                        }
                        // Emit produced records to sink
                        for out in emitted.into_iter() {
                            sink.on_element(out).await?;
                        }
                        // Inform sink about watermark
                        sink.on_watermark(wm).await?;
                    }
                }
            }
            Ok::<_, Error>(())
        });

        // Await tasks
        src_handle
            .await
            .map_err(|e| Error::Anyhow(anyhow::anyhow!(e)))??;
        op_handle.await.map_err(|e| Error::Anyhow(anyhow::anyhow!(e)))??;
        Ok(())
    }
}

pub mod prelude {
    pub use super::{
        CheckpointMeta, Context, EventTime, Executor, KvState, Operator, Record, Result, Sink, SnapshotId,
        Source, Watermark,
    };
}

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

    struct TestSource;
    #[async_trait::async_trait]
    impl Source for TestSource {
        async fn run(&mut self, ctx: &mut dyn Context) -> Result<()> {
            ctx.collect(Record::from_value(serde_json::json!({"n":1})));
            Ok(())
        }
    }

    struct TestOp;
    #[async_trait::async_trait]
    impl Operator for TestOp {
        async fn on_element(&mut self, ctx: &mut dyn Context, mut record: Record) -> Result<()> {
            record.value["n"] = serde_json::json!(record.value["n"].as_i64().unwrap() + 1);
            ctx.collect(record);
            Ok(())
        }
    }

    struct TestSink(pub std::sync::Arc<std::sync::Mutex<Vec<serde_json::Value>>>);
    #[async_trait::async_trait]
    impl Sink for TestSink {
        async fn on_element(&mut self, record: Record) -> Result<()> {
            self.0.lock().unwrap().push(record.value);
            Ok(())
        }
    }

    #[tokio::test]
    async fn executor_wires_stages() {
        let mut exec = Executor::new();
        let out = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
        exec.source(TestSource)
            .operator(TestOp)
            .sink(TestSink(out.clone()));
        exec.run().await.unwrap();
        let got = out.lock().unwrap().clone();
        assert_eq!(got.len(), 1);
        assert_eq!(got[0]["n"], serde_json::json!(2));
    }
}