From d60b08e8061856d653ce953af64684a16517ac34 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 21 Oct 2026 03:35:03 +0000 Subject: [PATCH 1/4] Recover task lease streams without replacing ownership --- crates/lance-context-master/src/task_store.rs | 306 +++++++++++++++++- 1 file changed, 295 insertions(+), 11 deletions(-) diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 4a9ed7f..55c5358 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -1411,24 +1411,90 @@ impl EtcdTaskStore { async fn start_keepalive(&self, lease_id: i64) -> lance::Result { let mut client = self.client.clone(); - let (mut keeper, mut stream) = client - .lease_keep_alive(lease_id) - .await - .map_err(etcd_error("start lease keepalive"))?; let interval = Duration::from_secs((self.lease_ttl / 3).max(1) as u64); + let ttl = Duration::from_secs(self.lease_ttl as u64); + let started = tokio::time::Instant::now(); + let (mut keeper, mut stream) = + tokio::time::timeout(interval, client.lease_keep_alive(lease_id)) + .await + .map_err(|_| lance::Error::io("start lease keepalive timed out"))? + .map_err(etcd_error("start lease keepalive"))?; + // etcd-client consumes a positive-TTL response during the handshake. + // Measure from request start, not receipt, to exclude network delay. + let mut confirmed_until = started + ttl; let (stop_tx, mut stop_rx) = oneshot::channel(); let task = tokio::spawn(async move { loop { tokio::select! { - _ = &mut stop_rx => break, - _ = tokio::time::sleep(interval) => { - if keeper.keep_alive().await.is_err() { - break; + _ = &mut stop_rx => return, + _ = tokio::time::sleep(interval) => {} + } + let sent = tokio::time::Instant::now(); + if sent >= confirmed_until { + tracing::warn!(lease_id, "task lease renewal confirmation expired"); + return; + } + let response = + tokio::time::timeout_at(confirmed_until.min(sent + interval), async { + keeper.keep_alive().await?; + stream.message().await + }) + .await; + match response { + Ok(Ok(Some(response))) => { + if response.id() != lease_id || response.ttl() <= 0 { + tracing::warn!(lease_id, "task lease is no longer renewable"); + return; } - match tokio::time::timeout(interval, stream.message()).await { - Ok(Ok(Some(response))) if response.ttl() > 0 => {} - _ => break, + confirmed_until = sent + Duration::from_secs(response.ttl() as u64); + continue; + } + failure => tracing::warn!( + lease_id, + ?failure, + "task lease stream interrupted; reconnecting the same lease" + ), + } + + // A broken stream is not proof that the lease expired. Recover + // only this existing lease within its last confirmed lifetime: + // never grant a new lease or recreate any ownership keys here. + loop { + let reconnect_started = tokio::time::Instant::now(); + if reconnect_started >= confirmed_until { + tracing::warn!( + lease_id, + "task lease stream recovery exhausted its confirmation window" + ); + return; + } + let result = tokio::time::timeout_at( + confirmed_until.min(reconnect_started + interval), + client.lease_keep_alive(lease_id), + ) + .await; + if let Ok(Ok((replacement, replacement_stream))) = result { + if replacement.id() != lease_id { + tracing::warn!(lease_id, "task lease reconnect returned another lease"); + return; } + keeper = replacement; + stream = replacement_stream; + confirmed_until = reconnect_started + ttl; + metrics::counter!("master_task_lease_stream_recoveries_total").increment(1); + tracing::info!( + lease_id, + "task lease stream recovered without changing ownership" + ); + break; + } + // Bound reconnect traffic during an etcd outage. The lease + // identity staggers retries without changing their deadline. + let retry_delay = + interval / 2 + Duration::from_millis(lease_id.unsigned_abs() % 500); + tokio::select! { + _ = &mut stop_rx => return, + _ = tokio::time::sleep_until(confirmed_until.min(tokio::time::Instant::now() + retry_delay)) => {} } } } @@ -1970,6 +2036,224 @@ mod tests { use super::*; use tempfile::TempDir; + struct LeaseTestProxy { + endpoint: String, + online: tokio::sync::watch::Sender, + task: JoinHandle<()>, + } + + impl LeaseTestProxy { + async fn start(upstream: &str) -> Self { + let upstream = upstream.strip_prefix("http://").unwrap().to_owned(); + assert!(upstream.starts_with("127.0.0.1:") || upstream.starts_with("localhost:")); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let (online, status) = tokio::sync::watch::channel(true); + let task = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (mut downstream, _) = accepted.unwrap(); + if !*status.borrow() { + continue; + } + let mut status = status.clone(); + let upstream = upstream.clone(); + connections.spawn(async move { + tokio::select! { + _ = status.changed() => {}, + _ = async { + if let Ok(mut upstream) = tokio::net::TcpStream::connect(upstream).await { + let _ = tokio::io::copy_bidirectional(&mut downstream, &mut upstream).await; + } + } => {}, + } + }); + } + _ = connections.join_next(), if !connections.is_empty() => {} + } + } + }); + Self { + endpoint, + online, + task, + } + } + } + + impl Drop for LeaseTestProxy { + fn drop(&mut self) { + self.task.abort(); + } + } + + async fn lease_proxy_store() -> (TempDir, LeaseTestProxy, TaskStore, Client) { + let upstream = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let upstream = upstream.split(',').next().unwrap(); + let proxy = LeaseTestProxy::start(upstream).await; + let observer = Client::connect([upstream], None).await.unwrap(); + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.etcd.etcd_endpoints = vec![proxy.endpoint.clone()]; + cfg.etcd.etcd_prefix = format!("/lease-stream-test/{}", generate_id()); + cfg.etcd_lease_ttl_secs = 6; + let store = TaskStore::open(&cfg).await.unwrap(); + (dir, proxy, store, observer) + } + + #[tokio::test] + #[ignore = "requires loopback ETCD_TEST_ENDPOINTS"] + async fn task_lease_stream_reconnect_preserves_the_same_claim() { + let (_dir, proxy, store, mut observer) = lease_proxy_store().await; + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + let claim = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let lease = claim.backend.lease_id; + let key = claim.backend.claim_key.clone(); + let token = claim.backend.token.clone(); + proxy.online.send(false).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + proxy.online.send(true).unwrap(); + // Survive more than the original lease lifetime after a real TCP cut. + tokio::time::sleep(Duration::from_secs(10)).await; + let values = observer.get(key.clone(), None).await.unwrap(); + assert_eq!( + values.kvs().len(), + 1, + "a transient stream failure expired the claim" + ); + assert_eq!(values.kvs()[0].value(), token.as_bytes()); + assert_eq!( + values.kvs()[0].lease(), + lease, + "recovery must not replace the lease" + ); + assert!(!claim.backend.keepalive.task.is_finished()); + // Dropping a claim must still stop renewal, including after reconnect. + drop(claim); + tokio::time::timeout(Duration::from_secs(9), async { + while !observer + .get(key.clone(), None) + .await + .unwrap() + .kvs() + .is_empty() + { + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .unwrap(); + observer + .delete( + store.inner.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + + #[tokio::test] + #[ignore = "requires loopback ETCD_TEST_ENDPOINTS"] + async fn task_lease_stream_recovery_stops_after_a_persistent_outage() { + let (_dir, proxy, store, mut observer) = lease_proxy_store().await; + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + let claim = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let key = claim.backend.claim_key.clone(); + proxy.online.send(false).unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while !claim.backend.keepalive.task.is_finished() + || !observer + .get(key.clone(), None) + .await + .unwrap() + .kvs() + .is_empty() + { + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .unwrap(); + proxy.online.send(true).unwrap(); + tokio::time::sleep(Duration::from_secs(2)).await; + assert!(observer.get(key, None).await.unwrap().kvs().is_empty()); + drop(claim); + observer + .delete( + store.inner.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + + #[tokio::test] + #[ignore = "requires loopback ETCD_TEST_ENDPOINTS"] + async fn task_lease_stream_recovery_never_resurrects_revoked_ownership() { + let (_dir, proxy, store, mut observer) = lease_proxy_store().await; + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + let claim = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let lease = claim.backend.lease_id; + let key = claim.backend.claim_key.clone(); + proxy.online.send(false).unwrap(); + observer.lease_revoke(lease).await.unwrap(); + observer + .put(key.clone(), "replacement-owner", None) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + proxy.online.send(true).unwrap(); + tokio::time::timeout(Duration::from_secs(9), async { + while !claim.backend.keepalive.task.is_finished() { + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .unwrap(); + let values = observer.get(key, None).await.unwrap(); + assert_eq!(values.kvs()[0].value(), b"replacement-owner"); + assert_eq!(values.kvs()[0].lease(), 0); + assert!( + observer + .lease_time_to_live(lease, None) + .await + .unwrap() + .ttl() + <= 0 + ); + drop(claim); + observer + .delete( + store.inner.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[test] fn completion_cleanup_never_changes_the_transaction_outcome() { let task = new_task(TaskKind::MergeWal, "hot", Vec::new()); From 17bae86d42230e39e6097c5dbcedfe70d7bb6ade Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 21 Oct 2026 03:42:42 +0000 Subject: [PATCH 2/4] Verify fault proxy restores connectivity before lease assertions --- crates/lance-context-master/src/task_store.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 55c5358..d4f6f1c 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -2048,14 +2048,14 @@ mod tests { assert!(upstream.starts_with("127.0.0.1:") || upstream.starts_with("localhost:")); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let endpoint = format!("http://{}", listener.local_addr().unwrap()); - let (online, status) = tokio::sync::watch::channel(true); + let (online, mut status) = tokio::sync::watch::channel(true); let task = tokio::spawn(async move { let mut connections = tokio::task::JoinSet::new(); loop { tokio::select! { accepted = listener.accept() => { let (mut downstream, _) = accepted.unwrap(); - if !*status.borrow() { + if !*status.borrow_and_update() { continue; } let mut status = status.clone(); @@ -2122,6 +2122,17 @@ mod tests { proxy.online.send(false).unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; proxy.online.send(true).unwrap(); + // Prove that the proxy really restored transport before checking lease + // recovery. New connections must not inherit an unseen old cut event. + tokio::time::timeout(Duration::from_secs(2), async { + let mut probe = Client::connect([proxy.endpoint.as_str()], None) + .await + .unwrap(); + let values = probe.get(key.clone(), None).await.unwrap(); + assert_eq!(values.kvs()[0].value(), token.as_bytes()); + }) + .await + .expect("proxy did not restore etcd connectivity"); // Survive more than the original lease lifetime after a real TCP cut. tokio::time::sleep(Duration::from_secs(10)).await; let values = observer.get(key.clone(), None).await.unwrap(); From 25dccc595db2282ff2fb3172bf5e19adcbcbeed5 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 21 Oct 2026 04:19:03 +0000 Subject: [PATCH 3/4] Recover task leases after delayed renewal acknowledgements --- crates/lance-context-master/src/task_store.rs | 130 ++++++++++++++---- 1 file changed, 106 insertions(+), 24 deletions(-) diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index d4f6f1c..4454df9 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -1413,15 +1413,12 @@ impl EtcdTaskStore { let mut client = self.client.clone(); let interval = Duration::from_secs((self.lease_ttl / 3).max(1) as u64); let ttl = Duration::from_secs(self.lease_ttl as u64); - let started = tokio::time::Instant::now(); let (mut keeper, mut stream) = tokio::time::timeout(interval, client.lease_keep_alive(lease_id)) .await .map_err(|_| lance::Error::io("start lease keepalive timed out"))? .map_err(etcd_error("start lease keepalive"))?; // etcd-client consumes a positive-TTL response during the handshake. - // Measure from request start, not receipt, to exclude network delay. - let mut confirmed_until = started + ttl; let (stop_tx, mut stop_rx) = oneshot::channel(); let task = tokio::spawn(async move { loop { @@ -1429,24 +1426,17 @@ impl EtcdTaskStore { _ = &mut stop_rx => return, _ = tokio::time::sleep(interval) => {} } - let sent = tokio::time::Instant::now(); - if sent >= confirmed_until { - tracing::warn!(lease_id, "task lease renewal confirmation expired"); - return; - } - let response = - tokio::time::timeout_at(confirmed_until.min(sent + interval), async { - keeper.keep_alive().await?; - stream.message().await - }) - .await; + let response = tokio::time::timeout(interval, async { + keeper.keep_alive().await?; + stream.message().await + }) + .await; match response { Ok(Ok(Some(response))) => { if response.id() != lease_id || response.ttl() <= 0 { tracing::warn!(lease_id, "task lease is no longer renewable"); return; } - confirmed_until = sent + Duration::from_secs(response.ttl() as u64); continue; } failure => tracing::warn!( @@ -1456,20 +1446,24 @@ impl EtcdTaskStore { ), } - // A broken stream is not proof that the lease expired. Recover - // only this existing lease within its last confirmed lifetime: - // never grant a new lease or recreate any ownership keys here. + // A missing response is not proof that the server failed to + // renew. Even beyond our last confirmation, etcd may still + // hold this lease. Bound recovery from the stream failure, + // letting etcd decide whether this same lease is renewable. + // This never grants a lease or recreates ownership keys, and + // renewal does not replace a task's publication ownership checks. + let recovery_until = tokio::time::Instant::now() + ttl; loop { let reconnect_started = tokio::time::Instant::now(); - if reconnect_started >= confirmed_until { + if reconnect_started >= recovery_until { tracing::warn!( lease_id, - "task lease stream recovery exhausted its confirmation window" + "task lease stream recovery exhausted its retry window" ); return; } let result = tokio::time::timeout_at( - confirmed_until.min(reconnect_started + interval), + recovery_until.min(reconnect_started + interval), client.lease_keep_alive(lease_id), ) .await; @@ -1480,7 +1474,6 @@ impl EtcdTaskStore { } keeper = replacement; stream = replacement_stream; - confirmed_until = reconnect_started + ttl; metrics::counter!("master_task_lease_stream_recoveries_total").increment(1); tracing::info!( lease_id, @@ -1494,7 +1487,7 @@ impl EtcdTaskStore { interval / 2 + Duration::from_millis(lease_id.unsigned_abs() % 500); tokio::select! { _ = &mut stop_rx => return, - _ = tokio::time::sleep_until(confirmed_until.min(tokio::time::Instant::now() + retry_delay)) => {} + _ = tokio::time::sleep_until(recovery_until.min(tokio::time::Instant::now() + retry_delay)) => {} } } } @@ -2039,6 +2032,7 @@ mod tests { struct LeaseTestProxy { endpoint: String, online: tokio::sync::watch::Sender, + responses: tokio::sync::watch::Sender, task: JoinHandle<()>, } @@ -2049,6 +2043,7 @@ mod tests { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let endpoint = format!("http://{}", listener.local_addr().unwrap()); let (online, mut status) = tokio::sync::watch::channel(true); + let (responses, response_status) = tokio::sync::watch::channel(true); let task = tokio::spawn(async move { let mut connections = tokio::task::JoinSet::new(); loop { @@ -2059,13 +2054,33 @@ mod tests { continue; } let mut status = status.clone(); + let mut response_status = response_status.clone(); let upstream = upstream.clone(); connections.spawn(async move { tokio::select! { _ = status.changed() => {}, _ = async { if let Ok(mut upstream) = tokio::net::TcpStream::connect(upstream).await { - let _ = tokio::io::copy_bidirectional(&mut downstream, &mut upstream).await; + let (mut client_read, mut client_write) = downstream.split(); + let (mut server_read, mut server_write) = upstream.split(); + let replies = async { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mut buffer = [0; 16 * 1024]; + loop { + let count = server_read.read(&mut buffer).await?; + if count == 0 { + return Ok::<(), std::io::Error>(()); + } + if response_status.wait_for(|enabled| *enabled).await.is_err() { + return Ok(()); + } + client_write.write_all(&buffer[..count]).await?; + } + }; + tokio::select! { + _ = tokio::io::copy(&mut client_read, &mut server_write) => {}, + _ = replies => {}, + } } } => {}, } @@ -2078,6 +2093,7 @@ mod tests { Self { endpoint, online, + responses, task, } } @@ -2090,6 +2106,10 @@ mod tests { } async fn lease_proxy_store() -> (TempDir, LeaseTestProxy, TaskStore, Client) { + let _ = tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .with_test_writer() + .try_init(); let upstream = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); let upstream = upstream.split(',').next().unwrap(); let proxy = LeaseTestProxy::start(upstream).await; @@ -2172,6 +2192,68 @@ mod tests { .unwrap(); } + #[tokio::test] + #[ignore = "requires loopback ETCD_TEST_ENDPOINTS"] + async fn task_lease_stream_recovers_when_renewal_replies_are_delayed() { + let (_dir, proxy, store, mut observer) = lease_proxy_store().await; + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + let claim = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let key = claim.backend.claim_key.clone(); + let token = claim.backend.token.clone(); + let lease = claim.backend.lease_id; + // Let the first ordinary renewal complete. Then forward requests to + // etcd while withholding responses: an unacknowledged renewal can + // extend the real lease beyond the client's last confirmation. + tokio::time::sleep(Duration::from_secs(3)).await; + proxy.responses.send(false).unwrap(); + tokio::time::sleep(Duration::from_secs(6)).await; + let values = observer.get(key.clone(), None).await.unwrap(); + assert_eq!( + values.kvs().len(), + 1, + "fixture must retain a live server lease" + ); + assert_eq!(values.kvs()[0].value(), token.as_bytes()); + assert_eq!(values.kvs()[0].lease(), lease); + assert!( + observer + .lease_time_to_live(lease, None) + .await + .unwrap() + .ttl() + > 0 + ); + proxy.responses.send(true).unwrap(); + assert!(!claim.backend.keepalive.task.is_finished(), + "unacknowledged renewals must be retried on the same lease within a bounded recovery window"); + tokio::time::sleep(Duration::from_secs(8)).await; + let values = observer.get(key, None).await.unwrap(); + assert_eq!( + values.kvs().len(), + 1, + "restored responses did not resume renewal" + ); + assert_eq!(values.kvs()[0].value(), token.as_bytes()); + assert_eq!(values.kvs()[0].lease(), lease); + assert!(!claim.backend.keepalive.task.is_finished()); + drop(claim); + observer.lease_revoke(lease).await.unwrap(); + observer + .delete( + store.inner.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[tokio::test] #[ignore = "requires loopback ETCD_TEST_ENDPOINTS"] async fn task_lease_stream_recovery_stops_after_a_persistent_outage() { From 970b07ac5ef5f7eada84d7007fbef676d8e80586 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 21 Oct 2026 04:27:51 +0000 Subject: [PATCH 4/4] Separate lease fault timing from delayed observer reads --- crates/lance-context-master/src/task_store.rs | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 4454df9..9f82aa5 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -2214,6 +2214,10 @@ mod tests { tokio::time::sleep(Duration::from_secs(3)).await; proxy.responses.send(false).unwrap(); tokio::time::sleep(Duration::from_secs(6)).await; + // Restore responses before any observer RPC: a slow linearizable + // read must not silently extend the injected six-second outage. + let renewal_still_running = !claim.backend.keepalive.task.is_finished(); + proxy.responses.send(true).unwrap(); let values = observer.get(key.clone(), None).await.unwrap(); assert_eq!( values.kvs().len(), @@ -2230,8 +2234,7 @@ mod tests { .ttl() > 0 ); - proxy.responses.send(true).unwrap(); - assert!(!claim.backend.keepalive.task.is_finished(), + assert!(renewal_still_running, "unacknowledged renewals must be retried on the same lease within a bounded recovery window"); tokio::time::sleep(Duration::from_secs(8)).await; let values = observer.get(key, None).await.unwrap(); @@ -2270,13 +2273,22 @@ mod tests { let key = claim.backend.claim_key.clone(); proxy.online.send(false).unwrap(); tokio::time::timeout(Duration::from_secs(10), async { - while !claim.backend.keepalive.task.is_finished() - || !observer - .get(key.clone(), None) - .await - .unwrap() - .kvs() - .is_empty() + while !claim.backend.keepalive.task.is_finished() { + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("renewal task exceeded its bounded recovery window"); + // Client termination and server-side lease deletion are separate: + // etcd may delay its expiry transaction or this linearizable read. + // A slow observer must not mask whether the client stopped in time. + tokio::time::timeout(Duration::from_secs(10), async { + while !observer + .get(key.clone(), None) + .await + .unwrap() + .kvs() + .is_empty() { tokio::time::sleep(Duration::from_millis(100)).await; }