diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 8ccb3bc..0a67b7d 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -213,7 +213,7 @@ async fn run_task(state: &Arc, mut claim: TaskClaim, timing: TaskCl } } let commit_start = std::time::Instant::now(); - let finished = state.task_store.finish(claim, outcome).await; + let finished = finish_task(state, claim, outcome).await; metrics::histogram!("master_task_phase_duration_seconds", "kind" => kind, "phase" => "commit") .record(commit_start.elapsed().as_secs_f64()); if let Err(error) = finished { @@ -240,6 +240,37 @@ async fn run_task(state: &Arc, mut claim: TaskClaim, timing: TaskCl } } +/// Publish completion before admitting any continuation for the same target. +async fn finish_task( + state: &Arc, + claim: TaskClaim, + outcome: Result, +) -> lance::Result<()> { + let target = claim.task.target.clone(); + let continue_resident = claim.task.kind == TaskKind::MergeWal + && outcome.is_ok() + && crate::local_merge::selected(&state.config, &target); + state.task_store.finish(claim, outcome).await?; + if continue_resident { + // The previous task and its write authority are durably released first. + // Retain ordinary queue fairness, dedupe and permits for the successor. + // Any failure here leaves durable demand for the periodic recovery scan; + // it must not turn a completed storage operation into a failed task. + let coordinator = state.task_store.merge_coordinator(); + let result = async { + if let Some(request) = coordinator.get_merge_request(&target).await? { + enqueue_merge_request(state, &coordinator, &request).await?; + } + Ok::<_, String>(()) + } + .await; + if let Err(error) = result { + tracing::warn!(%target, %error, "resident merge continuation enqueue deferred to recovery"); + } + } + Ok(()) +} + /// Whether a task error is the deterministic "manifest names a data file /// that is not there" failure. Every scan of the base table then fails the /// same way, from the master's compaction and from every worker's merge @@ -1860,6 +1891,196 @@ mod tests { ); } + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_completion_enqueues_durable_demand_without_periodic_scan() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.local_merge.master_local_merge_targets = vec!["exp".into()]; + cfg.catchup.shards = vec!["writer".into()]; + let state = MasterState::new(cfg).await.unwrap(); + let coordinator = state.task_store.merge_coordinator(); + let first = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let claim = state.task_store.claim_next().await.unwrap().unwrap(); + coordinator.request_merge("exp").await.unwrap(); + finish_task( + &state, + claim, + Ok("local WAL merge reclaimed 2 generations".into()), + ) + .await + .unwrap(); + // No scheduler, scan, or explicit enqueue call runs after completion. + let next = state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap() + .expect("successful native pass must enqueue its durable continuation"); + assert_ne!(first.id, next); + let successor = state.task_store.claim_next().await.unwrap().unwrap(); + assert_eq!(successor.task.id, next); + assert!(coordinator.request_page(None).await.unwrap().0.is_empty()); + // An empty pass with no new demand must not create a zero-work chain. + finish_task( + &state, + successor, + Ok("local WAL merge reclaimed 0 generations".into()), + ) + .await + .unwrap(); + assert!(state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap() + .is_none()); + } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_completion_does_not_fast_retry_failures_or_legacy_work() { + for resident in [false, true] { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + if resident { + cfg.local_merge.master_local_merge_targets = vec!["exp".into()]; + cfg.catchup.shards = vec!["writer".into()]; + } + let state = MasterState::new(cfg).await.unwrap(); + let coordinator = state.task_store.merge_coordinator(); + enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let claim = state.task_store.claim_next().await.unwrap().unwrap(); + coordinator.request_merge("exp").await.unwrap(); + let outcome = if resident { + Err("injected merge failure".into()) + } else { + Ok("legacy RPC succeeded".into()) + }; + finish_task(&state, claim, outcome).await.unwrap(); + assert!(state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap() + .is_none()); + assert!(coordinator + .get_merge_request("exp") + .await + .unwrap() + .is_some()); + } + } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_completion_with_lost_claim_preserves_successor_and_demand() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.local_merge.master_local_merge_targets = vec!["exp".into()]; + cfg.catchup.shards = vec!["writer".into()]; + let state = MasterState::new(cfg).await.unwrap(); + let coordinator = state.task_store.merge_coordinator(); + let task = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let old = state.task_store.claim_next().await.unwrap().unwrap(); + coordinator.request_merge("exp").await.unwrap(); + let demand = coordinator.get_merge_request("exp").await.unwrap().unwrap(); + // Revoke this test's real etcd lease, then elect a replacement for the + // same canonical task before letting the stale executor report success. + let mut client = state.task_store.etcd_client().clone(); + let key = format!("{}/claims/{}", state.config.etcd.prefix(), task.id); + let response = client.get(key, None).await.unwrap(); + let lease = response.kvs().first().unwrap().lease(); + assert_ne!(lease, 0); + client.lease_revoke(lease).await.unwrap(); + // Normal admission recovers orphaned tasks before claiming. + let successor = state.task_store.claim_next().await.unwrap().unwrap(); + assert_eq!(successor.task.id, task.id); + let error = finish_task(&state, old, Ok("stale success".into())) + .await + .unwrap_err(); + assert!(error.to_string().contains("lost its etcd claim"), "{error}"); + assert_eq!( + state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap(), + Some(task.id) + ); + assert_eq!( + coordinator + .get_merge_request("exp") + .await + .unwrap() + .unwrap() + .revision, + demand.revision + ); + state + .task_store + .finish(successor, Ok("replacement completed".into())) + .await + .unwrap(); + } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_target_lookup_preserves_newer_demand_after_enqueue() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(config(&dir)).await.unwrap(); + let coordinator = state.task_store.merge_coordinator(); + coordinator.request_merge("exp-0").await.unwrap(); + assert!(coordinator + .get_merge_request("exp") + .await + .unwrap() + .is_none()); + coordinator.request_merge("exp").await.unwrap(); + let before = coordinator.get_merge_request("exp").await.unwrap().unwrap(); + coordinator.request_merge("exp").await.unwrap(); + let after = coordinator.get_merge_request("exp").await.unwrap().unwrap(); + assert!(after.revision > before.revision); + enqueue_merge_request(&state, &coordinator, &before) + .await + .unwrap(); + assert_eq!( + coordinator + .get_merge_request("exp") + .await + .unwrap() + .unwrap() + .revision, + after.revision + ); + assert!(coordinator + .get_merge_request("exp-0") + .await + .unwrap() + .is_some()); + let first = state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap(); + enqueue_merge_request(&state, &coordinator, &after) + .await + .unwrap(); + assert_eq!( + state + .task_store + .get_active_id(TaskKind::MergeWal, "exp") + .await + .unwrap(), + first + ); + assert!(coordinator + .get_merge_request("exp") + .await + .unwrap() + .is_none()); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn demand_during_running_merge_retains_a_followup_without_stats_sweep() { diff --git a/crates/lance-context-merge/src/rollout.rs b/crates/lance-context-merge/src/rollout.rs index 7ae74cc..856e139 100644 --- a/crates/lance-context-merge/src/rollout.rs +++ b/crates/lance-context-merge/src/rollout.rs @@ -75,6 +75,27 @@ impl Coordinator { Ok(()) } + /// Read one target's demand after completing a resident merge. The normal + /// paged scan remains the crash/retry fallback; this lookup adds no scan. + pub async fn get_merge_request(&self, target: &str) -> Result> { + let response = self + .client + .clone() + .get(self.request_key(target), None) + .await + .map_err(|e| e.to_string())?; + let Some(kv) = response.kvs().first() else { + return Ok(None); + }; + if kv.value() != target.as_bytes() { + return Err("merge request target does not match its key".into()); + } + Ok(Some(MergeRequest { + target: target.to_owned(), + revision: kv.mod_revision(), + })) + } + pub async fn request_page( &self, after: Option<&str>,