From 7c92d4f181c5bc8bdab74e2efcf44de19ad6c552 Mon Sep 17 00:00:00 2001 From: kid Date: Fri, 24 Jul 2026 01:39:06 +0800 Subject: [PATCH] fix(mem-wal): reject mismatched shard specs before epoch claim --- rust/lance/src/dataset/mem_wal/manifest.rs | 317 +++++++++++++++++++-- 1 file changed, 295 insertions(+), 22 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/manifest.rs b/rust/lance/src/dataset/mem_wal/manifest.rs index a2099e97d44..2a80afd860c 100644 --- a/rust/lance/src/dataset/mem_wal/manifest.rs +++ b/rust/lance/src/dataset/mem_wal/manifest.rs @@ -440,6 +440,31 @@ impl ShardManifestStore { // Epoch-based Writer Fencing // ======================================================================== + fn validate_claimable( + &self, + manifest: &ShardManifest, + incoming_shard_spec_id: u32, + ) -> Result<()> { + // A `Sealed` manifest is the drop-table 2PC in-doubt marker. + // Sophon's reconcile uses the "sealed" marker to distinguish this + // state from an ordinary epoch fence or invalid writer configuration. + if manifest.status == ShardStatus::Sealed { + return Err(Error::invalid_input(format!( + "shard {} is sealed; refusing claim (drop in flight)", + self.shard_id + ))); + } + + if manifest.shard_spec_id != incoming_shard_spec_id { + return Err(Error::invalid_input(format!( + "cannot claim shard {} with incoming shard_spec_id {}; \ + stored shard_spec_id {} is immutable", + self.shard_id, incoming_shard_spec_id, manifest.shard_spec_id + ))); + } + Ok(()) + } + /// Claim a shard by incrementing its writer epoch. /// /// This establishes single-writer semantics by: @@ -464,9 +489,10 @@ impl ShardManifestStore { /// /// # Errors /// - /// Returns an error if another writer claimed an equal-or-higher - /// epoch than our target, or if the manifest stays contended past - /// the retry budget. + /// Returns an error without claiming or mutating the manifest if the shard + /// is sealed or the incoming shard spec conflicts with its stored identity. + /// Also returns an error if another writer claimed an equal-or-higher epoch + /// than our target, or if the manifest stays contended past the retry budget. #[instrument(name = "manifest_claim_epoch", level = "info", skip_all, fields(shard_id = %self.shard_id, shard_spec_id))] pub async fn claim_epoch(&self, shard_spec_id: u32) -> Result<(u64, ShardManifest)> { const MAX_CLAIM_RETRIES: usize = 16; @@ -476,19 +502,8 @@ impl ShardManifestStore { // writer's epoch, and the tip it finds is what our write builds on. let current = self.refresh_latest().await?; - // A sealed shard is mid-drop (drop-table 2PC). Refuse the claim - // with a distinguishable error rather than minting a new epoch, - // so a caller that skips its own status check still cannot - // resurrect a shard being dropped. Sophon's reconcile keys on - // the "sealed" marker in this message to tell it apart from an - // ordinary epoch fence. - if let Some(m) = ¤t - && m.status == ShardStatus::Sealed - { - return Err(Error::invalid_input(format!( - "shard {} is sealed; refusing claim (drop in flight)", - self.shard_id - ))); + if let Some(manifest) = ¤t { + self.validate_claimable(manifest, shard_spec_id)?; } let (next_version, next_epoch, base_manifest) = match current { @@ -526,11 +541,13 @@ impl ShardManifestStore { return Ok((next_epoch, new_manifest)); } Err(write_err) => { - let latest_epoch = self - .refresh_latest() - .await? - .map(|m| m.writer_epoch) - .unwrap_or(0); + // Refresh uncached: a claim must see another process's + // write, and the tip it finds is what the retry builds on. + let latest = self.refresh_latest().await?; + if let Some(manifest) = &latest { + self.validate_claimable(manifest, shard_spec_id)?; + } + let latest_epoch = latest.map(|m| m.writer_epoch).unwrap_or(0); if latest_epoch >= next_epoch { return Err(Error::io(format!( "Failed to claim shard {} (version {}): another writer claimed epoch {} (>= our target {}): {}", @@ -711,6 +728,40 @@ mod tests { } } + /// Arm a store policy that lands `sneaked` on disk the instant a writer + /// tries to publish the same version: a local `put_if_absent` stages the + /// object first, so the policy fires on that staging PUT — after the + /// writer's read, before its conditional rename — and the rename then + /// finds the version already taken. + /// + /// The destination is derived from the store's temp dir, not by + /// reconstructing a filesystem path from the intercepted object-store + /// path: the mapping between the two differs by platform, while the + /// on-disk layout (`mem_wal_path` / `shard_manifest_path`) is fixed. + fn arm_version_sneak( + policy: &Arc>, + sneaked: &ShardManifest, + store_root: &std::path::Path, + ) { + let sneaked_filename = manifest_filename(sneaked.version); + let sneaked_bytes = pb::ShardManifest::from(sneaked).encode_to_vec(); + let destination = store_root.join(format!( + "_mem_wal/{}/manifest/{}", + sneaked.shard_id.as_hyphenated(), + sneaked_filename + )); + policy.lock().unwrap().set_before_policy( + "sneak-peer-manifest", + Arc::new(move |method: &str, path: &Path| { + if method == "put" && path.as_ref().contains(sneaked_filename.as_str()) { + std::fs::create_dir_all(destination.parent().unwrap()).unwrap(); + std::fs::write(&destination, sneaked_bytes.clone()).unwrap(); + } + Ok(()) + }), + ); + } + /// A warm cache must not hide a successor's claim from `check_fenced`. #[tokio::test] async fn check_fenced_sees_a_peer_through_a_warm_cache() { @@ -994,6 +1045,228 @@ mod tests { assert_eq!(second.wal_entry_position_last_seen, 42); } + #[tokio::test] + async fn test_claim_epoch_preserves_matching_shard_spec() { + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let manifest_store = ShardManifestStore::new(store, &base_path, shard_id, 2); + + let (first_epoch, first) = manifest_store.claim_epoch(1).await.unwrap(); + assert_eq!(first_epoch, 1); + assert_eq!(first.version, 1); + assert_eq!(first.shard_spec_id, 1); + + let (second_epoch, second) = manifest_store.claim_epoch(1).await.unwrap(); + assert_eq!(second_epoch, 2); + assert_eq!(second.version, 2); + assert_eq!(second.shard_spec_id, 1); + } + + #[rstest::rstest] + #[case::manual_shard(0, 1)] + #[case::configured_shard_rejects_manual_claim(1, 0)] + #[tokio::test] + async fn test_claim_epoch_rejects_mismatched_shard_spec( + #[case] stored_shard_spec_id: u32, + #[case] incoming_shard_spec_id: u32, + ) { + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let manifest_store = ShardManifestStore::new(store, &base_path, shard_id, 2); + + let initial = manifest_store + .initialize_shard(stored_shard_spec_id, HashMap::new()) + .await + .unwrap(); + + let error = manifest_store + .claim_epoch(incoming_shard_spec_id) + .await + .unwrap_err(); + assert!(matches!(&error, Error::InvalidInput { .. })); + let message = error.to_string(); + assert!( + message.contains(&shard_id.to_string()) + && message.contains(&format!("incoming shard_spec_id {incoming_shard_spec_id}")) + && message.contains(&format!("stored shard_spec_id {stored_shard_spec_id}")), + "error must include the shard id and both shard spec ids, got: {error}" + ); + + let after = manifest_store.latest().await.unwrap().unwrap(); + assert_eq!( + after, initial, + "a rejected claim must not change the manifest" + ); + } + + #[tokio::test] + async fn test_concurrent_claims_reject_mismatched_shard_spec() { + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let manifest_store = ShardManifestStore::new(store, &base_path, shard_id, 2); + + let (first, second) = + tokio::join!(manifest_store.claim_epoch(0), manifest_store.claim_epoch(1)); + let (winners, losers): (Vec<_>, Vec<_>) = + [first, second].into_iter().partition(Result::is_ok); + assert_eq!( + winners.len(), + 1, + "exactly one claim must succeed, got: {winners:?} / {losers:?}" + ); + let claimed = winners.into_iter().next().unwrap().unwrap(); + let error = losers.into_iter().next().unwrap().unwrap_err(); + + assert!(matches!(&error, Error::InvalidInput { .. })); + let message = error.to_string(); + assert!( + message.contains(&shard_id.to_string()) + && message.contains("incoming shard_spec_id") + && message.contains("stored shard_spec_id"), + "error must include the shard id and both shard spec ids, got: {error}" + ); + + let latest = manifest_store.latest().await.unwrap().unwrap(); + assert_eq!(latest, claimed.1); + assert_eq!(latest.version, 1); + assert_eq!(latest.writer_epoch, claimed.0); + assert_eq!(manifest_store.list_versions().await.unwrap(), vec![1]); + } + + /// A claim whose write loses the version race must re-validate whatever + /// actually landed before it can retry: the re-read in the conflict path + /// is the last chance to notice a peer claimed with a different shard + /// spec, or a drop's `Sealed` marker slid in between our read and write. + /// + /// The peer's manifest is sneaked in between the claim's read and its + /// conditional rename (see [`arm_version_sneak`]), so this test loses + /// that race deterministically where the concurrent test above only + /// usually does. + #[rstest::rstest] + #[case::different_shard_spec( + 0, + 1, + ShardStatus::Active, + vec!["incoming shard_spec_id 0", "stored shard_spec_id 1"] + )] + #[case::sealed_by_concurrent_drop(0, 0, ShardStatus::Sealed, vec!["sealed"])] + #[tokio::test] + async fn test_claim_epoch_revalidates_after_write_conflict( + #[case] incoming_shard_spec_id: u32, + #[case] sneaked_shard_spec_id: u32, + #[case] sneaked_status: ShardStatus, + #[case] expected_fragments: Vec<&'static str>, + ) { + let (store, base_path, temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + + // The peer takes version 1 with our target epoch, so the spec/sealed + // rejection must fire even where an epoch war would also apply. + let mut sneaked = create_test_manifest(shard_id, 1, 1); + sneaked.shard_spec_id = sneaked_shard_spec_id; + sneaked.status = sneaked_status; + + let policy = Arc::new(Mutex::new(ProxyObjectStorePolicy::new())); + arm_version_sneak(&policy, &sneaked, temp_dir.path()); + let mut proxied = (*store).clone(); + proxied.inner = Arc::new(ProxyObjectStore::new(store.inner.clone(), policy)); + let manifest_store = ShardManifestStore::new(Arc::new(proxied), &base_path, shard_id, 2); + + let error = manifest_store + .claim_epoch(incoming_shard_spec_id) + .await + .unwrap_err(); + assert!(matches!(&error, Error::InvalidInput { .. })); + let message = error.to_string(); + assert!( + message.contains(&shard_id.to_string()), + "error must name the shard, got: {error}" + ); + for fragment in expected_fragments { + assert!( + message.contains(fragment), + "error must contain {fragment:?}, got: {error}" + ); + } + + let after = manifest_store.latest().await.unwrap().unwrap(); + assert_eq!( + after, sneaked, + "a rejected claim must not change the manifest" + ); + assert_eq!(manifest_store.list_versions().await.unwrap(), vec![1]); + } + + /// The conflict path a claim must survive: a peer's `initialize_shard` + /// (same spec, epoch 0) lands between our read and our write. The + /// revalidation passes, the peer's epoch does not beat our target, and + /// the retry builds on the version that won the race. + #[tokio::test] + async fn test_claim_epoch_retries_past_benign_version_bump() { + let (store, base_path, temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + + let mut sneaked = create_test_manifest(shard_id, 1, 0); + sneaked.shard_spec_id = 7; + sneaked.wal_entry_position_last_seen = 42; + + let policy = Arc::new(Mutex::new(ProxyObjectStorePolicy::new())); + arm_version_sneak(&policy, &sneaked, temp_dir.path()); + let mut proxied = (*store).clone(); + proxied.inner = Arc::new(ProxyObjectStore::new(store.inner.clone(), policy)); + let manifest_store = ShardManifestStore::new(Arc::new(proxied), &base_path, shard_id, 2); + + let (epoch, claimed) = manifest_store.claim_epoch(7).await.unwrap(); + assert_eq!(epoch, 1, "an epoch-0 bump is not a real claimant"); + assert_eq!(claimed.version, 2); + assert_eq!(claimed.shard_spec_id, 7); + assert_eq!( + claimed.wal_entry_position_last_seen, 42, + "the retry must build on the version that won the race" + ); + assert_eq!( + manifest_store.list_versions().await.unwrap(), + vec![1, 2], + "the peer's version and ours must both be durable" + ); + } + + /// Writes that keep failing with no manifest on disk must exhaust the + /// retry budget: the conflict revalidation has nothing to check, so the + /// claim surfaces the contention error instead of spinning forever. + #[tokio::test] + async fn test_claim_epoch_exhausts_retries_when_writes_keep_failing() { + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let v1_filename = manifest_filename(1); + let policy = Arc::new(Mutex::new(ProxyObjectStorePolicy::new())); + policy.lock().unwrap().set_before_policy( + "503", + Arc::new(move |method: &str, path: &Path| { + if method == "put" && path.as_ref().contains(v1_filename.as_str()) { + return Err(object_store::Error::Generic { + store: "test", + source: "503 slow down".into(), + } + .into()); + } + Ok(()) + }), + ); + let mut proxied = (*store).clone(); + proxied.inner = Arc::new(ProxyObjectStore::new(store.inner.clone(), policy)); + let manifest_store = ShardManifestStore::new(Arc::new(proxied), &base_path, shard_id, 2); + + let error = manifest_store.claim_epoch(0).await.unwrap_err(); + assert!(matches!(error, Error::IO { .. }), "got: {error}"); + assert!( + error.to_string().contains("after 16 retries"), + "expected the contention error, got: {error}" + ); + assert!(manifest_store.latest().await.unwrap().is_none()); + assert!(manifest_store.list_versions().await.unwrap().is_empty()); + } + #[tokio::test] async fn test_claim_epoch_refuses_sealed_manifest() { // A `Sealed` manifest is the drop-table 2PC in-doubt marker: @@ -1018,7 +1291,7 @@ mod tests { // The claim is refused with the distinguishable "sealed" error — // and the manifest is left untouched (no new epoch minted). - let err = manifest_store.claim_epoch(0).await.unwrap_err(); + let err = manifest_store.claim_epoch(1).await.unwrap_err(); assert!( err.to_string().contains("sealed"), "expected a distinguishable sealed-refusal error, got: {err}"