tidaldb/tidal-net/tests/snapshot_fetch.rs

//! m11p5 §2: `FetchSnapshot` over real sockets.
//!
//! A joiner (or a node behind a compacted leader) installs a staged snapshot
//! artifact before resuming the catch-up stream. These tests drive the REAL
//! `GrpcTransport` over real sockets with a fake [`SnapshotSource`] serving a
//! staged tempdir, and assert the wire contract end to end:
//!
//! - the manifest's BLAKE3 + sizes match the bytes the client receives, and
//!   the reassembled files are byte-identical;
//! - `needed=false` emits a header and NO file chunks;
//! - `Busy` maps to a RETRYABLE `UNAVAILABLE` (never `FAILED_PRECONDITION`,
//!   which would mis-route a joiner into the reseed class);
//! - the same-term term fence refuses a stale/newer puller;
//! - a `snapshot-required` catch-up trailer surfaces through the transport to
//!   the late-bound [`SnapshotRequiredSink`] (the seam the node latches its
//!   reseed marker on in stage B), while the standing retry timer keeps firing.

#![allow(clippy::unwrap_used)] // test assertions on known-good fixtures

use std::{
    collections::HashMap,
    net::SocketAddr,
    path::PathBuf,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

use tidal_net::{
    GrpcTransport, GrpcTransportConfig,
    client::fetch_snapshot_standalone,
    proto::snapshot_chunk::Chunk,
    sources::{
        SegmentChunk, SegmentReadError, SegmentSource, ServingSources, SnapshotRequiredSink,
        SnapshotSource, SnapshotStageError, SnapshotStaging,
    },
};
use tidaldb::replication::{shard::ShardId, transport::Transport};

/// Bind port 0 to obtain a free, OS-assigned address (tonic cannot bind 0
/// directly, so we resolve a concrete port up front).
fn free_addr() -> SocketAddr {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
    listener.local_addr().expect("local_addr")
}

/// A staged snapshot artifact: a tempdir holding `files` (relative path ->
/// bytes), with the manifest's BLAKE3/size computed from the bytes on disk.
struct StagedDir {
    _dir: tempfile::TempDir,
    root: PathBuf,
    manifest: Vec<(String, u64, [u8; 32])>,
}

impl StagedDir {
    fn new(files: &[(&str, &[u8])]) -> Self {
        let dir = tempfile::tempdir().expect("tempdir");
        let root = dir.path().to_path_buf();
        let mut manifest = Vec::new();
        for (rel, bytes) in files {
            let abs = root.join(rel);
            if let Some(parent) = abs.parent() {
                std::fs::create_dir_all(parent).expect("mkdir");
            }
            std::fs::write(&abs, bytes).expect("write file");
            let hash = blake3::hash(bytes);
            manifest.push(((*rel).to_string(), bytes.len() as u64, *hash.as_bytes()));
        }
        Self {
            _dir: dir,
            root,
            manifest,
        }
    }
}

/// A `SnapshotSource` serving a fixed [`StagedDir`] with a recorded
/// release-count (so a test can assert the per-consumer pin grace is signaled
/// exactly once). `needed` toggles the no-snapshot-needed path.
struct FakeSnapshots {
    staged: Arc<StagedDir>,
    needed: bool,
    snapshot_seq: u64,
    releases: Arc<AtomicU64>,
}

impl SnapshotSource for FakeSnapshots {
    fn stage(&self, _from_seqno: u64) -> Result<SnapshotStaging, SnapshotStageError> {
        Ok(SnapshotStaging {
            needed: self.needed,
            snapshot_seq: self.snapshot_seq,
            root: self.staged.root.clone(),
            files: if self.needed {
                self.staged.manifest.clone()
            } else {
                Vec::new()
            },
        })
    }
    fn release(&self) {
        self.releases.fetch_add(1, Ordering::AcqRel);
    }
}

/// A `SnapshotSource` that is always `Busy` — exercises the retryable
/// `UNAVAILABLE` mapping.
struct BusySnapshots;
impl SnapshotSource for BusySnapshots {
    fn stage(&self, _from_seqno: u64) -> Result<SnapshotStaging, SnapshotStageError> {
        Err(SnapshotStageError::Busy)
    }
    fn release(&self) {}
}

fn leader_with_snapshots(listen: SocketAddr, source: Arc<dyn SnapshotSource>) -> GrpcTransport {
    let sources = ServingSources::default();
    sources.set_snapshot_source(source);
    GrpcTransport::new_with_sources(
        GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: listen,
            insecure: true,
            ..GrpcTransportConfig::default()
        },
        sources,
    )
    .expect("leader transport")
}

/// Drain a `FetchSnapshot` stream into `(header, reassembled files)`.
async fn drain_snapshot(
    mut stream: tonic::Streaming<tidal_net::proto::SnapshotChunk>,
) -> (tidal_net::proto::SnapshotHeader, HashMap<String, Vec<u8>>) {
    let mut header = None;
    let mut files: HashMap<String, Vec<u8>> = HashMap::new();
    while let Some(msg) = stream.message().await.expect("stream chunk") {
        match msg.chunk.expect("chunk variant") {
            Chunk::Header(h) => header = Some(h),
            Chunk::File(f) => {
                let buf = files.entry(f.path.clone()).or_default();
                assert_eq!(
                    f.offset as usize,
                    buf.len(),
                    "file chunks must arrive in contiguous offset order"
                );
                buf.extend_from_slice(&f.data);
            }
        }
    }
    (header.expect("a header chunk must arrive first"), files)
}

/// The happy path: a 2-file artifact streams; the manifest BLAKE3 + sizes
/// match the received bytes exactly; the reassembled files are byte-identical;
/// and the per-consumer pin is released once on stream end.
#[test]
#[allow(clippy::significant_drop_tightening)] // transports intentionally live to test end
fn fetch_snapshot_streams_verified_files() {
    let file_a: &[u8] = b"the quick brown fox jumps over the lazy dog";
    let file_b: Vec<u8> = (0u8..=255).cycle().take(3_000_000).collect(); // > 1 MiB: multi-chunk
    let staged = Arc::new(StagedDir::new(&[
        ("ledger/000001.sst", file_a),
        ("wal/seg-0001.seg", &file_b),
    ]));
    let releases = Arc::new(AtomicU64::new(0));
    let leader_addr = free_addr();
    let _leader = leader_with_snapshots(
        leader_addr,
        Arc::new(FakeSnapshots {
            staged: Arc::clone(&staged),
            needed: true,
            snapshot_seq: 4096,
            releases: Arc::clone(&releases),
        }),
    );

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        let stream = fetch_snapshot_standalone(&leader_addr.to_string(), None, ShardId(0), 1, 0)
            .await
            .expect("fetch opens");
        let (header, files) = drain_snapshot(stream).await;

        assert!(header.needed, "a staged artifact reports needed=true");
        assert_eq!(header.snapshot_seq, 4096);
        assert_eq!(header.files.len(), 2, "the manifest lists both files");

        // Every manifest entry's size + BLAKE3 matches the received bytes.
        for entry in &header.files {
            let bytes = files
                .get(&entry.path)
                .unwrap_or_else(|| panic!("file {} streamed", entry.path));
            assert_eq!(
                bytes.len() as u64,
                entry.size,
                "received byte count must match the manifest size for {}",
                entry.path
            );
            let hash = blake3::hash(bytes);
            assert_eq!(
                hash.as_bytes().as_slice(),
                entry.blake3.as_slice(),
                "received BLAKE3 must match the manifest hash for {}",
                entry.path
            );
        }
        // Byte equality against the original sources.
        assert_eq!(files.get("ledger/000001.sst").unwrap().as_slice(), file_a);
        assert_eq!(files.get("wal/seg-0001.seg").unwrap().as_slice(), &file_b);
    });

    // The per-consumer pin (§2.1) is released exactly once when the stream
    // ends. Allow a moment for the server task's guard to drop after the
    // client finished draining.
    let deadline = Instant::now() + Duration::from_secs(5);
    while releases.load(Ordering::Acquire) == 0 && Instant::now() < deadline {
        std::thread::sleep(Duration::from_millis(10));
    }
    assert_eq!(
        releases.load(Ordering::Acquire),
        1,
        "the staging pin must be released exactly once on stream end"
    );
}

/// `needed=false`: the live WAL still serves the puller. The header says so and
/// NO file chunks follow; the no-pin staging is not released.
#[test]
#[allow(clippy::significant_drop_tightening)]
fn fetch_snapshot_not_needed_emits_header_only() {
    let staged = Arc::new(StagedDir::new(&[("ignored", b"unused")]));
    let releases = Arc::new(AtomicU64::new(0));
    let leader_addr = free_addr();
    let _leader = leader_with_snapshots(
        leader_addr,
        Arc::new(FakeSnapshots {
            staged,
            needed: false,
            snapshot_seq: 0,
            releases: Arc::clone(&releases),
        }),
    );

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        let stream = fetch_snapshot_standalone(&leader_addr.to_string(), None, ShardId(0), 1, 0)
            .await
            .expect("fetch opens");
        let (header, files) = drain_snapshot(stream).await;
        assert!(!header.needed, "the WAL can still serve: needed=false");
        assert!(
            header.files.is_empty(),
            "no manifest on the needed=false path"
        );
        assert!(
            files.is_empty(),
            "no file chunks follow a needed=false header"
        );
    });

    // A needed=false staging holds no pin, so release is never called.
    std::thread::sleep(Duration::from_millis(100));
    assert_eq!(
        releases.load(Ordering::Acquire),
        0,
        "a needed=false staging must not release a pin it never took"
    );
}

/// `Busy` staging maps to a RETRYABLE `UNAVAILABLE` — never
/// `FAILED_PRECONDITION` (which would mis-route a joiner into the reseed
/// class, §2).
#[test]
#[allow(clippy::significant_drop_tightening)]
fn fetch_snapshot_busy_is_unavailable() {
    let leader_addr = free_addr();
    let _leader = leader_with_snapshots(leader_addr, Arc::new(BusySnapshots));

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        // The open-time refusal surfaces as a Grpc status. `let-else` on the
        // Result directly (Streaming is not Debug, so `expect_err` is out).
        let Err(tidal_net::GrpcTransportError::Grpc(status)) =
            fetch_snapshot_standalone(&leader_addr.to_string(), None, ShardId(0), 1, 0).await
        else {
            panic!("Busy must surface a gRPC status (retryable open refusal)");
        };
        assert_eq!(
            status.code(),
            tonic::Code::Unavailable,
            "Busy is retryable UNAVAILABLE, not FAILED_PRECONDITION"
        );
    });
}

/// No snapshot source wired (a pre-m11p5 peer, §3.7): the RPC answers
/// `Unimplemented` so the joiner reports it loudly and retries the next seed.
#[test]
#[allow(clippy::significant_drop_tightening)]
fn fetch_snapshot_unimplemented_without_source() {
    let leader_addr = free_addr();
    let _leader = GrpcTransport::new(GrpcTransportConfig {
        local_shard: ShardId(0),
        listen_addr: leader_addr,
        insecure: true,
        ..GrpcTransportConfig::default()
    })
    .expect("leader transport");

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        let Err(tidal_net::GrpcTransportError::Grpc(status)) =
            fetch_snapshot_standalone(&leader_addr.to_string(), None, ShardId(0), 1, 0).await
        else {
            panic!("a node with no snapshot source must refuse with a gRPC status");
        };
        assert_eq!(status.code(), tonic::Code::Unimplemented);
    });
}

/// The same-term term fence (§2): a leader at term 5 refuses a puller stamping
/// term 3 with a `rejoin` typed trailer (rejoin the current term, not reseed).
#[test]
#[allow(clippy::significant_drop_tightening)]
fn fetch_snapshot_term_fence_refuses_stale_puller() {
    // An election-hooks stub claiming a fixed term.
    struct FixedTerm(u64);
    impl tidal_net::ElectionHooks for FixedTerm {
        fn self_claim(&self) -> (u64, u16) {
            (self.0, 0)
        }
        fn observe_leader_claim(
            &self,
            _term: u64,
            _leader_region: u16,
            _first_seq: u64,
        ) -> Result<(), tidal_net::ClaimRejection> {
            Ok(())
        }
        fn on_heartbeat(
            &self,
            _term: u64,
            _leader_region: u16,
            _stream_baseline: u64,
            _prev_log: tidaldb::replication::LogPosition,
        ) -> tidal_net::HeartbeatExchange {
            tidal_net::HeartbeatExchange {
                term: self.0,
                accepted: true,
            }
        }
        fn on_vote(&self, _rpc: tidaldb::replication::VoteRpc) -> tidaldb::replication::VoteReply {
            tidaldb::replication::VoteReply {
                term: self.0,
                granted: false,
            }
        }
        fn on_timeout_now(&self, _term: u64, _leader_region: u16) -> bool {
            false
        }
        fn on_observed_term(&self, _term: u64) {}
        fn report_term_acceptable(&self, reporter_term: u64) -> bool {
            reporter_term == self.0
        }
    }

    let staged = Arc::new(StagedDir::new(&[("f", b"x")]));
    let leader_addr = free_addr();
    let sources = ServingSources::default();
    sources.set_snapshot_source(Arc::new(FakeSnapshots {
        staged,
        needed: true,
        snapshot_seq: 1,
        releases: Arc::new(AtomicU64::new(0)),
    }));
    sources.set_election_hooks(Arc::new(FixedTerm(5)));
    let _leader = GrpcTransport::new_with_sources(
        GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: leader_addr,
            insecure: true,
            ..GrpcTransportConfig::default()
        },
        sources,
    )
    .expect("leader transport");

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        // Puller stamps term 3, below the source's 5.
        let Err(tidal_net::GrpcTransportError::Grpc(status)) =
            fetch_snapshot_standalone(&leader_addr.to_string(), None, ShardId(0), 1, 3).await
        else {
            panic!("a stale-term puller must be refused with a gRPC status");
        };
        assert_eq!(status.code(), tonic::Code::FailedPrecondition);
        assert_eq!(
            status
                .metadata()
                .get("x-tidal-catchup")
                .map(|v| v.as_bytes().to_vec()),
            Some(b"rejoin".to_vec()),
            "a stale puller gets the `rejoin` typed trailer, never reseed"
        );
    });
}

// ── Typed catch-up refusal → snapshot-required sink (§2.4) ──────────────────

/// A leader-side `SegmentSource` that always answers `Unavailable` — the
/// server maps this to `FAILED_PRECONDITION` + the `snapshot-required` trailer.
struct UnservableSegments;
impl SegmentSource for UnservableSegments {
    fn source_shard(&self) -> ShardId {
        ShardId(0)
    }
    fn stream_baseline(&self) -> u64 {
        0
    }
    fn flushed_seq(&self) -> u64 {
        100
    }
    fn collect_from(
        &self,
        _from_seq: u64,
        _max_events: u64,
        _max_bytes: usize,
    ) -> Result<Vec<SegmentChunk>, SegmentReadError> {
        Err(SegmentReadError::Unavailable {
            detail: "WAL segment format unknown".into(),
        })
    }
}

/// A scripted `SnapshotRequiredSink` recording whether it was invoked.
struct ScriptedSink {
    fired: Arc<AtomicBool>,
    shard: Arc<Mutex<Option<ShardId>>>,
}
impl SnapshotRequiredSink for ScriptedSink {
    fn snapshot_required(&self, shard: ShardId, _from_seqno: u64) {
        *self.shard.lock().unwrap() = Some(shard);
        self.fired.store(true, Ordering::Release);
    }
}

fn follower_config(listen: SocketAddr, leader: SocketAddr) -> GrpcTransportConfig {
    GrpcTransportConfig {
        local_shard: ShardId(1),
        listen_addr: listen,
        peers: std::iter::once((ShardId(0), leader.to_string())).collect(),
        insecure: true,
        catchup_retry_interval: Duration::from_millis(300),
        ..GrpcTransportConfig::default()
    }
}

/// The seam stage B latches its reseed marker on: a `StreamSegments` pull that
/// is refused `snapshot-required` surfaces through the transport to the
/// late-bound [`SnapshotRequiredSink`]. The standing retry timer keeps firing
/// regardless (the m11p4 re-arm-on-skip liveness fix must never be removed),
/// which is exactly why the sink can be invoked more than once — we assert it
/// fires AT LEAST once and carries the right shard.
#[test]
#[allow(clippy::significant_drop_tightening)] // transports intentionally live to test end
fn snapshot_required_trailer_surfaces_to_the_sink() {
    let leader_addr = free_addr();
    // The leader's segment source always answers Unavailable → the server maps
    // it to FAILED_PRECONDITION + the `snapshot-required` trailer.
    let leader_sources = ServingSources {
        segments: Some(Arc::new(UnservableSegments)),
        ..ServingSources::default()
    };
    let _leader = GrpcTransport::new_with_sources(
        GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: leader_addr,
            insecure: true,
            ..GrpcTransportConfig::default()
        },
        leader_sources,
    )
    .expect("leader transport");

    // The follower wires its snapshot-required sink BEFORE constructing the
    // transport (the cell is shared with the catch-up runner). Cloning the
    // ServingSources shares the same OnceLock cells, so setting on the clone
    // is visible to the runner.
    let fired = Arc::new(AtomicBool::new(false));
    let seen_shard = Arc::new(Mutex::new(None));
    let follower_sources = ServingSources::default();
    follower_sources.set_snapshot_required_sink(Arc::new(ScriptedSink {
        fired: Arc::clone(&fired),
        shard: Arc::clone(&seen_shard),
    }));
    let follower = GrpcTransport::new_with_sources(
        follower_config(free_addr(), leader_addr),
        follower_sources,
    )
    .expect("follower transport");

    // Trigger a catch-up pull: the leader refuses snapshot-required, the
    // transport reads the trailer and invokes the sink.
    follower.request_catchup(ShardId(0), 1);

    let deadline = Instant::now() + Duration::from_secs(10);
    while !fired.load(Ordering::Acquire) {
        assert!(
            Instant::now() < deadline,
            "the snapshot-required trailer never reached the sink: a follower \
             behind a compacted leader would never latch its reseed marker"
        );
        std::thread::sleep(Duration::from_millis(25));
    }
    assert_eq!(
        *seen_shard.lock().unwrap(),
        Some(ShardId(0)),
        "the sink learns which source shard demanded a snapshot"
    );
}