tidaldb/tidal/tests/m6_crash_surfaces.rs

//! Milestone 7: Mid-write crash recovery for the M6 surfaces.
//!
//! The four crash points below are wired into the M6-surface write paths but
//! were never *armed* by any test — `m7_crash_m6.rs` only exercises clean
//! `close()` + reopen roundtrips, and `m7_crash_property.rs` only arms the WAL /
//! signal-ledger crash points. This file closes that gap: for each M6 crash
//! point we open a DB, perform the corresponding write, ARM the injector so the
//! process panics mid-write, let `Drop` run the durability path during unwind,
//! reopen, and assert the recovered surface satisfies the WAL crash tests'
//! durability bounds — **no fully-acked write lost, no write duplicated**.
//!
//! Crash points under test (exact `CrashPoint` variants):
//! 1. `CrashPoint::CohortLedgerUpdate`    — `CohortSignalLedger::record`, after
//!    the hot-tier `on_signal` CAS, before the warm-tier `increment`
//!    (`src/cohort/ledger.rs:100`).
//! 2. `CrashPoint::CollectionIndexUpdate` — `TidalDb::add_to_collection`, after
//!    `persist_collection`, before return (`src/db/collections.rs:78`).
//! 3. `CrashPoint::CoEngagementUpdate`    — `CoEngagementIndex::record_positive`,
//!    after the edge-insertion loop, before the eviction check
//!    (`src/entities/co_engagement.rs:98`).
//! 4. `CrashPoint::CheckpointPostFlush`   — `signals::checkpoint`, *after*
//!    `write_batch` + `flush` (`src/signals/checkpoint/mod.rs:118`).
//!
//! ## Durability model
//!
//! Cohort + co-engagement state is **checkpoint-only** (no WAL): `record` /
//! `record_positive` mutate in-memory state and are persisted by the final
//! checkpoint that `shutdown_inner` runs from both `close()` AND `Drop`. When an
//! armed injector panics mid-write, the unwind drops the `TidalDb`, `Drop` runs
//! `shutdown_inner`, and the *current* in-memory state is checkpointed. The
//! injector is one-shot, so the crash aborts the in-flight call (and the
//! enclosing write loop) but the durability path still runs to completion.
//!
//! Because the crash points sit at precise pre/post boundaries we can assert
//! *exact* recovered counts rather than just bounds:
//! - `CohortLedgerUpdate` fires before `warm.increment`, so a crash on call
//!   `K+1` (after `K` complete) leaves the warm counter at exactly `K`.
//! - `CollectionIndexUpdate` fires after `persist_collection`, so a crash on the
//!   add of item `K+1` (after `K` complete adds) leaves items `1..=K+1` durable.
//! - `CoEngagementUpdate` fires after the edge-insertion loop, so edges from the
//!   crashing call are already in-memory and checkpointed by `Drop`.
//! - `CheckpointPostFlush` fires after the flush completes, so the checkpoint is
//!   fully durable and every surface must recover exactly.

#![allow(clippy::unwrap_used, clippy::cast_precision_loss)]

use std::{collections::HashMap, time::Duration};

use tidaldb::{
    TidalDb,
    cohort::{CohortDef, Predicate},
    entities::Visibility,
    schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
    testing::{CrashInjector, CrashPoint, crash_injector::run_with_crash},
};

// ── Schema + metadata helpers ────────────────────────────────────────────────

/// M6 schema with view + like signals across all windows and velocity enabled.
fn m6_schema() -> tidaldb::schema::Schema {
    let mut builder = SchemaBuilder::new();
    for &(name, half_life_days) in &[("view", 7u64), ("like", 14)] {
        let _ = builder
            .signal(
                name,
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(half_life_days * 24 * 3600),
                },
            )
            .windows(&[
                Window::OneHour,
                Window::TwentyFourHours,
                Window::SevenDays,
                Window::AllTime,
            ])
            .velocity(true)
            .add();
    }
    builder.build().expect("m6 schema must be valid")
}

fn item_metadata(category: &str, creator_id: u64) -> HashMap<String, String> {
    let mut meta = HashMap::new();
    meta.insert("category".to_string(), category.to_string());
    meta.insert("creator_id".to_string(), creator_id.to_string());
    meta
}

fn user_metadata(locale: &str, primary_category: &str) -> HashMap<String, String> {
    let mut m = HashMap::new();
    m.insert("locale".to_string(), locale.to_string());
    m.insert("primary_category".to_string(), primary_category.to_string());
    m
}

fn open_db(dir: &std::path::Path, schema: &tidaldb::schema::Schema) -> TidalDb {
    TidalDb::builder()
        .with_data_dir(dir)
        .with_schema(schema.clone())
        .open()
        .unwrap()
}

// ── Test 1: CohortLedgerUpdate mid-write crash ───────────────────────────────
//
// Arms `CrashPoint::CohortLedgerUpdate` (src/cohort/ledger.rs:100). We record
// N cohort-attributed `like` signals onto a SINGLE item for a single cohort, so
// the per-key warm counter accumulates. The injector fires on cohort `record`
// call K+1 — after `hot.on_signal`, before `warm.increment` — so that call's
// warm increment is lost and the loop aborts. `Drop` checkpoints the in-memory
// state (warm == K). On reopen, the recovered cohort windowed count must be
// EXACTLY K: every fully-acked record present (no loss), the in-flight one
// dropped and never re-applied (no duplicate — cohort state is checkpoint-only,
// so there is no WAL-replay double-apply).

#[test]
fn cohort_ledger_update_crash_recovers_exactly_acked() {
    let dir = tempfile::tempdir().unwrap();
    let schema = m6_schema();

    let total_likes = 8u64;
    let crash_after = 5u64; // K = 5 complete records, crash on the 6th.
    let user = 1001u64;
    let item = EntityId::new(1);
    let cohort = "tech_en";

    {
        let db = open_db(dir.path(), &schema);

        db.define_cohort(CohortDef {
            name: cohort.to_string(),
            predicate: Predicate::And(vec![
                Predicate::Eq {
                    field: "locale".into(),
                    value: "en".into(),
                },
                Predicate::Eq {
                    field: "primary_category".into(),
                    value: "tech".into(),
                },
            ]),
        })
        .unwrap();

        db.write_user(EntityId::new(user), &user_metadata("en", "tech"))
            .unwrap();
        db.write_item_with_metadata(item, &item_metadata("tech", 100))
            .unwrap();

        // Each `like` with user context drives try_cohort_attribution -> record.
        // `like` is a positive-engagement signal, so it ALSO drives co-engagement,
        // but only cohort `record` crosses CohortLedgerUpdate.
        let injector = CrashInjector::new(CrashPoint::CohortLedgerUpdate, crash_after);
        let ts = Timestamp::now();
        let outcome = run_with_crash(&injector, || {
            for _ in 0..total_likes {
                db.signal_with_context("like", item, 1.0, ts, Some(user), Some(100))
                    .unwrap();
            }
            // db dropped here on the normal path; on crash, dropped during unwind.
            drop(db);
        });

        assert_eq!(
            outcome,
            Err(CrashPoint::CohortLedgerUpdate),
            "injector must fire at CohortLedgerUpdate"
        );
        assert!(injector.has_fired(), "injector should have fired");
    }

    // Reopen: cohort warm count must be EXACTLY the number of fully-acked records.
    {
        let db = open_db(dir.path(), &schema);
        let ledger = db.cohort_ledger();

        let count = ledger
            .read_windowed_count(cohort, item, "like", Window::AllTime)
            .unwrap();
        assert_eq!(
            count, crash_after,
            "cohort recovered count must equal acked records ({crash_after}): \
             no acked write lost, no in-flight/duplicate write applied; got {count}"
        );

        // Decay score must be finite + non-negative (recovered hot tier).
        if let Some(score) = ledger.read_decay_score(cohort, item, "like", 0).unwrap() {
            assert!(
                score.is_finite(),
                "recovered cohort decay score not finite: {score}"
            );
            assert!(
                score >= 0.0,
                "recovered cohort decay score negative: {score}"
            );
        }

        db.close().unwrap();
    }
}

// ── Test 2: CollectionIndexUpdate mid-write crash ────────────────────────────
//
// Arms `CrashPoint::CollectionIndexUpdate` (src/db/collections.rs:78). The crash
// point fires AFTER `persist_collection`, so on the add of item K+1 the in-memory
// index already holds item K+1 and the updated collection has already been
// persisted to durable storage. We crash on add K+1, the loop aborts, `Drop`
// re-persists the in-memory index. On reopen: items 1..=K+1 must be present
// (every acked add + the persisted-then-crashed add), items K+2.. must be absent
// (their adds never ran). No membership lost, no phantom membership.

#[test]
fn collection_index_update_crash_recovers_persisted_adds() {
    let dir = tempfile::tempdir().unwrap();
    let schema = m6_schema();
    let owner = EntityId::new(42);

    let total_items = 8u64; // items 10, 20, ..., 80
    let crash_after = 5u64; // K = 5 complete adds, crash on add #6 (item 60).

    let collection_name = "favorites";

    {
        let db = open_db(dir.path(), &schema);
        let cid = db
            .create_collection(owner, collection_name, Visibility::Private)
            .unwrap();

        let injector = CrashInjector::new(CrashPoint::CollectionIndexUpdate, crash_after);
        let outcome = run_with_crash(&injector, || {
            for i in 1..=total_items {
                db.add_to_collection(cid, EntityId::new(i * 10)).unwrap();
            }
            drop(db);
        });

        assert_eq!(
            outcome,
            Err(CrashPoint::CollectionIndexUpdate),
            "injector must fire at CollectionIndexUpdate"
        );
        assert!(injector.has_fired(), "injector should have fired");
    }

    // Reopen: items 1..=K+1 present (acked adds + the persisted-then-crashed add),
    // items K+2.. absent.
    {
        let db = open_db(dir.path(), &schema);
        let collections = db.list_collections(owner).unwrap();
        assert_eq!(
            collections.len(),
            1,
            "the collection must survive the mid-write crash"
        );
        let col = collections
            .iter()
            .find(|c| c.name == collection_name)
            .expect("favorites collection must survive");

        let index = db.collection_index();

        // The crash fires AFTER persist on add K+1, so item K+1 is durable too.
        let durable_through = crash_after + 1;
        for i in 1..=durable_through {
            let item = i * 10;
            assert!(
                index.contains(col.id, u32::try_from(item).unwrap()),
                "item {item} (acked or persisted-then-crashed add) must survive restart"
            );
        }
        for i in (durable_through + 1)..=total_items {
            let item = i * 10;
            assert!(
                !index.contains(col.id, u32::try_from(item).unwrap()),
                "item {item} (add never ran) must NOT be present — no phantom membership"
            );
        }

        assert_eq!(
            col.item_count as u64, durable_through,
            "collection item_count must equal durable adds ({durable_through}); got {}",
            col.item_count
        );

        db.close().unwrap();
    }
}

// ── Test 3: CoEngagementUpdate mid-write crash ───────────────────────────────
//
// Arms `CrashPoint::CoEngagementUpdate` (src/entities/co_engagement.rs:98). The
// crash point fires AFTER the edge-insertion loop, before the eviction check.
// We drive co-engagement by liking items in sequence for one user, which on each
// `like` adds edges (new_item -> each prior item). We crash on `record_positive`
// call K+1: that call's edges were already inserted before the crash point, so
// `Drop` checkpoints all edges through call K+1. On reopen every edge produced
// by calls 1..=K+1 must survive with its exact weight; later calls never ran so
// their edges must be absent. No edge lost, no duplicate-weight inflation.

#[test]
fn co_engagement_update_crash_recovers_inserted_edges() {
    let dir = tempfile::tempdir().unwrap();
    let schema = m6_schema();

    let user = 7u64;
    // Like items in sequence: 10, 20, 30, 40, 50. Each like adds edges to all
    // prior liked items. record_positive call #k corresponds to liking item k*10.
    let items = [10u64, 20, 30, 40, 50];
    // K = 3 complete record_positive calls (likes of 10, 20, 30), crash on the
    // 4th call (like of 40). The 4th call's edges (40->10, 40->20, 40->30) are
    // inserted BEFORE the crash point, so they survive too.
    let crash_after = 3u64;

    {
        let db = open_db(dir.path(), &schema);
        for &id in &items {
            db.write_item_with_metadata(EntityId::new(id), &item_metadata("jazz", 100))
                .unwrap();
        }

        let injector = CrashInjector::new(CrashPoint::CoEngagementUpdate, crash_after);
        let ts = Timestamp::now();
        let outcome = run_with_crash(&injector, || {
            for &id in &items {
                db.signal_with_context("like", EntityId::new(id), 1.0, ts, Some(user), Some(100))
                    .unwrap();
            }
            drop(db);
        });

        assert_eq!(
            outcome,
            Err(CrashPoint::CoEngagementUpdate),
            "injector must fire at CoEngagementUpdate"
        );
        assert!(injector.has_fired(), "injector should have fired");
    }

    // Reopen: edges from calls 1..=K+1 must survive. The crash fires on call
    // K+1 (like of item items[K]) AFTER its edges are inserted, so edges out of
    // items[K] back to items[0..K] are durable; edges out of items[K+1..] are not.
    {
        let db = open_db(dir.path(), &schema);
        let co_eng = db.co_engagement();

        let crash_idx = usize::try_from(crash_after).unwrap(); // index of the crashing like
        // Every directed edge (items[j] -> items[i]) for i < j and j <= crash_idx
        // must have weight 1.0 (each pair co-engaged exactly once).
        for j in 1..=crash_idx {
            for i in 0..j {
                let score = co_eng.score(EntityId::new(items[j]), EntityId::new(items[i]));
                assert!(
                    (score - 1.0).abs() < f32::EPSILON,
                    "edge ({}->{}) from acked/crashing like must be 1.0; got {score}",
                    items[j],
                    items[i]
                );
                // Asymmetric: reverse edge must be absent (no phantom edge).
                let reverse = co_eng.score(EntityId::new(items[i]), EntityId::new(items[j]));
                assert!(
                    (reverse - 0.0).abs() < f32::EPSILON,
                    "reverse edge ({}->{}) must be 0.0 (asymmetric); got {reverse}",
                    items[i],
                    items[j]
                );
            }
        }

        // Likes after the crash never ran, so their edges must be absent.
        for j in (crash_idx + 1)..items.len() {
            for i in 0..j {
                let score = co_eng.score(EntityId::new(items[j]), EntityId::new(items[i]));
                assert!(
                    (score - 0.0).abs() < f32::EPSILON,
                    "edge ({}->{}) from a never-run like must be absent; got {score}",
                    items[j],
                    items[i]
                );
            }
        }

        db.close().unwrap();
    }
}

// ── Test 4: CheckpointPostFlush crash ────────────────────────────────────────
//
// Arms `CrashPoint::CheckpointPostFlush` (src/signals/checkpoint/mod.rs:118).
// This point fires AFTER `write_batch` + `flush`, so the signal checkpoint is
// already fully durable on disk when the crash occurs. We write a full M6
// workload (signals + a collection), then `close()` — which runs the final
// checkpoint and crashes right after the signal-ledger flush.
//
// ## Why the signal surface recovers with at-least-once (not exactly-once)
//
// `shutdown_inner` flushes the signal checkpoint FIRST, then (later, after the
// crash point) writes the WAL checkpoint *marker* that fences WAL replay. The
// crash fires between those two steps, so the marker for this close is never
// written. On reopen, `db/open.rs:172-187` restores the durable checkpoint
// (which captured every view: count == views_per_item) AND then replays the WAL
// events that post-date the *previous* marker — i.e. the same view events — via
// `apply_wal_event`, re-adding them on top. This is the documented checkpoint +
// WAL overlap contract (see `m7_crash_property.rs` Test 2): the signal ledger
// guarantees **no acked write is lost**, with at-most-double overlap when a
// checkpoint is interrupted before its WAL marker lands. It does NOT promise
// exactly-once across an interrupted checkpoint — asserting that would be the
// wrong invariant. We therefore bound the signal surface no-lost + at-most-2x
// and assert exact recovery only for the collection surface, which is persisted
// per-add independently of the signal checkpoint/WAL and so is unaffected.

#[test]
fn checkpoint_post_flush_crash_recovers_durable_state() {
    let dir = tempfile::tempdir().unwrap();
    let schema = m6_schema();
    let owner = EntityId::new(1);

    let item_count = 6u64;
    let views_per_item = 3u64;

    {
        let db = open_db(dir.path(), &schema);

        for id in 1..=item_count {
            db.write_item_with_metadata(EntityId::new(id), &item_metadata("tech", 100))
                .unwrap();
        }

        // Signal-ledger surface: deterministic per-item view counts.
        let ts = Timestamp::now();
        for id in 1..=item_count {
            for _ in 0..views_per_item {
                db.signal("view", EntityId::new(id), 1.0, ts).unwrap();
            }
        }

        // Collection surface (persisted per-add, independent of signal checkpoint).
        let cid = db
            .create_collection(owner, "best-of", Visibility::Public)
            .unwrap();
        for id in [1u64, 3, 5] {
            db.add_to_collection(cid, EntityId::new(id)).unwrap();
        }

        // Crash AFTER the signal checkpoint's flush completes. close() consumes
        // the db and runs shutdown_inner -> ledger.checkpoint -> flush -> fire.
        // The flush already persisted the checkpoint, so reopen must recover it.
        let injector = CrashInjector::new(CrashPoint::CheckpointPostFlush, 0);
        let outcome = run_with_crash(&injector, move || {
            let _ = db.close();
        });

        assert_eq!(
            outcome,
            Err(CrashPoint::CheckpointPostFlush),
            "injector must fire at CheckpointPostFlush"
        );
        assert!(injector.has_fired(), "injector should have fired");
    }

    // Reopen: durable checkpoint yields no-lost signal counts (with bounded
    // checkpoint+WAL overlap); collection surface recovers exactly.
    {
        let db = open_db(dir.path(), &schema);

        for id in 1..=item_count {
            let count = db
                .read_windowed_count(EntityId::new(id), "view", Window::AllTime)
                .unwrap();
            // No-lost: the durable checkpoint captured every acked view.
            assert!(
                count >= views_per_item,
                "item {id}: post-flush crash lost acked views — expected >= {views_per_item}; got {count}"
            );
            // Bounded overlap: an interrupted checkpoint (marker not yet written)
            // replays the same WAL events once on top — at most 2x, never runaway.
            assert!(
                count <= 2 * views_per_item,
                "item {id}: checkpoint+WAL overlap exceeded the 2x bound — \
                 expected <= {}; got {count}",
                2 * views_per_item
            );
            let score = db.read_decay_score(EntityId::new(id), "view", 0).unwrap();
            assert!(score.is_some(), "item {id} decay score must survive");
            let score = score.unwrap();
            assert!(
                score.is_finite() && score > 0.0,
                "item {id} score must be finite + positive: {score}"
            );
        }

        let collections = db.list_collections(owner).unwrap();
        assert_eq!(
            collections.len(),
            1,
            "collection must survive post-flush crash"
        );
        let col = &collections[0];
        assert_eq!(col.name, "best-of");
        assert_eq!(col.item_count, 3, "collection must recover all 3 items");
        let index = db.collection_index();
        for id in [1u64, 3, 5] {
            assert!(
                index.contains(col.id, u32::try_from(id).unwrap()),
                "collection item {id} must survive"
            );
        }
        assert!(
            !index.contains(col.id, 2),
            "item 2 was never added — must not appear"
        );

        db.close().unwrap();
    }
}