fix: preserve merge demand across running passes and racing flushes

crates/lance-context-master/src/scheduler.rs

@@ -658,6 +658,23 @@ async fn should_probe_failure(
     false
 }
 
+async fn enqueue_merge_request(
+    state: &Arc<MasterState>,
+    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<MasterState>) -> 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<MasterState>, id: &str) -> TaskRecord {
         for _ in 0..100 {

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]

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<String>, Option<String>)> {
+    pub async fn request_page(
+        &self,
+        after: Option<&str>,
+    ) -> Result<(Vec<MergeRequest>, Option<String>)> {
         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<bool> {
+        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
     }
 }
 

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::<Vec<_>>(),
             ["generic:s"]
         );
         assert_eq!(