diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 0fed3217..5fea2ec1 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -658,6 +658,23 @@ async fn should_probe_failure( false } +async fn enqueue_merge_request( + state: &Arc, + coordinator: &lance_context_merge::Coordinator, + request: &lance_context_merge::rollout::MergeRequest, +) -> Result<(), String> { + if !state.config.merge_rollout.owned(&request.target) { + return Ok(()); + } + let task = enqueue(state, TaskKind::MergeWal, &request.target) + .await + .map_err(|e| e.to_string())?; + if task.state == lance_context_api::TaskState::Queued { + coordinator.acknowledge_request(request).await?; + } + Ok(()) +} + /// Spawn the scheduler pollers plus the optional periodic auto-sweep. /// /// Returns the handle of the *general* poller only. When a separate WAL-merge @@ -680,16 +697,11 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { match coordinator.request_page(request_cursor.as_deref()).await { Ok((targets, next)) => { request_cursor = next; - for target in targets { - if retry_state.config.merge_rollout.owned(&target) { - match enqueue(&retry_state, TaskKind::MergeWal, &target).await { - Ok(_) => { - let _ = coordinator.acknowledge_request(&target).await; - } - Err(error) => { - tracing::warn!(%target, %error, "worker merge demand enqueue failed") - } - } + for request in targets { + if let Err(error) = + enqueue_merge_request(&retry_state, &coordinator, &request).await + { + tracing::warn!(target = %request.target, %error, "worker merge demand enqueue failed"); } } } @@ -1026,6 +1038,41 @@ mod tests { ); } + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn demand_during_running_merge_retains_a_followup_without_stats_sweep() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.merge_rollout.owned_targets = vec!["hot".into()]; + let state = MasterState::new(cfg).await.unwrap(); + let coordinator = state.task_store.merge_coordinator(); + let first = enqueue(&state, TaskKind::MergeWal, "hot").await.unwrap(); + let claim = state.task_store.claim_next().await.unwrap().unwrap(); + assert_eq!(claim.task.id, first.id); + coordinator.request_merge("hot").await.unwrap(); + let request = coordinator.request_page(None).await.unwrap().0.remove(0); + enqueue_merge_request(&state, &coordinator, &request) + .await + .unwrap(); + assert_eq!(coordinator.request_page(None).await.unwrap().0.len(), 1); + state + .task_store + .finish(claim, Ok("old pass complete".into())) + .await + .unwrap(); + enqueue_merge_request(&state, &coordinator, &request) + .await + .unwrap(); + let next = state + .task_store + .get_active_id(TaskKind::MergeWal, "hot") + .await + .unwrap() + .unwrap(); + assert_ne!(first.id, next); + assert!(coordinator.request_page(None).await.unwrap().0.is_empty()); + } + /// Wait until the task reaches a terminal state, returning its final record. async fn await_terminal(state: &Arc, id: &str) -> TaskRecord { for _ in 0..100 { diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index 87b193d5..377f1276 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -501,7 +501,8 @@ mod tests { coordinator.request_merge("table").await.unwrap(); } let (rows, next) = coordinator.request_page(None).await.unwrap(); - assert_eq!(rows, ["table"]); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].target, "table"); assert!(next.is_none()); let after = coordinator .failure("table", "worker") @@ -511,10 +512,17 @@ mod tests { let before = last.unwrap(); assert_eq!(after.consecutive_attempts, 4); assert_eq!(after.next_retry_ms, before.next_retry_ms); - coordinator.acknowledge_request("table").await.unwrap(); + // A later flush must survive the old demand acknowledgement. + coordinator.request_merge("table").await.unwrap(); + assert!(!coordinator.acknowledge_request(&rows[0]).await.unwrap()); + let newer = coordinator.request_page(None).await.unwrap().0; + assert!(coordinator.acknowledge_request(&newer[0]).await.unwrap()); assert!(coordinator.request_page(None).await.unwrap().0.is_empty()); coordinator.request_merge("table").await.unwrap(); - assert_eq!(coordinator.request_page(None).await.unwrap().0, ["table"]); + assert_eq!( + coordinator.request_page(None).await.unwrap().0[0].target, + "table" + ); } #[tokio::test] diff --git a/crates/lance-context-merge/src/rollout.rs b/crates/lance-context-merge/src/rollout.rs index 9f0fcbe1..30d2571d 100644 --- a/crates/lance-context-merge/src/rollout.rs +++ b/crates/lance-context-merge/src/rollout.rs @@ -47,6 +47,14 @@ impl MergeRollout { } } +/// One coalesced demand snapshot. Its revision identifies exactly what a +/// scheduler may acknowledge; later flushes must survive an older enqueue. +#[derive(Clone, Debug)] +pub struct MergeRequest { + pub target: String, + pub revision: i64, +} + impl Coordinator { fn request_key(&self, target: &str) -> String { execution_key(&self.prefix, target).replace("/merge-executions/", "/merge-requests/") @@ -55,15 +63,18 @@ impl Coordinator { /// Coalesce worker count/timer triggers without resetting failure budgets. pub async fn request_merge(&self, target: &str) -> Result<()> { let key = self.request_key(target); - self.transact( - vec![Compare::version(key.as_str(), CompareOp::Equal, 0)], - vec![TxnOp::put(key, target, None)], - ) - .await?; + self.client + .clone() + .put(key, target, None) + .await + .map_err(|e| e.to_string())?; Ok(()) } - pub async fn request_page(&self, after: Option<&str>) -> Result<(Vec, Option)> { + pub async fn request_page( + &self, + after: Option<&str>, + ) -> Result<(Vec, Option)> { let prefix = format!("{}/merge-requests/", self.prefix.trim_end_matches('/')); let (start, options) = match after { None => (prefix.clone(), GetOptions::new().with_prefix()), @@ -83,7 +94,10 @@ impl Coordinator { let rows = response .kvs() .iter() - .map(|kv| String::from_utf8_lossy(kv.value()).into_owned()) + .map(|kv| MergeRequest { + target: String::from_utf8_lossy(kv.value()).into_owned(), + revision: kv.mod_revision(), + }) .collect(); let next = response .more() @@ -97,15 +111,20 @@ impl Coordinator { Ok((rows, next)) } - /// Called only after a durable enqueue. A subsequent flush tick reasserts - /// demand if the active task already passed this worker's shard. - pub async fn acknowledge_request(&self, target: &str) -> Result<()> { - self.client - .clone() - .delete(self.request_key(target), None) - .await - .map_err(|e| e.to_string())?; - Ok(()) + /// Only acknowledge a queued task, never an already-running pass that may + /// have visited this shard. A concurrent flush advances the revision and + /// survives this compare-and-delete even if enqueue/start raced it. + pub async fn acknowledge_request(&self, request: &MergeRequest) -> Result { + let key = self.request_key(&request.target); + self.transact( + vec![Compare::mod_revision( + key.clone(), + CompareOp::Equal, + request.revision, + )], + vec![TxnOp::delete(key, None)], + ) + .await } } diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index 24274752..79aeeef2 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -434,7 +434,14 @@ mod tests { ) .await; assert_eq!( - coordinator.request_page(None).await.unwrap().0, + coordinator + .request_page(None) + .await + .unwrap() + .0 + .into_iter() + .map(|r| r.target) + .collect::>(), ["generic:s"] ); assert_eq!(