lodestone-types 0.0.16

//! Bounded process-local stage measurements for development performance work.
//!
//! Measurements are reconstructible observations, never correctness or
//! acknowledgement authority. Each thread owns its counter cell, so hot paths
//! do not contend on one process-wide atomic. Snapshots walk a bounded registry
//! of retained cells and never reset cumulative state, including after a worker
//! thread exits.

use std::cell::RefCell;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};

use serde::Serialize;

const MAX_REGISTERED_THREADS: usize = 128;

/// Stable process-local performance stages used by profile-v2 measurements.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(usize)]
pub enum PerformanceStageV2 {
    SearchParse = 0,
    SearchAdmission = 1,
    SearchMask = 2,
    SearchExecute = 3,
    SearchMaterialize = 4,
    SearchEncode = 5,
    IngestParseValidate = 6,
    IngestWalEncode = 7,
    ObjectFileSync = 8,
    ObjectDirectorySync = 9,
    ObjectPutReadback = 10,
    RaftAppend = 11,
    CommittedIndexPersistence = 12,
    TantivyCommit = 13,
    TantivyReload = 14,
    ProjectionRebuild = 15,
    VisibilityProof = 16,
    IngestResponseMaterialize = 17,
}

impl PerformanceStageV2 {
    /// Complete stable stage inventory in serialized order.
    pub const ALL: [Self; 18] = [
        Self::SearchParse,
        Self::SearchAdmission,
        Self::SearchMask,
        Self::SearchExecute,
        Self::SearchMaterialize,
        Self::SearchEncode,
        Self::IngestParseValidate,
        Self::IngestWalEncode,
        Self::ObjectFileSync,
        Self::ObjectDirectorySync,
        Self::ObjectPutReadback,
        Self::RaftAppend,
        Self::CommittedIndexPersistence,
        Self::TantivyCommit,
        Self::TantivyReload,
        Self::ProjectionRebuild,
        Self::VisibilityProof,
        Self::IngestResponseMaterialize,
    ];

    /// Stable snake-case label emitted in receipts and reports.
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::SearchParse => "search_parse",
            Self::SearchAdmission => "search_admission",
            Self::SearchMask => "search_mask",
            Self::SearchExecute => "search_execute",
            Self::SearchMaterialize => "search_materialize",
            Self::SearchEncode => "search_encode",
            Self::IngestParseValidate => "ingest_parse_validate",
            Self::IngestWalEncode => "ingest_wal_encode",
            Self::ObjectFileSync => "object_file_sync",
            Self::ObjectDirectorySync => "object_directory_sync",
            Self::ObjectPutReadback => "object_put_readback",
            Self::RaftAppend => "raft_append",
            Self::CommittedIndexPersistence => "committed_index_persistence",
            Self::TantivyCommit => "tantivy_commit",
            Self::TantivyReload => "tantivy_reload",
            Self::ProjectionRebuild => "projection_rebuild",
            Self::VisibilityProof => "visibility_proof",
            Self::IngestResponseMaterialize => "ingest_response_materialize",
        }
    }

    const fn index(self) -> usize {
        self as usize
    }
}

#[derive(Debug)]
struct ThreadCounters {
    observations: [AtomicU64; PerformanceStageV2::ALL.len()],
    elapsed_nanos: [AtomicU64; PerformanceStageV2::ALL.len()],
}

impl ThreadCounters {
    fn new() -> Self {
        Self {
            observations: std::array::from_fn(|_| AtomicU64::new(0)),
            elapsed_nanos: std::array::from_fn(|_| AtomicU64::new(0)),
        }
    }

    fn record(&self, stage: PerformanceStageV2, elapsed: Duration) {
        let index = stage.index();
        self.observations[index].fetch_add(1, Ordering::Relaxed);
        let elapsed = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
        let _ = self.elapsed_nanos[index].fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
            Some(current.saturating_add(elapsed))
        });
    }
}

#[derive(Debug)]
struct Registry {
    threads: Mutex<Vec<Arc<ThreadCounters>>>,
    overflow: Arc<ThreadCounters>,
    overflow_assignments: AtomicU64,
    observation_enabled: AtomicBool,
}

impl Registry {
    fn new() -> Self {
        Self {
            threads: Mutex::new(Vec::new()),
            overflow: Arc::new(ThreadCounters::new()),
            overflow_assignments: AtomicU64::new(0),
            observation_enabled: AtomicBool::new(true),
        }
    }

    fn register_thread(&self) -> Arc<ThreadCounters> {
        let mut threads = match self.threads.lock() {
            Ok(threads) => threads,
            Err(poisoned) => poisoned.into_inner(),
        };
        if threads.len() >= MAX_REGISTERED_THREADS {
            self.overflow_assignments.fetch_add(1, Ordering::Relaxed);
            return Arc::clone(&self.overflow);
        }
        let counters = Arc::new(ThreadCounters::new());
        threads.push(Arc::clone(&counters));
        counters
    }

    fn registered_threads(&self) -> Vec<Arc<ThreadCounters>> {
        let threads = match self.threads.lock() {
            Ok(threads) => threads,
            Err(poisoned) => poisoned.into_inner(),
        };
        threads.clone()
    }
}

static REGISTRY: OnceLock<Registry> = OnceLock::new();

fn registry() -> &'static Registry {
    REGISTRY.get_or_init(Registry::new)
}

thread_local! {
    static THREAD_COUNTERS: RefCell<Option<Arc<ThreadCounters>>> = const { RefCell::new(None) };
}

fn thread_counters() -> Arc<ThreadCounters> {
    THREAD_COUNTERS.with(|slot| {
        if let Ok(slot) = slot.try_borrow()
            && let Some(counters) = slot.as_ref()
        {
            return Arc::clone(counters);
        }
        let counters = registry().register_thread();
        if let Ok(mut slot) = slot.try_borrow_mut() {
            *slot = Some(Arc::clone(&counters));
        }
        counters
    })
}

/// One cumulative reading for a stable stage.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct PerformanceStageReadingV2 {
    pub stage: &'static str,
    pub observations: u64,
    pub elapsed_nanos: u64,
}

/// One bounded, cumulative process-local snapshot.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct BoundedStageSnapshotV2 {
    pub format_version: u32,
    pub readings: Vec<PerformanceStageReadingV2>,
    pub registered_threads: u64,
    pub overflow_assignments: u64,
}

impl BoundedStageSnapshotV2 {
    /// Returns a saturating stage-by-stage delta from an earlier snapshot.
    #[must_use]
    pub fn since(&self, earlier: &Self) -> Self {
        let readings = self
            .readings
            .iter()
            .zip(&earlier.readings)
            .map(|(current, prior)| PerformanceStageReadingV2 {
                stage: current.stage,
                observations: current.observations.saturating_sub(prior.observations),
                elapsed_nanos: current.elapsed_nanos.saturating_sub(prior.elapsed_nanos),
            })
            .collect();
        Self {
            format_version: 2,
            readings,
            registered_threads: self.registered_threads,
            overflow_assignments: self.overflow_assignments.saturating_sub(earlier.overflow_assignments),
        }
    }

    /// Reads one stage by stable identity.
    #[must_use]
    pub fn reading(&self, stage: PerformanceStageV2) -> Option<&PerformanceStageReadingV2> {
        self.readings.get(stage.index())
    }
}

/// Takes a non-resetting process-local snapshot from at most 128 thread cells
/// plus one explicit overflow cell.
#[must_use]
pub fn bounded_stage_snapshot_v2() -> BoundedStageSnapshotV2 {
    let registry = registry();
    let threads = registry.registered_threads();
    let mut observations = [0_u64; PerformanceStageV2::ALL.len()];
    let mut elapsed_nanos = [0_u64; PerformanceStageV2::ALL.len()];
    for counters in threads.iter().chain(std::iter::once(&registry.overflow)) {
        for stage in PerformanceStageV2::ALL {
            let index = stage.index();
            observations[index] =
                observations[index].saturating_add(counters.observations[index].load(Ordering::Relaxed));
            elapsed_nanos[index] =
                elapsed_nanos[index].saturating_add(counters.elapsed_nanos[index].load(Ordering::Relaxed));
        }
    }
    BoundedStageSnapshotV2 {
        format_version: 2,
        readings: PerformanceStageV2::ALL
            .into_iter()
            .map(|stage| {
                let index = stage.index();
                PerformanceStageReadingV2 {
                    stage: stage.as_str(),
                    observations: observations[index],
                    elapsed_nanos: elapsed_nanos[index],
                }
            })
            .collect(),
        registered_threads: u64::try_from(threads.len()).unwrap_or(u64::MAX),
        overflow_assignments: registry.overflow_assignments.load(Ordering::Relaxed),
    }
}

/// Records an already measured exclusive duration.
pub fn record_performance_stage_v2(stage: PerformanceStageV2, elapsed: Duration) {
    if registry().observation_enabled.load(Ordering::Relaxed) {
        thread_counters().record(stage, elapsed);
    }
}

/// RAII timer for request-sized or I/O-sized stage boundaries.
pub struct PerformanceObservationV2 {
    stage: PerformanceStageV2,
    started: Option<Instant>,
}

impl PerformanceObservationV2 {
    #[must_use]
    pub fn start(stage: PerformanceStageV2) -> Self {
        let started = registry()
            .observation_enabled
            .load(Ordering::Relaxed)
            .then(Instant::now);
        Self { stage, started }
    }
}

impl Drop for PerformanceObservationV2 {
    fn drop(&mut self) {
        if let Some(started) = self.started {
            record_performance_stage_v2(self.stage, started.elapsed());
        }
    }
}

#[cfg(test)]
fn set_observation_enabled_for_test(enabled: bool) {
    registry().observation_enabled.store(enabled, Ordering::SeqCst);
}

#[cfg(test)]
mod tests {
    use std::hint::black_box;
    use std::sync::MutexGuard;

    use super::*;

    static PERFORMANCE_TEST_LOCK: Mutex<()> = Mutex::new(());

    struct RestoreObservation;

    impl Drop for RestoreObservation {
        fn drop(&mut self) {
            set_observation_enabled_for_test(true);
        }
    }

    fn request_sized_work() {
        let mut value = 0x9e37_79b9_7f4a_7c15_u64;
        for iteration in 0..250_000_u64 {
            value = value.rotate_left(7) ^ iteration.wrapping_mul(0xd6e8_feb8_6659_fd93);
            black_box(value);
        }
    }

    fn test_lock() -> MutexGuard<'static, ()> {
        match PERFORMANCE_TEST_LOCK.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }

    fn measure_request_window(observed: bool) -> u64 {
        set_observation_enabled_for_test(observed);
        let started = Instant::now();
        for _ in 0..64 {
            let _stage = PerformanceObservationV2::start(PerformanceStageV2::SearchExecute);
            request_sized_work();
        }
        u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)
    }

    fn percentile(values: &mut [u64], numerator: usize, denominator: usize) -> u64 {
        values.sort_unstable();
        let index = values
            .len()
            .saturating_sub(1)
            .saturating_mul(numerator)
            .div_ceil(denominator);
        values[index.min(values.len().saturating_sub(1))]
    }

    #[test]
    fn performance_v2_snapshot_is_bounded_monotonic_and_stage_exact() {
        let _guard = test_lock();
        let before = bounded_stage_snapshot_v2();
        record_performance_stage_v2(PerformanceStageV2::ObjectDirectorySync, Duration::from_nanos(17));
        let after = bounded_stage_snapshot_v2();
        let delta = after.since(&before);
        let reading = delta
            .reading(PerformanceStageV2::ObjectDirectorySync)
            .map(|reading| (reading.observations, reading.elapsed_nanos));
        assert_eq!(reading, Some((1, 17)));
        assert_eq!(delta.readings.len(), PerformanceStageV2::ALL.len());
        assert!(after.registered_threads <= MAX_REGISTERED_THREADS as u64);
    }

    #[test]
    fn performance_v2_snapshot_retains_completed_worker_observations() {
        let _guard = test_lock();
        let before = bounded_stage_snapshot_v2();
        let worker = std::thread::spawn(|| {
            record_performance_stage_v2(PerformanceStageV2::VisibilityProof, Duration::from_nanos(31));
        });
        assert!(worker.join().is_ok());

        let delta = bounded_stage_snapshot_v2().since(&before);
        let reading = delta
            .reading(PerformanceStageV2::VisibilityProof)
            .map(|reading| (reading.observations, reading.elapsed_nanos));
        assert_eq!(reading, Some((1, 31)));
    }

    #[test]
    fn performance_v2_no_observation_control_stays_within_three_percent() {
        let _guard = test_lock();
        let _restore = RestoreObservation;
        let mut observed = Vec::with_capacity(41);
        let mut unobserved = Vec::with_capacity(41);
        request_sized_work();
        for trial in 0..41 {
            if trial % 2 == 0 {
                observed.push(measure_request_window(true));
                unobserved.push(measure_request_window(false));
            } else {
                unobserved.push(measure_request_window(false));
                observed.push(measure_request_window(true));
            }
        }
        let observed_p50 = percentile(&mut observed.clone(), 50, 100);
        let observed_p95 = percentile(&mut observed, 95, 100);
        let unobserved_p50 = percentile(&mut unobserved.clone(), 50, 100);
        let unobserved_p95 = percentile(&mut unobserved, 95, 100);
        assert!(
            observed_p50 <= unobserved_p50.saturating_mul(103).div_ceil(100),
            "observed_p50={observed_p50} unobserved_p50={unobserved_p50}"
        );
        assert!(
            observed_p95 <= unobserved_p95.saturating_mul(103).div_ceil(100),
            "observed_p95={observed_p95} unobserved_p95={unobserved_p95}"
        );
    }
}