tidaldb/tidal/src/testing/cluster.rs

//! Simulated multi-region cluster for M8 UAT testing.
//!
//! Creates a set of ephemeral [`TidalDb`] instances wired with the real M8
//! distributed fabric: in-process transports, `spawn_receiver`, and
//! `ReplicationState`. Signal replication now traverses the full path:
//!
//! ```text
//! write_signal → encode_batch → channel send
//!                               ↓
//!                        spawn_receiver thread
//!                               ↓
//!                         apply_payload
//!                               ↓
//!                    SignalLedger::apply_wal_event
//!                    ReplicationState::advance
//! ```
//!
//! # Architecture
//!
//! * All nodes open with `NodeRole::Single` so they accept direct writes AND
//!   can be promoted to leader at any time.
//! * Every non-initial-leader region starts a `spawn_receiver` thread (via
//!   `db.start_replication(transport)`) that processes incoming WAL batches.
//! * `write_signal` encodes the event as a one-event WAL batch and ships it
//!   immediately to all non-partitioned followers.
//! * A `batch_log` records every shipped batch so `await_convergence` can
//!   re-deliver missed batches after a partition is healed.
//! * `await_convergence` ships any pending batches, then polls
//!   `ReplicationState::applied_seqno` until all active followers have caught
//!   up to the current leader's sequence number.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

use crate::db::TidalDb;
use crate::db::config::{NodeConfig, NodeRole};
use crate::query::retrieve::Retrieve;
use crate::query::search::Search;
use crate::replication::shard::{RegionId, ShardId};
use crate::replication::transport::WalSegmentPayload;
use crate::replication::{WalSegmentId, spawn_receiver};
use crate::schema::{EntityId, Schema, Timestamp};
use crate::signals::{NoopWalWriter, SignalLedger};
use crate::wal::format::batch::{EventRecord, encode_batch};

use super::cluster_transport::{BatchEntry, ReceiveOnlyTransport, redeliver_missed};

// ── Public API ────────────────────────────────────────────────────────────

/// A simulated node in the cluster.
pub struct SimulatedNode {
    /// The region this node belongs to.
    pub region_id: RegionId,
    /// The `TidalDb` instance for this node.
    pub db: TidalDb,
}

/// Configuration for building a [`SimulatedCluster`].
pub struct ClusterConfig {
    /// Region IDs to create.
    pub regions: Vec<RegionId>,
    /// Which region hosts the leader node.
    pub leader_region: RegionId,
    /// Schema shared by all nodes.
    pub schema: Schema,
    /// Ranking profiles shared by all nodes (optional, defaults to builtins only).
    pub profiles: Vec<crate::ranking::profile::RankingProfile>,
}

/// A simulated multi-region tidalDB cluster using the real M8 distributed
/// fabric.
///
/// Signal replication traverses the real WAL-batch encode → transport →
/// `apply_payload` → `ReplicationState::advance` pipeline instead of calling
/// `db.signal()` directly on followers.
pub struct SimulatedCluster {
    /// Current leader region ID. Mutable via `promote_leader`.
    leader_region: Mutex<RegionId>,
    /// All nodes indexed by region.
    nodes: HashMap<RegionId, SimulatedNode>,
    /// Per-follower channel senders (region → sender for WAL payloads).
    ///
    /// Only regions that have an active `spawn_receiver` thread appear here.
    /// When dropped, the corresponding receiver thread exits cleanly.
    follower_senders: HashMap<RegionId, crossbeam::channel::Sender<WalSegmentPayload>>,
    /// All batches ever shipped, for partition-recovery re-delivery.
    batch_log: Mutex<Vec<BatchEntry>>,
    /// Per-leader signal count used as the WAL seqno for that leader's batches.
    leader_seqnos: Mutex<HashMap<RegionId, u64>>,
    /// Total signals ever written to any leader (for [`relay_log_len`]).
    total_signals: AtomicU64,
    /// Regions currently isolated by a network partition.
    partitioned_regions: Arc<RwLock<HashSet<RegionId>>>,
    /// Schema kept for reference.
    #[allow(dead_code)]
    schema: Schema,
    /// Pre-computed map: signal type name → `u8` ID used in [`EventRecord`].
    signal_type_ids: HashMap<String, u8>,
}

impl SimulatedCluster {
    /// Build a cluster from the given configuration.
    ///
    /// All nodes are created immediately in ephemeral mode. Non-leader regions
    /// have a `spawn_receiver` thread started automatically.
    ///
    /// # Panics
    ///
    /// Panics if any `TidalDb` fails to open or if `start_replication` fails.
    #[must_use]
    pub fn build(config: ClusterConfig) -> Self {
        // Pre-compute signal type IDs via a temporary ledger.
        let scratch_ledger = SignalLedger::new(config.schema.clone(), Box::new(NoopWalWriter));
        let signal_type_ids: HashMap<String, u8> = config
            .schema
            .signals()
            .filter_map(|def| {
                scratch_ledger
                    .resolve_signal_type(def.name())
                    .ok()
                    .map(|id| (def.name().to_string(), id.as_u16() as u8))
            })
            .collect();

        // Derive all shard IDs (one per region).
        let all_shards: Vec<ShardId> = config.regions.iter().map(|r| ShardId(r.0)).collect();

        let mut nodes: HashMap<RegionId, SimulatedNode> = HashMap::new();
        let mut follower_senders: HashMap<RegionId, crossbeam::channel::Sender<WalSegmentPayload>> =
            HashMap::new();
        let mut leader_seqnos: HashMap<RegionId, u64> = HashMap::new();

        for &region in &config.regions {
            let shard_id = ShardId(region.0);
            let peer_shards: Vec<ShardId> = all_shards
                .iter()
                .copied()
                .filter(|&s| s != shard_id)
                .collect();

            let db = TidalDb::builder()
                .ephemeral()
                .with_schema(config.schema.clone())
                .with_profiles(config.profiles.clone())
                .with_cluster(NodeConfig {
                    role: NodeRole::Single,
                    shard_id,
                    peer_shards,
                    ..NodeConfig::default()
                })
                .open()
                .expect("ephemeral TidalDb with valid schema must open");

            // Wire a receiver for every non-leader region.
            if region != config.leader_region {
                let (tx, rx) = crossbeam::channel::bounded::<WalSegmentPayload>(1024);
                let transport = Arc::new(ReceiveOnlyTransport {
                    local_shard: shard_id,
                    rx,
                });
                // spawn_receiver directly: we already have Arc<ReceiveOnlyTransport>
                // which implements Transport. Use the replication state from the db.
                let ledger = db
                    .ledger()
                    .expect("ephemeral db with schema must have ledger")
                    .clone();
                let rep_state = db.replication_state().clone();
                let _handle = spawn_receiver(transport, ledger, rep_state);
                // Note: the JoinHandle is intentionally not stored — the receiver thread
                // will exit cleanly when `tx` (and all senders to `rx`) are dropped.
                follower_senders.insert(region, tx);
            }

            nodes.insert(
                region,
                SimulatedNode {
                    region_id: region,
                    db,
                },
            );
            leader_seqnos.insert(region, 0);
        }

        Self {
            leader_region: Mutex::new(config.leader_region),
            nodes,
            follower_senders,
            batch_log: Mutex::new(Vec::new()),
            leader_seqnos: Mutex::new(leader_seqnos),
            total_signals: AtomicU64::new(0),
            partitioned_regions: Arc::new(RwLock::new(HashSet::new())),
            schema: config.schema,
            signal_type_ids,
        }
    }

    /// The current leader region ID.
    #[must_use]
    pub fn leader_region(&self) -> RegionId {
        *self
            .leader_region
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// Reference to the leader node.
    ///
    /// # Panics
    ///
    /// Panics if the leader region has no node.
    #[must_use]
    pub fn leader(&self) -> &SimulatedNode {
        let region = self.leader_region();
        self.nodes
            .get(&region)
            .expect("leader region must have a node")
    }

    /// Reference to a node by region.
    ///
    /// # Panics
    ///
    /// Panics if no node exists for the region.
    #[must_use]
    pub fn node(&self, region: RegionId) -> &SimulatedNode {
        self.nodes
            .get(&region)
            .unwrap_or_else(|| panic!("no node for region {region}"))
    }

    /// Write a signal to the cluster leader and ship it to active followers.
    ///
    /// The signal is immediately applied to the leader's `TidalDb`. A one-event
    /// WAL batch is encoded and shipped via the channel transport to all
    /// non-partitioned followers with active receivers. The batch is also
    /// recorded in the `batch_log` for partition-recovery re-delivery.
    ///
    /// # Errors
    ///
    /// Returns `TidalError` if the signal type is not registered in the schema
    /// or if the leader write fails.
    pub fn write_signal(
        &self,
        signal_type: &str,
        entity_id: EntityId,
        weight: f64,
    ) -> crate::Result<()> {
        let ts = Timestamp::now();
        let leader_region = self.leader_region();
        let leader_shard = ShardId(leader_region.0);

        // Write to the leader's signal ledger.
        self.nodes[&leader_region]
            .db
            .signal(signal_type, entity_id, weight, ts)?;

        // Encode as a one-event WAL batch.
        let type_id = *self.signal_type_ids.get(signal_type).ok_or_else(|| {
            crate::TidalError::invalid_input(format!("unknown signal type '{signal_type}'"))
        })?;

        let seqno = {
            let mut seqnos = self
                .leader_seqnos
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let v = {
                let s = seqnos.entry(leader_region).or_insert(0);
                *s += 1;
                *s
            };
            drop(seqnos);
            v
        };

        let events = [EventRecord {
            entity_id: entity_id.as_u64(),
            signal_type: type_id,
            weight: weight as f32,
            timestamp_nanos: ts.as_nanos(),
        }];
        let bytes =
            encode_batch(&events, seqno, ts.as_nanos()).expect("WAL batch encoding must not fail");

        // Ship immediately to non-partitioned followers that have active receivers.
        let partitioned = self
            .partitioned_regions
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone();

        for (&region, tx) in &self.follower_senders {
            if region == leader_region || partitioned.contains(&region) {
                continue;
            }
            let payload = WalSegmentPayload {
                id: WalSegmentId::new(crate::replication::RegionId::SINGLE, leader_shard, seqno),
                bytes: bytes.clone(),
                event_count: 1,
            };
            // Ignore send errors: the receiver may have exited (e.g. after a crash).
            let _ = tx.send(payload);
        }

        // Record in batch_log for partition-recovery re-delivery.
        self.batch_log
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push(BatchEntry {
                source_shard: leader_shard,
                seqno,
                bytes,
            });

        self.total_signals.fetch_add(1, Ordering::Relaxed);
        Ok(())
    }

    /// Write a signal directly to a specific region (bypassing the leader).
    ///
    /// Used to simulate partitioned writes during split-brain scenarios.
    ///
    /// # Panics
    ///
    /// Panics if the signal write fails.
    pub fn write_signal_to_region(
        &self,
        region: RegionId,
        signal_type: &str,
        entity_id: EntityId,
        weight: f64,
    ) {
        let ts = Timestamp::now();
        self.nodes
            .get(&region)
            .unwrap_or_else(|| panic!("no node for region {region}"))
            .db
            .signal(signal_type, entity_id, weight, ts)
            .expect("signal write must succeed");
    }

    /// Wait for all non-partitioned followers to apply all pending WAL batches.
    ///
    /// * Ships any batches that were missed during a partition (re-delivery from
    ///   the `batch_log`).
    /// * Polls `ReplicationState::applied_seqno` for each active follower until
    ///   it reaches the current leader's latest seqno.
    ///
    /// # Panics
    ///
    /// Panics if convergence is not reached within `timeout`.
    pub fn await_convergence(&self, timeout: Duration) {
        let deadline = Instant::now() + timeout;

        loop {
            assert!(
                Instant::now() <= deadline,
                "convergence timeout: cluster did not converge within {timeout:?}"
            );

            let leader_region = self.leader_region();
            let leader_shard = ShardId(leader_region.0);

            let target_seqno = self
                .leader_seqnos
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .get(&leader_region)
                .copied()
                .unwrap_or(0);

            let partitioned = self
                .partitioned_regions
                .read()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .clone();

            // Re-deliver any batches missed during a partition.
            {
                let log = self
                    .batch_log
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                for (&rgn, tx) in &self.follower_senders {
                    if rgn == leader_region || partitioned.contains(&rgn) {
                        continue;
                    }
                    redeliver_missed(tx, &self.nodes[&rgn].db, &log);
                }
            }

            // Check whether all active non-partitioned followers have converged.
            let mut all_converged = true;
            for (&region, node) in &self.nodes {
                if region == leader_region || partitioned.contains(&region) {
                    continue;
                }
                // Only check regions that have an active receiver thread.
                if !self.follower_senders.contains_key(&region) {
                    continue;
                }
                let applied = node
                    .db
                    .replication_state()
                    .applied_seqno(leader_shard)
                    .unwrap_or(0);
                if applied < target_seqno {
                    all_converged = false;
                    break;
                }
            }

            if all_converged {
                return;
            }
            std::thread::sleep(Duration::from_millis(1));
        }
    }

    /// Read decay score from a specific region's node.
    #[must_use]
    pub fn read_decay_score(
        &self,
        region: RegionId,
        entity_id: EntityId,
        signal_type: &str,
    ) -> Option<f64> {
        self.nodes.get(&region).and_then(|n| {
            n.db.read_decay_score(entity_id, signal_type, 0)
                .ok()
                .flatten()
        })
    }

    /// Broadcast item metadata to all nodes.
    pub fn write_item_with_metadata(
        &self,
        entity_id: EntityId,
        metadata: &HashMap<String, String>,
    ) -> crate::Result<()> {
        for node in self.nodes.values() {
            node.db.write_item_with_metadata(entity_id, metadata)?;
        }
        Ok(())
    }

    /// Broadcast embedding writes to all nodes.
    pub fn write_item_embedding(
        &self,
        entity_id: EntityId,
        embedding: &[f32],
    ) -> crate::Result<()> {
        for node in self.nodes.values() {
            node.db.write_item_embedding(entity_id, embedding)?;
        }
        Ok(())
    }

    /// Run a RETRIEVE query against a region.
    pub fn retrieve(
        &self,
        region: RegionId,
        query: &Retrieve,
    ) -> crate::Result<crate::query::retrieve::Results> {
        self.node(region).db.retrieve(query)
    }

    /// Run a SEARCH query against a region.
    pub fn search(
        &self,
        region: RegionId,
        query: &Search,
    ) -> crate::Result<crate::query::search::SearchResults> {
        self.node(region).db.search(query)
    }

    /// Item count helper for health checks.
    pub fn item_count(&self, region: RegionId) -> u64 {
        self.node(region).db.item_count()
    }

    /// Current leader seqno.
    pub fn leader_seqno(&self) -> u64 {
        let leader = self.leader_region();
        self.leader_seqnos
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(&leader)
            .copied()
            .unwrap_or(0)
    }

    /// Mark a region as partitioned.
    pub fn partition_region(&self, region: RegionId) {
        let mut partitions = self
            .partitioned_regions
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        partitions.insert(region);
    }

    /// Heal a partitioned region and re-deliver any missed WAL batches.
    pub fn heal_region(&self, region: RegionId) {
        {
            let mut partitions = self
                .partitioned_regions
                .write()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            partitions.remove(&region);
        }
        if let Some(tx) = self.follower_senders.get(&region) {
            let log = self
                .batch_log
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            redeliver_missed(tx, &self.nodes[&region].db, &log);
        }
    }

    /// Whether a region is currently partitioned.
    pub fn is_partitioned(&self, region: RegionId) -> bool {
        self.partitioned_regions
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .contains(&region)
    }

    /// Total number of signals ever written to any leader (replaces relay log
    /// length in the old signal-replay implementation).
    #[must_use]
    pub fn relay_log_len(&self) -> u64 {
        self.total_signals.load(Ordering::Acquire)
    }

    /// Number of WAL batches applied from the initial leader (`ShardId(0)`)
    /// on a specific region's replication state.
    ///
    /// This is equivalent to the event count for all tests that do not involve
    /// leader promotion, since each signal produces exactly one WAL batch.
    #[must_use]
    pub fn applied_count(&self, region: RegionId) -> u64 {
        // ShardId(0) is always the initial leader's shard (RegionId(0)).
        // Tests that check `applied_count` do not involve leader promotion,
        // so this is always the correct source shard.
        self.nodes.get(&region).map_or(0, |n| {
            n.db.replication_state()
                .applied_seqno(ShardId(0))
                .unwrap_or(0)
        })
    }

    /// Access the shared partitioned-regions set (for fault injection via
    /// [`crate::testing::faults::NetworkPartition`]).
    #[must_use]
    pub const fn partitioned_regions(&self) -> &Arc<RwLock<HashSet<RegionId>>> {
        &self.partitioned_regions
    }

    /// Promote a follower to leader.
    ///
    /// After promotion `write_signal` routes writes to the new leader and ships
    /// batches to all other regions that have active receivers. The new leader
    /// must already exist as a node in the cluster.
    ///
    /// # Panics
    ///
    /// Panics if `new_leader` does not correspond to an existing node.
    pub fn promote_leader(&self, new_leader: RegionId) {
        assert!(
            self.nodes.contains_key(&new_leader),
            "cannot promote non-existent region {new_leader}"
        );
        let mut leader = self
            .leader_region
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        *leader = new_leader;
    }

    /// All region IDs in the cluster.
    #[must_use]
    pub fn regions(&self) -> Vec<RegionId> {
        self.nodes.keys().copied().collect()
    }
}