tidaldb/tidal/src/testing/cluster_transport.rs

//! Internal transport types used by [`super::cluster::SimulatedCluster`].

use crate::{
    db::TidalDb,
    replication::{
        WalSegmentId,
        shard::{RegionId, ShardId},
        transport::{Transport, TransportError, WalSegmentPayload},
    },
};

// ── ChannelTransport: bidirectional crossbeam transport ─────────────────

/// Bidirectional crossbeam-channel transport for in-process cluster simulation.
///
/// Wraps both sender and receiver so `SimulatedCluster` can call
/// `send_segment` to ship WAL payloads and `spawn_receiver` can call
/// `recv_segment` to apply them.
pub(super) struct ChannelTransport {
    pub(super) local_shard: ShardId,
    pub(super) tx: crossbeam::channel::Sender<WalSegmentPayload>,
    pub(super) rx: crossbeam::channel::Receiver<WalSegmentPayload>,
}

impl Transport for ChannelTransport {
    fn send_segment(&self, _to: ShardId, payload: WalSegmentPayload) -> Result<(), TransportError> {
        self.tx.send(payload).map_err(|_| TransportError::Closed)
    }

    fn recv_segment(&self) -> Option<WalSegmentPayload> {
        self.rx.recv().ok()
    }

    fn local_shard(&self) -> ShardId {
        self.local_shard
    }
}

/// Receive-only crossbeam transport handed to the receiver thread.
///
/// Crucially this holds **only** the receiver end. The matching `Sender` lives
/// in [`super::cluster::SimulatedCluster`]; dropping the cluster drops every
/// sender, which closes the channel and makes [`recv_segment`] return `None`,
/// letting the receiver thread exit so its `SegmentReceiverHandle` can be
/// joined. (A bidirectional [`ChannelTransport`] cannot be joined on shutdown:
/// the receiver thread's own `Arc` keeps a `Sender` alive, so `recv` would
/// block forever.)
///
/// [`recv_segment`]: Transport::recv_segment
pub(super) struct RecvOnlyChannelTransport {
    pub(super) local_shard: ShardId,
    pub(super) rx: crossbeam::channel::Receiver<WalSegmentPayload>,
}

impl Transport for RecvOnlyChannelTransport {
    fn send_segment(
        &self,
        _to: ShardId,
        _payload: WalSegmentPayload,
    ) -> Result<(), TransportError> {
        // The receiver thread never ships; only the cluster's send side does.
        Err(TransportError::Closed)
    }

    fn recv_segment(&self) -> Option<WalSegmentPayload> {
        self.rx.recv().ok()
    }

    fn local_shard(&self) -> ShardId {
        self.local_shard
    }
}

// ── Batch log entry ──────────────────────────────────────────────────────

pub(super) struct BatchEntry {
    /// Which leader shard produced this batch.
    pub(super) source_shard: ShardId,
    /// 1-indexed sequence number scoped to `source_shard`.
    pub(super) seqno: u64,
    /// Encoded WAL batch bytes (from [`crate::wal::format::batch::encode_batch`]).
    pub(super) bytes: Vec<u8>,
}

// ── Payload construction ───────────────────────────────────────────────────

/// Build the single-event [`WalSegmentPayload`] for one replicated batch.
///
/// Both the eager ship path (`SimulatedCluster::write_signal`) and the recovery
/// path ([`redeliver_missed`]) emit exactly the same payload shape: a one-event
/// batch in [`RegionId::SINGLE`] whose `first_seq == last_seq == seqno`. Sharing
/// the construction here keeps the two sites from drifting on the load-bearing
/// `event_count`/`leader_last_seq` invariants (the receiver's monotonic
/// `advance` relies on `leader_last_seq == seqno`).
pub(super) const fn single_event_payload(
    source_shard: ShardId,
    seqno: u64,
    bytes: Vec<u8>,
) -> WalSegmentPayload {
    WalSegmentPayload {
        id: WalSegmentId::new(RegionId::SINGLE, source_shard, seqno),
        bytes,
        event_count: 1,
        // Single-event batch with first_seq = seqno → last WAL seq = seqno.
        leader_last_seq: seqno,
    }
}

// ── Re-delivery helper ────────────────────────────────────────────────────

/// Re-deliver all batch-log entries not yet applied to `db`, sending through
/// the given transport.
///
/// Called by `SimulatedCluster::heal_region` for immediate recovery after a
/// partition is healed, and by `await_convergence` during the polling loop.
///
/// # Load-bearing invariants
///
/// Re-delivery correctness rests on two invariants that hold today but are easy
/// to break silently:
///
/// 1. **Per-source FIFO ordering.** `log` is appended to in strictly increasing
///    `seqno` order per `source_shard` (`SimulatedCluster::write_signal` bumps
///    and commits the seqno under the `leader_seqnos` lock, then pushes the
///    `BatchEntry`), and we iterate `log` in that same order. The receiver
///    applies via `ReplicationState::advance`, which is monotonic per shard, so
///    a batch with `seqno = applied + 1` must arrive before `applied + 2` can be
///    accepted. Re-shipping out of order — or interleaving sources without
///    preserving each source's order — would leave gaps that never close. The
///    `entry.seqno > applied` gate below is what makes the scan safe to run
///    repeatedly (every poll) against the *same* growing log.
///
/// 2. **Idempotent re-application.** The same batch may be shipped more than
///    once: once eagerly by `write_signal` and again here on every
///    `await_convergence` poll and on `heal_region`. Re-applying an already-
///    applied `seqno` must be a no-op. We enforce that on the *send* side with
///    the `entry.seqno > applied` check, and the receiver enforces it again on
///    the *apply* side because `ReplicationState::advance` ignores any
///    `seqno <= applied`. Both guards are required: dropping the send-side check
///    would flood the transport with redundant payloads; relying on it alone
///    (without the receiver's monotonic advance) would double-apply on any race
///    between the `applied_seqno` read here and the receiver thread's advance.
pub(super) fn redeliver_missed(transport: &dyn Transport, db: &TidalDb, log: &[BatchEntry]) {
    for entry in log {
        // Read the follower's high-water-mark per source shard. The
        // `seqno > applied` gate enforces both invariants documented above:
        // it skips already-applied batches (idempotency) and, combined with the
        // FIFO append order of `log`, ships missing batches in seqno order.
        let applied = db
            .replication_state()
            .applied_seqno(entry.source_shard)
            .unwrap_or(0);
        if entry.seqno > applied {
            let payload =
                single_event_payload(entry.source_shard, entry.seqno, entry.bytes.clone());
            let _ = transport.send_segment(entry.source_shard, payload);
        }
    }
}

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

    /// Locks the load-bearing payload-shape invariants of the single source of
    /// truth that both `SimulatedCluster::write_signal` (W2 fix) and
    /// [`redeliver_missed`] now share: a one-event batch in [`RegionId::SINGLE`]
    /// whose `leader_last_seq == seqno`. The receiver's monotonic `advance`
    /// relies on `leader_last_seq == seqno`; a future edit that breaks either
    /// invariant on this single helper fails here instead of silently diverging
    /// between the eager-ship and recovery sites.
    #[test]
    fn single_event_payload_holds_replication_invariants() {
        let shard = ShardId(3);
        let seqno = 42;
        let bytes = vec![1u8, 2, 3, 4];
        let payload = single_event_payload(shard, seqno, bytes.clone());

        assert_eq!(payload.event_count, 1, "must be a single-event batch");
        assert_eq!(
            payload.leader_last_seq, seqno,
            "receiver advance relies on leader_last_seq == seqno"
        );
        assert_eq!(payload.id.region_id, RegionId::SINGLE);
        assert_eq!(payload.id.shard_id, shard);
        assert_eq!(payload.id.seqno, seqno);
        assert_eq!(payload.bytes, bytes);
    }
}