tidaldb/tidal-server/tests/cluster_sharding.rs

//! Tier-3 sharding × replication suite (m11p6 L4, REAL multi-process cluster).
//!
//! The m11p6 exit gate, run over real OS processes — 3 shard groups × RF=3, every
//! node a replica of every group, each group its own elected leader (balanced:
//! group `s` led by node `s`). Two pillars:
//!
//! 1. **Kill-node failover localizes to the dead node's groups** — SIGKILL a node
//!    under concurrent `ack=quorum` load and prove: (a) ONLY the groups that node
//!    led re-elect a new leader (the groups led by survivors keep theirs); (b) the
//!    re-election completes inside the failover budget (<10s); (c) READS NEVER STOP
//!    (a concurrent `/feed` poller on a survivor sees zero failures across the
//!    window — reads are served from local replicas, no leader needed); and
//!    (d) ZERO acknowledged-write loss per group — every item the client saw a
//!    2xx + `x-tidal-seq` for is present afterwards on its shard's NEW leader (the
//!    m11p4 vote restriction guarantees the elected leader holds every committed
//!    write).
//!
//! 2. **The L3 rebalance verbs move exactly one group** — `POST
//!    /cluster/shards/{id}/transfer` and `/cluster/promote?shard=` move one
//!    group's leadership and leave the others untouched, and `POST
//!    /cluster/shards/{id}/replicas` (remove) runs a per-group fenced conf-change.
//!
//! Kill-point count: `TIDAL_SHARDING_KILLPOINTS` (default 3 for CI; the recorded
//! exit-gate run sweeps more — see docs/planning/milestone-11/phase-6.md).
//!
//! Run: `cargo test -p tidal-server --features cluster-e2e --test cluster_sharding -- --nocapture`

#![cfg(feature = "cluster-e2e")]
#![allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::too_many_lines
)]

mod support;

use std::collections::BTreeMap;
use std::sync::{
    Arc,
    atomic::{AtomicBool, AtomicU64, Ordering},
};
use std::time::{Duration, Instant};

use support::{
    invariants::{AckLedger, item_token, post_acked},
    multiproc::{MultiProcCluster, convergence_budget},
};
use tidaldb::{replication::shard::ShardRouter, schema::EntityId};

/// Fast election (mirrors `cluster_election.rs`): heartbeat 100ms, timeout
/// 500–1000ms, lease 350ms — so a SIGKILL failover completes well inside the
/// budget. Applied to every group (each `ShardReplica` reads this block).
const FAST_ELECTION_YAML: &str = "election:\n  heartbeat_interval_ms: 100\n  election_timeout_min_ms: 500\n  election_timeout_max_ms: 1000\n  leader_lease_ms: 350";

/// "Failover < 10s p99" — the roadmap's m11p6 bound (also p4's).
const FAILOVER_BUDGET: Duration = Duration::from_secs(10);

const NODES: usize = 3;
const SHARDS: usize = 3;

/// CI-default kill points; the exit-gate run sets `TIDAL_SHARDING_KILLPOINTS` higher.
fn killpoints() -> usize {
    std::env::var("TIDAL_SHARDING_KILLPOINTS")
        .ok()
        .and_then(|v| v.parse().ok())
        .filter(|&n| n > 0)
        .unwrap_or(3)
}

/// Find the node index whose region name matches `name`.
fn idx_of(cluster: &MultiProcCluster, name: &str) -> usize {
    (0..cluster.len())
        .find(|&i| cluster.region_name(i) == name)
        .unwrap_or_else(|| panic!("no node named {name}"))
}

/// The m11p6 headline exit gate over real OS processes: kill a node under
/// `ack=quorum` load and prove only ITS shard-leaderships move (<10s), reads
/// never stop, and zero acknowledged writes are lost — across several kill points.
#[test]
fn mp_sharded_kill_node_moves_only_its_leaderships_zero_loss() {
    let mut cluster = MultiProcCluster::start_sharded(NODES, SHARDS, Some(FAST_ELECTION_YAML));
    // The gateway's entity→shard hash (the same FNV-1a router every node routes by).
    let router = ShardRouter::hash(SHARDS as u16).expect("build shard router");

    // Balanced placement: each group `s` is led (term 0) by node `s`.
    let initial = cluster.wait_shard_leaders_agreed(convergence_budget());
    for s in 0..SHARDS as u16 {
        assert_eq!(
            initial[&s],
            cluster.region_name(usize::from(s)),
            "group {s} should boot led by node {s} (balanced placement)"
        );
    }

    let mut entity: u64 = 1;
    for round in 0..killpoints() {
        // Wait for a fully-converged steady state (all groups agree their leader).
        let before = cluster.wait_shard_leaders_agreed(convergence_budget());
        // Target the leader of group `round % SHARDS` — "kill any node".
        let target_shard = (round % SHARDS) as u16;
        let target_name = before[&target_shard].clone();
        let target_idx = idx_of(&cluster, &target_name);
        // The groups this node currently leads (its leaderships must move) and the
        // rest (must stay put — the localization property).
        let led: Vec<u16> = (0..SHARDS as u16)
            .filter(|s| before[s] == target_name)
            .collect();
        let others: Vec<u16> = (0..SHARDS as u16)
            .filter(|s| before[s] != target_name)
            .collect();
        // Write + read through a SURVIVOR gateway so the client's entry node never
        // dies mid-request (it forwards group-leader writes; reads serve locally).
        let survivor = (0..NODES)
            .find(|&i| i != target_idx)
            .expect("a live survivor gateway distinct from the kill target");
        let survivor_base = cluster.node(survivor);

        // ── Concurrent reads that must never stop ─────────────────────────────
        let stop = Arc::new(AtomicBool::new(false));
        let read_failures = Arc::new(AtomicU64::new(0));
        let read_ok = Arc::new(AtomicU64::new(0));
        let reader = {
            let stop = Arc::clone(&stop);
            let fails = Arc::clone(&read_failures);
            let oks = Arc::clone(&read_ok);
            let base = survivor_base.clone();
            std::thread::spawn(move || {
                let client = reqwest::blocking::Client::builder()
                    .timeout(Duration::from_secs(3))
                    .build()
                    .unwrap();
                while !stop.load(Ordering::Acquire) {
                    match client
                        .get(format!("{base}/feed?profile=for_you&limit=24"))
                        .send()
                    {
                        Ok(r) if r.status().is_success() => {
                            oks.fetch_add(1, Ordering::Relaxed);
                        }
                        _ => {
                            fails.fetch_add(1, Ordering::Relaxed);
                        }
                    }
                    std::thread::sleep(Duration::from_millis(50));
                }
            })
        };

        // ── Two writer threads: ack=quorum items+views, entities spread across
        //    all groups by the gateway hash. Record only what the client saw acked.
        let mut writers = Vec::new();
        for w in 0..2u64 {
            let stop = Arc::clone(&stop);
            let base = survivor_base.clone();
            let first = entity + w * 100_000;
            writers.push(std::thread::spawn(move || {
                let client = reqwest::blocking::Client::builder()
                    .timeout(Duration::from_secs(3))
                    .build()
                    .unwrap();
                let mut acked: Vec<(u64, u64, Option<u64>)> = Vec::new();
                let mut e = first;
                while !stop.load(Ordering::Acquire) {
                    let item_seq = post_acked(
                        &client,
                        &base,
                        "/items",
                        "quorum",
                        &serde_json::json!({
                            "entity_id": e, "metadata": { "title": item_token(e) }
                        }),
                    );
                    let view_seq = post_acked(
                        &client,
                        &base,
                        "/signals",
                        "quorum",
                        &serde_json::json!({ "entity_id": e, "signal": "view", "weight": 1.0 }),
                    );
                    if let Some(s) = item_seq {
                        acked.push((e, s, view_seq));
                    }
                    e += 1;
                }
                acked
            }));
        }

        // Pseudo-random kill point per round (reproducible — no Math.random here).
        std::thread::sleep(Duration::from_millis(200 + (round as u64 * 131) % 500));
        cluster.kill_hard(target_idx);
        stop.store(true, Ordering::Release);

        // ── Collect the acked ledger and the per-shard coverage. ──────────────
        let mut ledger = AckLedger::new();
        let mut per_shard: BTreeMap<u16, usize> = BTreeMap::new();
        for wj in writers {
            for (e, item_seq, view_seq) in wj.join().expect("writer thread") {
                let s = router.route(EntityId::new(e));
                *per_shard.entry(s.0).or_default() += 1;
                ledger.record(e, item_seq, view_seq);
            }
        }
        reader.join().expect("reader thread");

        // ── (a)+(b) Failover localizes to the killed node's groups, <budget. ──
        let deadline = Instant::now() + FAILOVER_BUDGET;
        let after = loop {
            if let Some(now) = cluster.agreed_shard_leaders() {
                let led_moved = led
                    .iter()
                    .all(|s| now.get(s).is_some_and(|l| l != &target_name));
                let others_stable = others.iter().all(|s| now.get(s) == before.get(s));
                if led_moved && others_stable {
                    break now;
                }
            }
            assert!(
                Instant::now() <= deadline,
                "round {round}: killing {target_name} (led {led:?}) did not localize within \
                 {FAILOVER_BUDGET:?}: before={before:?}, now={:?}",
                cluster.agreed_shard_leaders()
            );
            std::thread::sleep(Duration::from_millis(100));
        };
        for s in &others {
            assert_eq!(
                after.get(s),
                before.get(s),
                "round {round}: group {s} (led by a survivor) must NOT change leader"
            );
        }

        // ── (c) Reads never stopped. ──────────────────────────────────────────
        assert_eq!(
            read_failures.load(Ordering::Relaxed),
            0,
            "round {round}: {} feed reads on survivor {survivor} FAILED during failover \
             (reads must never stop); {} succeeded",
            read_failures.load(Ordering::Relaxed),
            read_ok.load(Ordering::Relaxed)
        );
        assert!(
            read_ok.load(Ordering::Relaxed) > 0,
            "round {round}: the read poller never got a single 2xx — it was not exercising reads"
        );

        // A round with ZERO acked writes proves nothing — the failover would have
        // had no acknowledged state to lose. The writers run hundreds of quorum
        // writes in the pre-kill window, so an empty ledger means a setup fault
        // (writers never got a 2xx), not a passing round. The ledger spans every
        // group the gateway hash routed to (logged per-shard); the loss proof
        // below probes each acked write on ITS shard's new leader, so coverage of
        // a killed node's groups is whatever genuinely routed there this round.
        assert!(
            !ledger.is_empty(),
            "round {round}: no acked writes recorded — the kill tested nothing \
             (writers saw 0 quorum acks before the kill at {target_name})"
        );
        println!(
            "round {round}: killed {target_name} (led {led:?}); failover -> {after:?}; \
             {} acked writes, per-shard {per_shard:?}; reads {} ok / 0 failed",
            ledger.len(),
            read_ok.load(Ordering::Relaxed)
        );

        // ── (d) ZERO acknowledged-write loss: every acked item is present on ITS
        //    shard's NEW leader. A quorum ack means the write committed (a majority
        //    held it durably), and the m11p4 vote restriction guarantees the
        //    elected leader holds every committed write — so the post-failover
        //    shard leader is the authoritative place to prove presence (a still
        //    catching-up follower is not). Poll past the text index's 2s
        //    auto-commit (the leader scatter-reads all its hosted groups).
        for w in ledger.writes() {
            let shard = router.route(EntityId::new(w.entity_id)).0;
            let leader_name = after.get(&shard).expect("post-failover leader for shard");
            let leader_idx = idx_of(&cluster, leader_name);
            assert!(
                item_present(&cluster, leader_idx, w.entity_id, Duration::from_secs(15)),
                "round {round} (killed {target_name}): ACKNOWLEDGED LOSS — item {} \
                 (seq {}, shard {shard}, acked) is missing on shard {shard}'s new leader \
                 {leader_name} (node {leader_idx})",
                w.entity_id,
                w.item_seq
            );
        }

        // Restart the killed node so the cluster is whole for the next kill point
        // (one node down at a time keeps every group's quorum intact). The doubled
        // budget covers BOTH the rejoiner's boot+catch-up AND any group whose
        // leadership is still settling — stacked, not just convergence.
        entity += 10_000;
        cluster.restart(target_idx, &[]);
        let _ = cluster.wait_shard_leaders_agreed(convergence_budget() + convergence_budget());
    }
}

/// The L3 rebalance verbs over real processes: `POST /cluster/shards/{id}/transfer`
/// and `/cluster/promote?shard=` move EXACTLY one group's leadership, and the
/// per-group remove verb runs a fenced conf-change — proving `?shard=` selection
/// and the per-group reuse of the m11p4/m11p5 machinery end to end.
#[test]
fn mp_sharded_rebalance_verbs_move_one_group() {
    let cluster = MultiProcCluster::start_sharded(NODES, SHARDS, Some(FAST_ELECTION_YAML));
    let before = cluster.wait_shard_leaders_agreed(convergence_budget());
    for s in 0..SHARDS as u16 {
        assert_eq!(before[&s], cluster.region_name(usize::from(s)));
    }
    let node0 = cluster.region_name(0).to_string();
    let node1 = cluster.region_name(1).to_string();

    // ── Transfer group 0's leadership node0 -> node1 (RESTful rebalance verb). ──
    let resp = cluster.post(
        0,
        "/cluster/shards/0/transfer",
        &serde_json::json!({ "region": node1 }),
    );
    assert!(
        resp.status().is_success(),
        "shards/0/transfer must succeed, got {}",
        resp.status()
    );
    wait_until(FAILOVER_BUDGET, || {
        cluster
            .agreed_shard_leaders()
            .is_some_and(|m| m[&0] == node1 && m[&1] == before[&1] && m[&2] == before[&2])
    });
    let mid = cluster
        .agreed_shard_leaders()
        .expect("agreed after transfer");
    assert_eq!(
        mid[&0], node1,
        "group 0 leadership must have moved to node1"
    );
    assert_eq!(mid[&1], before[&1], "group 1 leadership must be untouched");
    assert_eq!(mid[&2], before[&2], "group 2 leadership must be untouched");

    // ── Move it back with the shard-scoped promote (`?shard=` selector). ──────
    let resp = cluster.post(
        0,
        "/cluster/promote?shard=0",
        &serde_json::json!({ "region": node0 }),
    );
    assert!(
        resp.status().is_success(),
        "promote?shard=0 must succeed, got {}",
        resp.status()
    );
    wait_until(FAILOVER_BUDGET, || {
        cluster
            .agreed_shard_leaders()
            .is_some_and(|m| m[&0] == node0)
    });

    // ── The per-group replica verb is wired + `?shard=`-scoped: assert its input
    //    validation and that `/cluster/members?shard=` selects the right group's
    //    roster. (The live add/remove conf-change reuses the m11p5 machinery
    //    per-group — proven by the membership suite — but a node's readiness is
    //    still node-global across its co-hosted groups, so removing a node from
    //    ONE group of a multi-group node is a tracked S>1 follow-up, NOT asserted
    //    here. See docs/planning/milestone-11/phase-6.md.)
    let bad_action = cluster.post(
        0,
        "/cluster/shards/0/replicas",
        &serde_json::json!({ "action": "frobnicate", "name": node1 }),
    );
    assert_eq!(
        bad_action.status().as_u16(),
        400,
        "an unknown shard-replica action must be a 400"
    );
    let add_missing_addrs = cluster.post(
        0,
        "/cluster/shards/0/replicas",
        &serde_json::json!({ "action": "add", "name": "region-9" }),
    );
    assert_eq!(
        add_missing_addrs.status().as_u16(),
        400,
        "add without grpc_addr/http_addr must be a 400"
    );
    // The `?shard=` selector resolves a hosted group's roster (full placement ⇒
    // every group lists all three nodes); an unhosted shard id is a 400.
    let roster = cluster.get_json(0, "/cluster/members?shard=2");
    assert_eq!(
        roster["members"].as_array().map(Vec::len),
        Some(NODES),
        "group 2's roster must list every node under full placement"
    );
    let unhosted = cluster.post(
        0,
        "/cluster/shards/9/transfer",
        &serde_json::json!({ "region": node0 }),
    );
    assert_eq!(
        unhosted.status().as_u16(),
        400,
        "targeting a shard this node does not host must be a 400"
    );
}

/// Whether item `entity` is searchable on node `idx` within `budget` — the
/// content presence probe (`/search?query=<token>`), polling past the text
/// index's ~2s auto-commit. A clean 2xx-with-no-hit through the whole budget is
/// genuine absence (returns false); transport/non-2xx is retried until the
/// budget, so a just-promoted leader still warming up is not charged as absence
/// prematurely.
fn item_present(cluster: &MultiProcCluster, idx: usize, entity: u64, budget: Duration) -> bool {
    let token = item_token(entity);
    let base = cluster.node(idx);
    let client = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(4))
        .build()
        .unwrap();
    let deadline = Instant::now() + budget;
    loop {
        if let Ok(resp) = client
            .get(format!("{base}/search?query={token}&limit=5"))
            .send()
            && resp.status().is_success()
        {
            let body: serde_json::Value = resp.json().unwrap_or(serde_json::Value::Null);
            let hit = body["items"].as_array().is_some_and(|items| {
                items
                    .iter()
                    .any(|it| it["entity_id"].as_u64() == Some(entity))
            });
            if hit {
                return true;
            }
        }
        if Instant::now() > deadline {
            return false;
        }
        std::thread::sleep(Duration::from_millis(200));
    }
}

/// Poll `cond` every 100ms until it returns true or `budget` elapses; returns the
/// final value of `cond` (so a caller can assert it true with context).
fn wait_until(budget: Duration, mut cond: impl FnMut() -> bool) -> bool {
    let deadline = Instant::now() + budget;
    loop {
        if cond() {
            return true;
        }
        if Instant::now() > deadline {
            return false;
        }
        std::thread::sleep(Duration::from_millis(100));
    }
}