538 lines
19 KiB
Rust
538 lines
19 KiB
Rust
//! 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"
|
|
);
|
|
}
|