tidaldb/tidal/src/testing/crash_injector.rs

//! Fault injection for crash recovery testing.
//!
//! Provides a `CrashInjector` that triggers controlled panics at specific
//! write-path locations identified by `CrashPoint`. Hooks are installed via
//! a thread-local slot and compile away entirely in release builds without
//! `test` or `test-utils` feature.
//!
//! # Design
//!
//! - **One-shot**: each injector fires at most once. After firing, subsequent
//!   crossings are no-ops. This prevents cascading panics during unwind.
//! - **Thread-local**: injectors are installed per-thread via `install_injector`.
//!   This avoids global mutable state and is compatible with `catch_unwind`.
//! - **Zero overhead**: all call sites are behind `#[cfg(any(test, feature = "test-utils"))]`,
//!   so production builds contain no trace of the injection infrastructure.

use std::{
    cell::RefCell,
    panic::{AssertUnwindSafe, catch_unwind},
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
};

/// Identifies a specific point in the write path where a crash can be injected.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CrashPoint {
    /// After WAL append, before in-memory aggregation update.
    WalPreAggregate,
    /// After in-memory aggregation update completes.
    WalPostAggregate,
    /// Before the checkpoint `WriteBatch` is flushed to storage.
    CheckpointPreFlush,
    /// After the checkpoint `WriteBatch` is flushed to storage.
    CheckpointPostFlush,
    /// Inside `HotSignalState::on_signal`, after decay score CAS loops,
    /// before `last_update_ns` CAS.
    SignalAggregationUpdate,
    /// Inside `CohortSignalLedger::record`, after hot-tier update,
    /// before warm-tier increment.
    CohortLedgerUpdate,
    /// Inside `TidalDb::add_to_collection`, after persist, before return.
    CollectionIndexUpdate,
    /// Inside `CoEngagementIndex::record_positive`, after edge insertion loop,
    /// before eviction check.
    CoEngagementUpdate,
}

/// Controlled fault injector for crash recovery testing.
///
/// Fires a panic after `fire_after_n` crossings of the target `CrashPoint`.
/// One-shot: once fired, all subsequent `maybe_crash` calls are no-ops.
pub struct CrashInjector {
    target: CrashPoint,
    fire_after_n: u64,
    crossing_count: AtomicU64,
    fired: AtomicBool,
    armed: AtomicBool,
}

impl CrashInjector {
    /// Create a new injector targeting `target` that fires after `fire_after_n` crossings.
    ///
    /// The injector is armed by default. The first `fire_after_n` crossings are
    /// counted but do not trigger a panic. The `(fire_after_n + 1)`-th crossing
    /// (i.e., when `crossing_count >= fire_after_n`) triggers the panic.
    #[must_use]
    pub fn new(target: CrashPoint, fire_after_n: u64) -> Arc<Self> {
        Arc::new(Self {
            target,
            fire_after_n,
            crossing_count: AtomicU64::new(0),
            fired: AtomicBool::new(false),
            armed: AtomicBool::new(true),
        })
    }

    /// Arm the injector so subsequent crossings can trigger a panic.
    pub fn arm(&self) {
        // Release: any subsequent Acquire load of `armed` will observe this write.
        self.armed.store(true, Ordering::Release);
    }

    /// Disarm the injector so crossings are counted but never trigger a panic.
    pub fn disarm(&self) {
        // Release: pairs with the Acquire load in `maybe_crash`.
        self.armed.store(false, Ordering::Release);
    }

    /// Returns `true` if the injector has already fired (panicked).
    #[must_use]
    pub fn has_fired(&self) -> bool {
        // Acquire: see the Release store in `maybe_crash` when it sets `fired = true`.
        self.fired.load(Ordering::Acquire)
    }

    /// Returns the total number of crossings observed so far.
    #[must_use]
    pub fn crossing_count(&self) -> u64 {
        // Acquire: see the latest AcqRel fetch_add from any thread.
        self.crossing_count.load(Ordering::Acquire)
    }

    /// Check whether this crossing should trigger a panic.
    ///
    /// - If `point` does not match the target: returns immediately.
    /// - If not armed: returns immediately.
    /// - If already fired: returns immediately (one-shot).
    /// - Otherwise increments the crossing count. If `count >= fire_after_n`,
    ///   sets `fired = true` and panics.
    ///
    /// # Panics
    ///
    /// Panics (intentionally) when the crossing count reaches `fire_after_n`
    /// and the injector is armed. This is the core mechanism for simulating
    /// process crashes at specific write-path locations.
    pub fn maybe_crash(&self, point: CrashPoint) {
        if point != self.target {
            return;
        }
        // Acquire: see any preceding Release from arm() / disarm().
        if !self.armed.load(Ordering::Acquire) {
            return;
        }
        // Acquire: see the Release store that sets `fired = true` below (one-shot guard).
        if self.fired.load(Ordering::Acquire) {
            return;
        }

        // AcqRel: acquire ensures we see prior increments from concurrent callers;
        // release ensures our increment is visible to subsequent Acquire loads.
        let count = self.crossing_count.fetch_add(1, Ordering::AcqRel);
        if count >= self.fire_after_n {
            // Release: pairs with the Acquire load in has_fired() and the guard above.
            self.fired.store(true, Ordering::Release);
            panic!(
                "CrashInjector: simulated crash at {:?} after {} crossings",
                self.target, count
            );
        }
    }
}

// ── Thread-local slot ────────────────────────────────────────────────────────

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

/// Install a `CrashInjector` into the current thread's slot.
///
/// Replaces any previously installed injector.
pub fn install_injector(injector: Arc<CrashInjector>) {
    INJECTOR.with(|cell| {
        *cell.borrow_mut() = Some(injector);
    });
}

/// Clear the current thread's injector slot.
pub fn clear_injector() {
    INJECTOR.with(|cell| {
        *cell.borrow_mut() = None;
    });
}

/// Check the thread-local injector at a specific crash point.
///
/// If no injector is installed, or the injector targets a different point,
/// this is a no-op. Designed to be called from `#[cfg(any(test, feature = "test-utils"))]`
/// guarded hooks throughout the write path.
#[inline]
pub fn check_crash_point(point: CrashPoint) {
    INJECTOR.with(|cell| {
        if let Some(ref injector) = *cell.borrow() {
            injector.maybe_crash(point);
        }
    });
}

// ── run_with_crash helper ────────────────────────────────────────────────────

/// Execute a closure with a crash injector installed, catching injector panics.
///
/// Returns `Ok(T)` if the closure completes without the injector firing.
/// Returns `Err(CrashPoint)` if the injector triggered a panic.
/// Re-raises any panic that was NOT caused by the injector (e.g., a real bug).
///
/// # Errors
///
/// Returns `Err(CrashPoint)` when the injector fires during execution of `f`.
pub fn run_with_crash<T, F: FnOnce() -> T>(
    injector: &Arc<CrashInjector>,
    f: F,
) -> std::result::Result<T, CrashPoint> {
    let target = injector.target;
    install_injector(Arc::clone(injector));
    let result = catch_unwind(AssertUnwindSafe(f));
    clear_injector();
    match result {
        Ok(val) => Ok(val),
        Err(_) if injector.has_fired() => Err(target),
        Err(payload) => std::panic::resume_unwind(payload),
    }
}

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

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

    #[test]
    fn injector_fires_at_threshold() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 3);
        inj.maybe_crash(CrashPoint::WalPreAggregate); // count=0, not fired
        inj.maybe_crash(CrashPoint::WalPreAggregate); // count=1, not fired
        inj.maybe_crash(CrashPoint::WalPreAggregate); // count=2, not fired
        assert!(!inj.has_fired());
        // count=3 >= fire_after_n=3: fires
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            inj.maybe_crash(CrashPoint::WalPreAggregate);
        }));
        assert!(result.is_err());
        assert!(inj.has_fired());
    }

    #[test]
    fn injector_ignores_wrong_crash_point() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 0);
        inj.maybe_crash(CrashPoint::CheckpointPreFlush);
        inj.maybe_crash(CrashPoint::CohortLedgerUpdate);
        assert!(!inj.has_fired());
    }

    #[test]
    fn injector_one_shot() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 0);
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            inj.maybe_crash(CrashPoint::WalPreAggregate);
        }));
        assert!(result.is_err());
        assert!(inj.has_fired());
        // Second crossing should not panic (one-shot):
        inj.maybe_crash(CrashPoint::WalPreAggregate); // must not panic
    }

    #[test]
    fn injector_disarm_prevents_fire() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 0);
        inj.disarm();
        inj.maybe_crash(CrashPoint::WalPreAggregate);
        assert!(!inj.has_fired());
        inj.arm();
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            inj.maybe_crash(CrashPoint::WalPreAggregate);
        }));
        assert!(result.is_err());
    }

    #[test]
    fn run_with_crash_returns_ok_on_complete() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 1000);
        let result = run_with_crash(&inj, || 42);
        assert_eq!(result, Ok(42));
    }

    #[test]
    fn run_with_crash_returns_err_on_fire() {
        let inj = CrashInjector::new(CrashPoint::WalPreAggregate, 0);
        let result = run_with_crash(&inj, || {
            check_crash_point(CrashPoint::WalPreAggregate);
            42
        });
        assert_eq!(result, Err(CrashPoint::WalPreAggregate));
    }
}