datum-core 0.10.5

//! Per-node fused-graph metrics and the disabled-path probe seam.

use super::FusedNodeAttributes;
use crate::MetricsLevel;
use std::{
    sync::atomic::{AtomicBool, Ordering},
    time::{Duration, Instant},
};

/// Immutable metrics for one fused graph node.
///
/// Scalar stages populate `elements_*`. Arrow batch stages populate `batches_*` and `rows_*`.
/// Typed fused kernels measure a whole execution chunk with a monotonic clock and publish an equal
/// share to the participating nodes; erased nodes measure their individual transitions. Stall time
/// is measured exactly around bounded-buffer/demand waits. Counts are always exact.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FusedNodeMetrics {
    pub stage_index: usize,
    pub stage_name: String,
    pub effective_name: String,
    pub metrics_level: MetricsLevel,
    pub elements_in: u64,
    pub elements_out: u64,
    pub batches_in: u64,
    pub batches_out: u64,
    pub rows_in: u64,
    pub rows_out: u64,
    pub processing_time: Duration,
    pub backpressure_stall_time: Duration,
    pub queue_depth: Option<usize>,
}

#[derive(Clone, Copy, Debug, Default)]
pub(super) struct MetricUnits {
    pub(super) elements: u64,
    pub(super) batches: u64,
    pub(super) rows: u64,
}

impl MetricUnits {
    pub(super) const fn scalar(elements: u64) -> Self {
        Self {
            elements,
            batches: 0,
            rows: 0,
        }
    }

    #[allow(dead_code)]
    pub(super) const fn batches(batches: u64, rows: u64) -> Self {
        Self {
            elements: 0,
            batches,
            rows,
        }
    }
}

#[derive(Debug)]
struct NodeState {
    report: FusedNodeMetrics,
    stall_started: Option<Instant>,
}

/// Fresh per-materialization state. Disabled runs are stack-only and own no node allocation.
#[derive(Debug)]
pub(super) struct FusedMetricsRun {
    enabled: AtomicBool,
    nodes: Vec<NodeState>,
}

#[derive(Debug)]
pub(super) enum NodeMeasurement {
    Disabled,
    Counts { stage_index: usize },
    Timing { stage_index: usize, start: Instant },
}

#[derive(Debug)]
pub(super) enum StallMeasurement {
    Disabled,
    Timing { stage_index: usize, start: Instant },
}

impl FusedMetricsRun {
    pub(super) fn disabled() -> Self {
        Self {
            enabled: AtomicBool::new(false),
            nodes: Vec::new(),
        }
    }

    pub(super) fn new(attributes: &[FusedNodeAttributes]) -> Self {
        let enabled = attributes.iter().any(|node| {
            node.attributes
                .metrics_level_hint()
                .unwrap_or(MetricsLevel::Off)
                != MetricsLevel::Off
        });
        if !enabled {
            return Self::disabled();
        }

        let nodes = attributes
            .iter()
            .map(|node| NodeState {
                report: FusedNodeMetrics {
                    stage_index: node.stage_index,
                    stage_name: node.stage_name.clone(),
                    effective_name: node.effective_name.clone(),
                    metrics_level: node
                        .attributes
                        .metrics_level_hint()
                        .unwrap_or(MetricsLevel::Off),
                    elements_in: 0,
                    elements_out: 0,
                    batches_in: 0,
                    batches_out: 0,
                    rows_in: 0,
                    rows_out: 0,
                    processing_time: Duration::ZERO,
                    backpressure_stall_time: Duration::ZERO,
                    queue_depth: None,
                },
                stall_started: None,
            })
            .collect();
        Self {
            enabled: AtomicBool::new(true),
            nodes,
        }
    }

    /// One relaxed load used to select the untouched hot loop or its instrumented twin.
    #[inline(always)]
    pub(super) fn enabled(&self) -> bool {
        metrics_gate_enabled(&self.enabled)
    }

    pub(super) fn timing_enabled(&self) -> bool {
        self.nodes.iter().any(|node| {
            matches!(
                node.report.metrics_level,
                MetricsLevel::Timing | MetricsLevel::Stalls
            )
        })
    }

    pub(super) fn complete_exact_counts(
        &mut self,
        stage_index: usize,
        input: MetricUnits,
        output: MetricUnits,
    ) {
        let Some(node) = self.nodes.get_mut(stage_index) else {
            return;
        };
        if node.report.metrics_level == MetricsLevel::Off {
            return;
        }
        node.report.elements_in = input.elements;
        node.report.elements_out = output.elements;
        node.report.batches_in = input.batches;
        node.report.batches_out = output.batches;
        node.report.rows_in = input.rows;
        node.report.rows_out = output.rows;
    }

    pub(super) fn complete_linear_metrics(
        &mut self,
        elements: u64,
        processing_time: Option<Duration>,
    ) {
        let timed_nodes = self
            .nodes
            .iter()
            .filter(|node| {
                matches!(
                    node.report.metrics_level,
                    MetricsLevel::Timing | MetricsLevel::Stalls
                )
            })
            .count();
        let processing_share = processing_time.and_then(|elapsed| {
            (timed_nodes != 0).then(|| {
                let nanos = elapsed.as_nanos() / timed_nodes as u128;
                Duration::from_nanos(nanos.min(u128::from(u64::MAX)) as u64)
            })
        });

        for node in &mut self.nodes {
            if node.report.metrics_level == MetricsLevel::Off {
                continue;
            }
            node.report.elements_in = elements;
            node.report.elements_out = elements;
            if let Some(share) = processing_share
                && matches!(
                    node.report.metrics_level,
                    MetricsLevel::Timing | MetricsLevel::Stalls
                )
            {
                node.report.processing_time = share;
            }
        }
    }

    pub(super) fn record_processing_time_share(
        &mut self,
        stage_indices: &[usize],
        elapsed: Duration,
    ) {
        let timed_nodes = stage_indices
            .iter()
            .filter(|&&stage_index| {
                self.nodes.get(stage_index).is_some_and(|node| {
                    matches!(
                        node.report.metrics_level,
                        MetricsLevel::Timing | MetricsLevel::Stalls
                    )
                })
            })
            .count();
        if timed_nodes == 0 {
            return;
        }
        let share = elapsed.as_nanos() / timed_nodes as u128;
        let share = Duration::from_nanos(share.min(u128::from(u64::MAX)) as u64);
        for &stage_index in stage_indices {
            let Some(node) = self.nodes.get_mut(stage_index) else {
                continue;
            };
            if matches!(
                node.report.metrics_level,
                MetricsLevel::Timing | MetricsLevel::Stalls
            ) {
                node.report.processing_time = node.report.processing_time.saturating_add(share);
            }
        }
    }

    /// Enter an enabled node site after the relaxed loop-entry gate has selected instrumentation.
    #[inline(always)]
    pub(super) fn begin_node(&mut self, stage_index: usize, input: MetricUnits) -> NodeMeasurement {
        let Some(node) = self.nodes.get_mut(stage_index) else {
            return NodeMeasurement::Disabled;
        };
        if node.report.metrics_level == MetricsLevel::Off {
            return NodeMeasurement::Disabled;
        }
        node.report.elements_in = node.report.elements_in.saturating_add(input.elements);
        node.report.batches_in = node.report.batches_in.saturating_add(input.batches);
        node.report.rows_in = node.report.rows_in.saturating_add(input.rows);

        match node.report.metrics_level {
            MetricsLevel::Off => NodeMeasurement::Disabled,
            MetricsLevel::Counts => NodeMeasurement::Counts { stage_index },
            MetricsLevel::Timing | MetricsLevel::Stalls => NodeMeasurement::Timing {
                stage_index,
                start: Instant::now(),
            },
        }
    }

    #[inline(always)]
    pub(super) fn end_node(&mut self, measurement: NodeMeasurement, output: MetricUnits) {
        let (stage_index, elapsed) = match measurement {
            NodeMeasurement::Disabled => return,
            NodeMeasurement::Counts { stage_index } => (stage_index, None),
            NodeMeasurement::Timing { stage_index, start } => (stage_index, Some(start.elapsed())),
        };
        let node = &mut self.nodes[stage_index].report;
        node.elements_out = node.elements_out.saturating_add(output.elements);
        node.batches_out = node.batches_out.saturating_add(output.batches);
        node.rows_out = node.rows_out.saturating_add(output.rows);
        if let Some(elapsed) = elapsed {
            node.processing_time = node.processing_time.saturating_add(elapsed);
        }
    }

    /// Observe a bounded stage transition. A no-output transition begins a stall; the next output
    /// closes it. This covers demand/peer waits in the erased bounded-buffer state machines.
    pub(super) fn observe_backpressure(&mut self, stage_index: usize, output_count: usize) {
        let Some(node) = self.nodes.get_mut(stage_index) else {
            return;
        };
        if node.report.metrics_level != MetricsLevel::Stalls {
            return;
        }
        if output_count == 0 {
            if node.stall_started.is_none() {
                node.stall_started = Some(Instant::now());
            }
            node.report.queue_depth = Some(node.report.queue_depth.unwrap_or(0).max(1));
        } else if let Some(started) = node.stall_started.take() {
            node.report.backpressure_stall_time = node
                .report
                .backpressure_stall_time
                .saturating_add(started.elapsed());
        }
    }

    /// Enter a bounded-buffer/demand wait site.
    #[inline(always)]
    #[allow(dead_code)]
    pub(super) fn begin_stall(&self, stage_index: usize) -> StallMeasurement {
        if !self.enabled.load(Ordering::Relaxed) {
            return StallMeasurement::Disabled;
        }
        if self
            .nodes
            .get(stage_index)
            .is_none_or(|node| node.report.metrics_level != MetricsLevel::Stalls)
        {
            return StallMeasurement::Disabled;
        }
        StallMeasurement::Timing {
            stage_index,
            start: Instant::now(),
        }
    }

    #[allow(dead_code)]
    pub(super) fn end_stall(&mut self, measurement: StallMeasurement, queue_depth: Option<usize>) {
        let StallMeasurement::Timing { stage_index, start } = measurement else {
            return;
        };
        let node = &mut self.nodes[stage_index].report;
        node.backpressure_stall_time = node.backpressure_stall_time.saturating_add(start.elapsed());
        if let Some(depth) = queue_depth {
            node.queue_depth = Some(node.queue_depth.unwrap_or(0).max(depth));
        }
    }

    pub(super) fn snapshot(&self) -> Vec<FusedNodeMetrics> {
        if !self.enabled.load(Ordering::Relaxed) {
            return Vec::new();
        }
        self.nodes
            .iter()
            .filter(|node| node.report.metrics_level != MetricsLevel::Off)
            .map(|node| node.report.clone())
            .collect()
    }
}

/// Stable, non-inlined shim used to inspect the disabled-path generated code.
#[doc(hidden)]
#[inline(never)]
pub fn metrics_disabled_path_probe(flag: &AtomicBool, value: u64) -> u64 {
    if !metrics_gate_enabled(flag) {
        return value.wrapping_mul(3).wrapping_add(1);
    }
    metrics_enabled_path_probe(value)
}

/// Same-signature no-site control for path-sensitive generated-code comparison.
#[doc(hidden)]
#[inline(never)]
pub fn metrics_no_site_probe(_flag: &AtomicBool, value: u64) -> u64 {
    value.wrapping_mul(3).wrapping_add(1)
}

#[inline(always)]
fn metrics_gate_enabled(flag: &AtomicBool) -> bool {
    if flag.load(Ordering::Relaxed) {
        return metrics_gate_enabled_cold();
    }
    false
}

#[cold]
#[inline(never)]
fn metrics_gate_enabled_cold() -> bool {
    true
}

#[cold]
#[inline(never)]
fn metrics_enabled_path_probe(value: u64) -> u64 {
    value.wrapping_mul(3).wrapping_add(2)
}