diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index a85fb6ee..a9e59ebd 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/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, target: &str) -> Result { +async fn run_legacy(state: &Arc, claim: &TaskClaim) -> Result { 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, 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, + 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 { + 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 { 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 { + 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()); diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index db6dad7d..7389bb78 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -237,6 +237,31 @@ impl Coordinator { { return Err("invalid merge execution reservation".into()); } + self.reserve_record(claim, execution).await + } + + /// Record an unfenced HTTP execution BEFORE sending its request. Protocol 0 + /// cannot be recovered using protocol 2 storage barriers. Running is + /// deliberate: a legacy worker cannot CAS Reserved -> Running, so even a + /// crash before send must conservatively retain ownership. + pub async fn admit_legacy( + &self, + claim: &ClaimProof, + target: &str, + endpoint: &str, + ) -> Result { + let mut execution = Execution::new(target, endpoint, "legacy-http", 1); + execution.protocol = 0; + execution.phase = Phase::Running; + if !self.reserve_record(claim, &execution).await? { + return Err( + "claim lost or execution already present before legacy merge admission".into(), + ); + } + Ok(execution) + } + + async fn reserve_record(&self, claim: &ClaimProof, execution: &Execution) -> Result { let key = execution_key(&self.prefix, &execution.target); self.transact( vec![