datum-core 0.9.2

//! Opt-in stream instrumentation for operational tooling.
//!
//! Instrumentation is explicit: ordinary stream blueprints do not carry a registry,
//! branch, timestamp call, or atomic counter. Calling [`Source::instrumented`] inserts
//! a small counting boundary into that source only. Each materialization registers a
//! fresh run in a [`StreamInstrumentationRegistry`], and future control-plane tools can
//! read snapshots from the registry without putting actors on the element path.
//!
//! [`Source::instrumented`]: crate::Source::instrumented

use crate::stream::{BoxStream, StreamError, StreamResult};
use std::{
    collections::BTreeMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU8, AtomicU64, Ordering},
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};

const STATE_RUNNING: u8 = 0;
const STATE_DRAINING: u8 = 1;
const STATE_COMPLETED: u8 = 2;
const STATE_FAILED: u8 = 3;

/// Stable identifier for one instrumented stream materialization.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct StreamInstrumentationId(u64);

impl StreamInstrumentationId {
    /// Return the numeric id assigned by the registry.
    #[must_use]
    pub const fn get(self) -> u64 {
        self.0
    }
}

/// The externally visible state of one instrumented stream materialization.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StreamInstrumentationState {
    /// The stream has been materialized and has not reached a terminal state.
    Running,
    /// The stream is draining or was dropped before the boundary observed normal completion.
    Draining,
    /// The boundary observed upstream completion.
    Completed,
    /// The boundary observed a non-cancellation stream failure.
    Failed,
}

impl StreamInstrumentationState {
    #[must_use]
    const fn from_code(code: u8) -> Self {
        match code {
            STATE_DRAINING => Self::Draining,
            STATE_COMPLETED => Self::Completed,
            STATE_FAILED => Self::Failed,
            _ => Self::Running,
        }
    }

    #[must_use]
    const fn code(self) -> u8 {
        match self {
            Self::Running => STATE_RUNNING,
            Self::Draining => STATE_DRAINING,
            Self::Completed => STATE_COMPLETED,
            Self::Failed => STATE_FAILED,
        }
    }

    #[must_use]
    const fn is_terminal(self) -> bool {
        matches!(self, Self::Completed | Self::Failed)
    }
}

/// Point-in-time values for one instrumented stream materialization.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StreamInstrumentationSnapshot {
    /// Stable id assigned when the stream materialized.
    pub id: StreamInstrumentationId,
    /// User-provided instrumentation name.
    pub name: String,
    /// Count of successful elements observed by the instrumentation boundary.
    pub elements_through: u64,
    /// Restart count recorded by the owner of this materialization.
    pub restarts: u64,
    /// Current stream state.
    pub state: StreamInstrumentationState,
    /// Wall-clock time when this materialization registered.
    pub started_at: SystemTime,
    /// Wall-clock time of the most recent state transition.
    pub state_changed_at: SystemTime,
    /// Wall-clock terminal timestamp when the boundary observed completion or failure.
    pub finished_at: Option<SystemTime>,
    /// Elapsed time since `started_at`, or until `finished_at` for terminal runs.
    pub uptime: Duration,
}

/// Cloneable registry handle for instrumented stream materializations.
///
/// The registry is not consulted by ordinary streams. It is only captured by streams that
/// explicitly call [`Source::instrumented`], and registration happens once per materialization.
///
/// [`Source::instrumented`]: crate::Source::instrumented
#[derive(Clone, Debug, Default)]
pub struct StreamInstrumentationRegistry {
    inner: Arc<RegistryInner>,
}

#[derive(Debug, Default)]
struct RegistryInner {
    next_id: AtomicU64,
    runs: Mutex<BTreeMap<StreamInstrumentationId, Arc<StreamInstrumentationCounters>>>,
}

impl StreamInstrumentationRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register one materialized stream run.
    ///
    /// This is public so a future control plane can create a run handle for work it
    /// supervises directly. Most stream users should prefer [`Source::instrumented`].
    ///
    /// [`Source::instrumented`]: crate::Source::instrumented
    #[must_use]
    pub fn register(&self, name: impl Into<String>) -> StreamInstrumentationRun {
        let id = StreamInstrumentationId(self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1);
        let counters = Arc::new(StreamInstrumentationCounters::new(id, name.into()));
        self.inner
            .runs
            .lock()
            .expect("stream instrumentation registry poisoned")
            .insert(id, Arc::clone(&counters));
        StreamInstrumentationRun { counters }
    }

    /// Return a snapshot for one materialized run, if it is still retained.
    #[must_use]
    pub fn snapshot(&self, id: StreamInstrumentationId) -> Option<StreamInstrumentationSnapshot> {
        self.inner
            .runs
            .lock()
            .expect("stream instrumentation registry poisoned")
            .get(&id)
            .map(|counters| counters.snapshot())
    }

    /// Return snapshots for every retained materialized run, ordered by id.
    #[must_use]
    pub fn snapshots(&self) -> Vec<StreamInstrumentationSnapshot> {
        self.inner
            .runs
            .lock()
            .expect("stream instrumentation registry poisoned")
            .values()
            .map(|counters| counters.snapshot())
            .collect()
    }

    /// Drop a retained run from the registry.
    ///
    /// Removing a run only affects future snapshots; existing [`StreamInstrumentationRun`]
    /// handles continue to own their counters.
    pub fn remove(&self, id: StreamInstrumentationId) -> bool {
        self.inner
            .runs
            .lock()
            .expect("stream instrumentation registry poisoned")
            .remove(&id)
            .is_some()
    }
}

/// Handle for updating one instrumented materialization.
#[derive(Clone, Debug)]
pub struct StreamInstrumentationRun {
    counters: Arc<StreamInstrumentationCounters>,
}

impl StreamInstrumentationRun {
    /// The id assigned by the registry.
    #[must_use]
    pub fn id(&self) -> StreamInstrumentationId {
        self.counters.id
    }

    /// Record one successfully observed element.
    pub fn record_element(&self) {
        self.record_elements(1);
    }

    /// Record multiple successfully observed elements.
    pub fn record_elements(&self, elements: u64) {
        self.counters
            .elements_through
            .fetch_add(elements, Ordering::Relaxed);
    }

    /// Record one restart for this materialization/job owner.
    pub fn record_restart(&self) {
        self.record_restarts(1);
    }

    /// Record multiple restarts for this materialization/job owner.
    pub fn record_restarts(&self, restarts: u64) {
        self.counters
            .restarts
            .fetch_add(restarts, Ordering::Relaxed);
    }

    /// Mark the run as running.
    pub fn mark_running(&self) {
        self.counters
            .mark_state(StreamInstrumentationState::Running);
    }

    /// Mark the run as draining.
    pub fn mark_draining(&self) {
        self.counters
            .mark_state(StreamInstrumentationState::Draining);
    }

    /// Mark the run as completed.
    pub fn mark_completed(&self) {
        self.counters
            .mark_state(StreamInstrumentationState::Completed);
    }

    /// Mark the run as failed.
    pub fn mark_failed(&self) {
        self.counters.mark_state(StreamInstrumentationState::Failed);
    }

    /// Return a point-in-time snapshot of this run.
    #[must_use]
    pub fn snapshot(&self) -> StreamInstrumentationSnapshot {
        self.counters.snapshot()
    }
}

#[derive(Debug)]
struct StreamInstrumentationCounters {
    id: StreamInstrumentationId,
    name: Arc<str>,
    elements_through: AtomicU64,
    restarts: AtomicU64,
    state: AtomicU8,
    started_at_millis: u64,
    state_changed_at_millis: AtomicU64,
    finished_at_millis: AtomicU64,
}

impl StreamInstrumentationCounters {
    fn new(id: StreamInstrumentationId, name: String) -> Self {
        let now = unix_time_millis(SystemTime::now());
        Self {
            id,
            name: Arc::from(name),
            elements_through: AtomicU64::new(0),
            restarts: AtomicU64::new(0),
            state: AtomicU8::new(STATE_RUNNING),
            started_at_millis: now,
            state_changed_at_millis: AtomicU64::new(now),
            finished_at_millis: AtomicU64::new(0),
        }
    }

    fn mark_state(&self, state: StreamInstrumentationState) {
        let now = unix_time_millis(SystemTime::now());
        self.state.store(state.code(), Ordering::Relaxed);
        self.state_changed_at_millis.store(now, Ordering::Relaxed);
        if state.is_terminal() {
            self.finished_at_millis.store(now, Ordering::Relaxed);
        }
    }

    fn snapshot(&self) -> StreamInstrumentationSnapshot {
        let state = StreamInstrumentationState::from_code(self.state.load(Ordering::Relaxed));
        let started_at = system_time_from_millis(self.started_at_millis);
        let state_changed_at =
            system_time_from_millis(self.state_changed_at_millis.load(Ordering::Relaxed));
        let finished_at_millis = self.finished_at_millis.load(Ordering::Relaxed);
        let finished_at =
            (finished_at_millis != 0).then(|| system_time_from_millis(finished_at_millis));
        let uptime_end = finished_at.unwrap_or_else(SystemTime::now);
        let uptime = uptime_end
            .duration_since(started_at)
            .unwrap_or(Duration::ZERO);

        StreamInstrumentationSnapshot {
            id: self.id,
            name: self.name.to_string(),
            elements_through: self.elements_through.load(Ordering::Relaxed),
            restarts: self.restarts.load(Ordering::Relaxed),
            state,
            started_at,
            state_changed_at,
            finished_at,
            uptime,
        }
    }
}

pub(crate) struct InstrumentedStream<T> {
    input: BoxStream<T>,
    run: StreamInstrumentationRun,
    terminal_observed: bool,
}

impl<T> InstrumentedStream<T> {
    pub(crate) fn new(input: BoxStream<T>, run: StreamInstrumentationRun) -> Self {
        Self {
            input,
            run,
            terminal_observed: false,
        }
    }
}

impl<T> Iterator for InstrumentedStream<T> {
    type Item = StreamResult<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.terminal_observed {
            return None;
        }

        match self.input.next() {
            Some(Ok(item)) => {
                self.run.record_element();
                Some(Ok(item))
            }
            Some(Err(error)) => {
                self.terminal_observed = true;
                if matches!(error, StreamError::Cancelled) {
                    self.run.mark_draining();
                } else {
                    self.run.mark_failed();
                }
                Some(Err(error))
            }
            None => {
                self.terminal_observed = true;
                self.run.mark_completed();
                None
            }
        }
    }
}

impl<T> Drop for InstrumentedStream<T> {
    fn drop(&mut self) {
        if !self.terminal_observed {
            self.run.mark_draining();
        }
    }
}

fn unix_time_millis(time: SystemTime) -> u64 {
    time.duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_millis().min(u128::from(u64::MAX)) as u64)
        .unwrap_or(0)
}

fn system_time_from_millis(millis: u64) -> SystemTime {
    UNIX_EPOCH + Duration::from_millis(millis)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Keep, NotUsed, Sink, Source};
    use std::{
        sync::{
            Arc,
            atomic::{AtomicBool, Ordering},
        },
        thread,
    };

    #[test]
    fn enabled_counters_track_successful_stream() {
        let registry = StreamInstrumentationRegistry::new();

        let values = Source::from_iter(0_u64..4)
            .instrumented("success", &registry)
            .run_collect()
            .expect("instrumented stream succeeds");

        assert_eq!(values, vec![0, 1, 2, 3]);
        let snapshots = registry.snapshots();
        assert_eq!(snapshots.len(), 1);
        let snapshot = &snapshots[0];
        assert_eq!(snapshot.name, "success");
        assert_eq!(snapshot.elements_through, 4);
        assert_eq!(snapshot.restarts, 0);
        assert_eq!(snapshot.state, StreamInstrumentationState::Completed);
        assert!(snapshot.finished_at.is_some());
        assert!(snapshot.uptime >= Duration::ZERO);
    }

    #[test]
    fn enabled_counters_track_failure_after_successful_elements() {
        let registry = StreamInstrumentationRegistry::new();
        let error = StreamError::Failed("boom".into());

        let result = Source::from_iter([1_u64, 2])
            .concat(Source::failed(error.clone()))
            .instrumented("failure", &registry)
            .run_collect();

        assert_eq!(result, Err(error));
        let snapshot = registry
            .snapshots()
            .into_iter()
            .next()
            .expect("snapshot registered");
        assert_eq!(snapshot.elements_through, 2);
        assert_eq!(snapshot.state, StreamInstrumentationState::Failed);
        assert!(snapshot.finished_at.is_some());
    }

    #[test]
    fn enabled_counters_mark_cancelled_stream_as_draining() {
        let registry = StreamInstrumentationRegistry::new();
        let emitted = Arc::new(AtomicBool::new(false));
        let source = {
            let emitted = Arc::clone(&emitted);
            Source::from_materialized_factory(move |_materializer| {
                let emitted = Arc::clone(&emitted);
                Ok((
                    Box::new(std::iter::from_fn(move || {
                        if !emitted.swap(true, Ordering::SeqCst) {
                            return Some(Ok(1_u64));
                        }
                        loop {
                            if crate::stream::current_stream_cancelled()
                                .as_ref()
                                .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
                            {
                                return Some(Err(StreamError::Cancelled));
                            }
                            thread::park_timeout(Duration::from_millis(1));
                        }
                    })) as BoxStream<u64>,
                    NotUsed,
                ))
            })
        };

        let completion = source
            .instrumented("cancel", &registry)
            .to_mat(Sink::ignore(), Keep::right)
            .run()
            .expect("stream materializes");

        wait_for_snapshot(&registry, |snapshot| snapshot.elements_through == 1);
        drop(completion);
        let snapshot = wait_for_snapshot(&registry, |snapshot| {
            snapshot.state == StreamInstrumentationState::Draining
        });

        assert_eq!(snapshot.elements_through, 1);
        assert_eq!(snapshot.state, StreamInstrumentationState::Draining);
    }

    #[test]
    fn disabled_behavior_matches_plain_source() {
        let plain = Source::from_iter(0_u64..8)
            .map(|item| item * 2)
            .run_collect()
            .expect("plain stream succeeds");
        let registry = StreamInstrumentationRegistry::new();
        let instrumented = Source::from_iter(0_u64..8)
            .map(|item| item * 2)
            .instrumented("enabled", &registry)
            .run_collect()
            .expect("instrumented stream succeeds");

        assert_eq!(plain, instrumented);
        assert!(StreamInstrumentationRegistry::new().snapshots().is_empty());
    }

    #[test]
    fn run_handle_records_restarts() {
        let registry = StreamInstrumentationRegistry::new();
        let run = registry.register("job");

        run.record_restart();
        run.record_restarts(2);

        let snapshot = registry.snapshot(run.id()).expect("snapshot retained");
        assert_eq!(snapshot.restarts, 3);
        assert_eq!(snapshot.state, StreamInstrumentationState::Running);
    }

    fn wait_for_snapshot(
        registry: &StreamInstrumentationRegistry,
        predicate: impl Fn(&StreamInstrumentationSnapshot) -> bool,
    ) -> StreamInstrumentationSnapshot {
        for _ in 0..500 {
            if let Some(snapshot) = registry.snapshots().into_iter().next()
                && predicate(&snapshot)
            {
                return snapshot;
            }
            thread::sleep(Duration::from_millis(2));
        }
        panic!("timed out waiting for instrumentation snapshot");
    }
}