fix(master): retain ownership across ambiguous legacy merge replies

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

@@ -56,6 +56,10 @@ pub(crate) async fn run_merge_wal(
         if let Some(old) = coordinator.get(target).await? {
             // Turning off new admission must never disable recovery of durable
             // worker records. Observe useful progress before considering cancel.
+            if old.protocol == 0 && old.phase == Phase::Finished {
+                recover_execution(state, &coordinator, &proof, old).await?;
+                return Ok("released acknowledged legacy merge; new pass may be scheduled".into());
+            }
             if old.maintenance.is_some() || old.protocol != 2 {
                 return Err("owned execution still present; refusing legacy downgrade".into());
             }
@@ -70,7 +74,7 @@ pub(crate) async fn run_merge_wal(
             }
             return Ok("reconciled recorded worker merge; new fenced admission disabled".into());
         }
-        return run_legacy(state, target).await;
+        return run_legacy(state, claim).await;
     }
     if state.config.worker_endpoints.is_empty()
         && !crate::local_merge::selected(&state.config, target)
@@ -150,17 +154,88 @@ pub(crate) async fn run_merge_wal(
 
 // Compatibility path is chosen only by explicit configuration, never after a
 // failed owned RPC. It retains legacy limitations until that table is drained.
-async fn run_legacy(state: &Arc<MasterState>, target: &str) -> Result<String, String> {
+async fn run_legacy(state: &Arc<MasterState>, claim: &TaskClaim) -> Result<String, String> {
     if state.config.worker_endpoints.is_empty() {
         return Err("no worker endpoints configured".into());
     }
-    run_legacy_workers(&state.http, &state.config.worker_endpoints, target).await
+    let coordinator = state.task_store.merge_coordinator();
+    let proof = state.task_store.merge_claim(claim);
+    run_legacy_workers_inner(
+        &state.http,
+        &state.config.worker_endpoints,
+        &claim.task.target,
+        LegacyOwnership::Durable {
+            coordinator: &coordinator,
+            proof: &proof,
+        },
+    )
+    .await
+}
+
+// The untracked variant exists only for HTTP classification unit tests. Every
+// production call must carry the canonical claim and durable coordinator.
+enum LegacyOwnership<'a> {
+    Durable {
+        coordinator: &'a Coordinator,
+        proof: &'a ClaimProof,
+    },
+    #[cfg(test)]
+    HttpTest,
+}
+
+impl LegacyOwnership<'_> {
+    async fn admit(&self, target: &str, endpoint: &str) -> Result<Option<Execution>, String> {
+        match self {
+            Self::Durable { coordinator, proof } => coordinator
+                .admit_legacy(proof, target, endpoint)
+                .await
+                .map(Some),
+            #[cfg(test)]
+            Self::HttpTest => Ok(None),
+        }
+    }
+
+    async fn acknowledged(
+        &self,
+        execution: Option<Execution>,
+        reclaimed: usize,
+    ) -> Result<(), String> {
+        match self {
+            Self::Durable { coordinator, proof } => {
+                let running = execution.ok_or("missing legacy ownership record")?;
+                if !coordinator.finish(&running, Ok(reclaimed)).await? {
+                    return Err(
+                        "legacy completion ownership changed; refusing further admission".into(),
+                    );
+                }
+                if !coordinator
+                    .release(proof, &running.finished(Ok(reclaimed)))
+                    .await?
+                {
+                    return Err("claim lost before releasing acknowledged legacy execution".into());
+                }
+                Ok(())
+            }
+            #[cfg(test)]
+            Self::HttpTest => Ok(()),
+        }
+    }
 }
 
+#[cfg(test)]
 async fn run_legacy_workers(
     http: &reqwest::Client,
     endpoints: &[String],
     target: &str,
+) -> Result<String, String> {
+    run_legacy_workers_inner(http, endpoints, target, LegacyOwnership::HttpTest).await
+}
+
+async fn run_legacy_workers_inner(
+    http: &reqwest::Client,
+    endpoints: &[String],
+    target: &str,
+    ownership: LegacyOwnership<'_>,
 ) -> Result<String, String> {
     let mut reclaimed = 0u64;
     let mut completed = 0usize;
@@ -171,6 +246,7 @@ async fn run_legacy_workers(
             Some(name) => format!("{endpoint}/api/v1/generic/{name}/merge-wal"),
             None => format!("{endpoint}/api/v1/internal/merge-wal/{target}"),
         };
+        let execution = ownership.admit(target, endpoint).await?;
         let started = std::time::Instant::now();
         let result = async {
             let response = http.post(url).send().await.map_err(|error| {
@@ -203,12 +279,15 @@ async fn run_legacy_workers(
         metrics::counter!("master_merge_wal_workers_total", "result" => if result.is_ok() { "ok" } else { "failed" }).increment(1);
         match result {
             Ok(n) => {
+                ownership.acknowledged(execution, n as usize).await?;
                 completed += 1;
                 reclaimed += n;
                 // Preserve real progress even when a later worker is unavailable.
                 metrics::counter!("master_merge_wal_generations_reclaimed_total").increment(n);
             }
             Err((true, error)) => {
+                // Redirects are disabled; this connection never admitted work.
+                ownership.acknowledged(execution, 0).await?;
                 tracing::warn!(%endpoint, %target, %error, "worker connection failed before merge admission; continuing healthy workers");
                 unavailable.push(format!("{endpoint}: {error}"));
             }
@@ -588,6 +667,9 @@ async fn reconcile_with_grace(
     cancel_immediately: bool,
     grace: Duration,
 ) -> Result<usize, String> {
+    if initial.protocol != 2 {
+        return Err("legacy merge ownership unresolved; terminal acknowledgement or verified retirement required".into());
+    }
     let mut deadlines = Deadlines::new(&initial, tokio::time::Instant::now());
     let mut cancel_started = cancel_immediately.then(tokio::time::Instant::now);
     let mut next_cancel = tokio::time::Instant::now();
@@ -1332,6 +1414,228 @@ mod legacy_connection_tests {
         Arc,
     };
 
+    async fn durable_store() -> (tempfile::TempDir, crate::task_store::TaskStore, String) {
+        use clap::Parser;
+        let dir = tempfile::tempdir().unwrap();
+        let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap();
+        assert!(
+            endpoint.starts_with("http://127.0.0.1:") || endpoint.starts_with("http://localhost:")
+        );
+        let mut cfg = crate::config::MasterConfig::parse_from(["test"]);
+        cfg.data_dir = dir.path().to_string_lossy().into();
+        cfg.etcd.etcd_endpoints = vec![endpoint];
+        cfg.etcd.etcd_prefix = format!(
+            "/legacy-ownership-tests/{}",
+            lance_context_core::generate_id()
+        );
+        let store = crate::task_store::TaskStore::open(&cfg).await.unwrap();
+        (dir, store, cfg.etcd.etcd_prefix)
+    }
+
+    #[tokio::test]
+    #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
+    async fn legacy_timeout_retains_ownership_after_task_finish_and_blocks_replay() {
+        use super::{run_legacy_workers_inner, LegacyOwnership};
+        use crate::task_store::TaskKinds;
+        use lance_context_api::TaskKind;
+        let (_dir, store, prefix) = durable_store().await;
+        let calls = Arc::new(AtomicUsize::new(0));
+        let observed = calls.clone();
+        let release = Arc::new(tokio::sync::Notify::new());
+        let barrier = release.clone();
+        let (endpoint, server) = serve(Router::new().fallback(post(move || {
+            let observed = observed.clone();
+            let barrier = barrier.clone();
+            async move {
+                observed.fetch_add(1, Ordering::SeqCst);
+                barrier.notified().await;
+                Json(serde_json::json!({"reclaimed": 3}))
+            }
+        })))
+        .await;
+        store
+            .enqueue(TaskKind::MergeWal, "table", vec![])
+            .await
+            .unwrap();
+        let claim = store
+            .claim_next_of_kinds(TaskKinds::MERGE_WAL)
+            .await
+            .unwrap()
+            .unwrap();
+        let coordinator = store.merge_coordinator();
+        let proof = store.merge_claim(&claim);
+        let http = reqwest::Client::builder()
+            .redirect(reqwest::redirect::Policy::none())
+            .timeout(std::time::Duration::from_millis(200))
+            .build()
+            .unwrap();
+        let outcome = run_legacy_workers_inner(
+            &http,
+            std::slice::from_ref(&endpoint),
+            "table",
+            LegacyOwnership::Durable {
+                coordinator: &coordinator,
+                proof: &proof,
+            },
+        )
+        .await;
+        assert!(outcome.as_ref().unwrap_err().contains("stopped at"));
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+        let running = coordinator.get("table").await.unwrap().unwrap();
+        assert_eq!(running.protocol, 0);
+        assert_eq!(running.phase, lance_context_merge::Phase::Running);
+        store.finish(claim, outcome).await.unwrap();
+        let lock = store
+            .etcd_client()
+            .clone()
+            .get(lance_context_merge::target_lock_key(&prefix, "table"), None)
+            .await
+            .unwrap();
+        assert_eq!(
+            lock.kvs()[0].lease(),
+            0,
+            "ambiguous writer protection must outlive task lease"
+        );
+        assert_eq!(
+            lock.kvs()[0].value(),
+            lance_context_merge::execution_owner(&running).as_bytes()
+        );
+        store
+            .enqueue(TaskKind::MergeWal, "table", vec![])
+            .await
+            .unwrap();
+        let retry = store
+            .claim_next_of_kinds(TaskKinds::MERGE_WAL)
+            .await
+            .unwrap()
+            .unwrap();
+        let retry_proof = store.merge_claim(&retry);
+        let retry_result = run_legacy_workers_inner(
+            &http,
+            &[endpoint],
+            "table",
+            LegacyOwnership::Durable {
+                coordinator: &coordinator,
+                proof: &retry_proof,
+            },
+        )
+        .await;
+        assert!(retry_result
+            .unwrap_err()
+            .contains("execution already present"));
+        assert_eq!(
+            calls.load(Ordering::SeqCst),
+            1,
+            "reconciler claim cannot POST another writer"
+        );
+        assert!(
+            coordinator.freeze(&retry_proof, &running).await.is_err(),
+            "protocol 2 barriers cannot retire legacy work"
+        );
+        store
+            .finish(retry, Err("legacy ownership unresolved".into()))
+            .await
+            .unwrap();
+        assert_eq!(coordinator.get("table").await.unwrap(), Some(running));
+        release.notify_one();
+        server.abort();
+    }
+
+    #[tokio::test]
+    #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
+    async fn legacy_success_releases_each_endpoint_and_task_lock() {
+        use super::{run_legacy_workers_inner, LegacyOwnership};
+        use crate::task_store::TaskKinds;
+        use lance_context_api::TaskKind;
+        let (_dir, store, prefix) = durable_store().await;
+        let calls = Arc::new(AtomicUsize::new(0));
+        let (endpoint, server) = healthy(calls.clone()).await;
+        store
+            .enqueue(TaskKind::MergeWal, "table", vec![])
+            .await
+            .unwrap();
+        let claim = store
+            .claim_next_of_kinds(TaskKinds::MERGE_WAL)
+            .await
+            .unwrap()
+            .unwrap();
+        let coordinator = store.merge_coordinator();
+        let proof = store.merge_claim(&claim);
+        let result = run_legacy_workers_inner(
+            &crate::state::worker_http_client().unwrap(),
+            &[endpoint.clone(), endpoint],
+            "table",
+            LegacyOwnership::Durable {
+                coordinator: &coordinator,
+                proof: &proof,
+            },
+        )
+        .await;
+        assert!(result.as_ref().unwrap().contains("merged 14 generations"));
+        assert_eq!(calls.load(Ordering::SeqCst), 2);
+        assert!(coordinator.get("table").await.unwrap().is_none());
+        store.finish(claim, result).await.unwrap();
+        assert!(store
+            .etcd_client()
+            .clone()
+            .get(lance_context_merge::target_lock_key(&prefix, "table"), None)
+            .await
+            .unwrap()
+            .kvs()
+            .is_empty());
+        server.abort();
+    }
+
+    #[tokio::test]
+    #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
+    async fn legacy_admission_survives_claim_lease_loss_before_http() {
+        use crate::task_store::TaskKinds;
+        use lance_context_api::TaskKind;
+        let (_dir, store, prefix) = durable_store().await;
+        store
+            .enqueue(TaskKind::MergeWal, "table", vec![])
+            .await
+            .unwrap();
+        let claim = store
+            .claim_next_of_kinds(TaskKinds::MERGE_WAL)
+            .await
+            .unwrap()
+            .unwrap();
+        let proof = store.merge_claim(&claim);
+        let coordinator = store.merge_coordinator();
+        let running = coordinator
+            .admit_legacy(&proof, "table", "http://uncontacted-worker")
+            .await
+            .unwrap();
+        store
+            .etcd_client()
+            .clone()
+            .lease_revoke(proof.lease_id)
+            .await
+            .unwrap();
+        assert_eq!(
+            coordinator.get("table").await.unwrap(),
+            Some(running.clone())
+        );
+        assert!(!coordinator.cancel_reserved(&running).await.unwrap());
+        assert!(coordinator.freeze(&proof, &running).await.is_err());
+        let lock = store
+            .etcd_client()
+            .clone()
+            .get(lance_context_merge::target_lock_key(&prefix, "table"), None)
+            .await
+            .unwrap();
+        assert_eq!(lock.kvs()[0].lease(), 0);
+        assert_eq!(
+            lock.kvs()[0].value(),
+            lance_context_merge::execution_owner(&running).as_bytes()
+        );
+        assert!(coordinator
+            .admit_legacy(&proof, "other", "http://uncontacted-worker")
+            .await
+            .is_err());
+    }
+
     async fn serve(app: Router) -> (String, tokio::task::JoinHandle<()>) {
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
         let url = format!("http://{}", listener.local_addr().unwrap());