From 0c52334480516082d0914e5ab0f7deff1808518f Mon Sep 17 00:00:00 2001 From: Henry Guo Date: Thu, 28 May 2026 01:55:46 +0800 Subject: [PATCH] fix(lock): retry transient distributed lock timeouts (#3101) * fix(lock): retry transient distributed lock timeouts * fix(lock): avoid retry during pending lock cleanup * fix(lock): preserve acquire timeout budget --------- Co-authored-by: Henry Guo Co-authored-by: houseme --- crates/ecstore/src/rpc/remote_locker.rs | 39 +++++-- crates/lock/src/distributed_lock.rs | 120 ++++++++++++++++++--- crates/lock/src/namespace/tests.rs | 132 ++++++++++++++++++++++++ 3 files changed, 269 insertions(+), 22 deletions(-) diff --git a/crates/ecstore/src/rpc/remote_locker.rs b/crates/ecstore/src/rpc/remote_locker.rs index 82bd84a71..9cea3154a 100644 --- a/crates/ecstore/src/rpc/remote_locker.rs +++ b/crates/ecstore/src/rpc/remote_locker.rs @@ -101,7 +101,12 @@ impl RemoteClient { } Err(_) => { let reason = format!("RPC timed out after {:?}", timeout_duration); - self.evict_connection(op, &reason).await; + warn!( + addr = %self.addr, + op, + reason, + "Remote lock RPC timed out without evicting cached connection" + ); Err(LockError::timeout(format!("remote lock RPC {op} on {}", self.addr), timeout_duration)) } } @@ -494,7 +499,7 @@ mod tests { } #[tokio::test] - async fn test_remote_client_acquire_lock_respects_request_timeout_and_evicts_connection() { + async fn test_remote_client_acquire_lock_respects_request_timeout_without_evicting_connection() { ensure_test_rpc_secret(); let (addr, accept_task) = spawn_hanging_listener().await; cache_lazy_channel(&addr).await; @@ -513,15 +518,15 @@ mod tests { assert!(!response.success, "timed out lock acquisition should fail"); assert_eq!(response.error.as_deref(), Some("Lock acquisition timeout")); assert!( - !GLOBAL_CONN_MAP.read().await.contains_key(&addr), - "timeout should evict cached connection" + GLOBAL_CONN_MAP.read().await.contains_key(&addr), + "local timeout should keep cached connection" ); accept_task.abort(); } #[tokio::test] - async fn test_remote_client_acquire_locks_batch_respects_request_timeout_and_evicts_connection() { + async fn test_remote_client_acquire_locks_batch_respects_request_timeout_without_evicting_connection() { ensure_test_rpc_secret(); let (addr, accept_task) = spawn_hanging_listener().await; cache_lazy_channel(&addr).await; @@ -541,13 +546,33 @@ mod tests { assert!(!responses[0].success, "timed out batch lock acquisition should fail"); assert_eq!(responses[0].error.as_deref(), Some("Lock acquisition timeout")); assert!( - !GLOBAL_CONN_MAP.read().await.contains_key(&addr), - "batch timeout should evict cached connection" + GLOBAL_CONN_MAP.read().await.contains_key(&addr), + "local batch timeout should keep cached connection" ); accept_task.abort(); } + #[tokio::test] + async fn test_remote_client_rpc_status_error_evicts_connection() { + let addr = "http://127.0.0.1:1".to_string(); + cache_lazy_channel(&addr).await; + assert!(GLOBAL_CONN_MAP.read().await.contains_key(&addr)); + + let client = RemoteClient::new(addr.clone()); + let result = client + .execute_rpc("lock", Duration::from_millis(50), async { + Err::<(), _>(tonic::Status::unavailable("connection unavailable")) + }) + .await; + + assert!(result.is_err(), "RPC status errors should be returned"); + assert!( + !GLOBAL_CONN_MAP.read().await.contains_key(&addr), + "RPC status errors should evict cached connection" + ); + } + #[test] fn test_remote_client_zero_timeout_is_clamped() { assert_eq!(RemoteClient::rpc_timeout(Duration::ZERO), Duration::from_millis(1)); diff --git a/crates/lock/src/distributed_lock.rs b/crates/lock/src/distributed_lock.rs index 431422c1d..30a930cc3 100644 --- a/crates/lock/src/distributed_lock.rs +++ b/crates/lock/src/distributed_lock.rs @@ -30,6 +30,8 @@ use uuid::Uuid; const UNLOCK_RETRY_ATTEMPTS: usize = 3; const UNLOCK_RETRY_BACKOFF: Duration = Duration::from_millis(100); +const LOCK_ACQUIRE_RETRY_ATTEMPTS: usize = 3; +const LOCK_ACQUIRE_RETRY_INITIAL_BACKOFF: Duration = Duration::from_millis(250); /// Generate a new aggregate lock ID for multiple client locks fn generate_aggregate_lock_id(resource: &ObjectKey) -> LockId { @@ -131,6 +133,12 @@ pub struct DistributedLock { type LockAcquireTaskResult = (usize, Result); +struct LockAcquireQuorumResult { + response: LockResponse, + individual_locks: Vec<(LockId, Arc)>, + pending_cleanup_spawned: bool, +} + impl DistributedLock { /// Create new distributed lock pub fn new(namespace: String, clients: Vec>, quorum: usize) -> Self { @@ -185,7 +193,11 @@ impl DistributedLock { } let required_quorum = self.required_quorum(request.lock_type); - let (resp, individual_locks) = self.acquire_lock_quorum(request).await?; + let LockAcquireQuorumResult { + response: resp, + individual_locks, + .. + } = self.acquire_lock_quorum_with_retry(request).await?; if resp.success { // Use aggregate lock_id from LockResponse's LockInfo // The aggregate id is what we expose to callers; individual_locks carries @@ -287,6 +299,20 @@ impl DistributedLock { pending } + fn lock_acquire_retry_backoff(attempt: usize) -> Duration { + LOCK_ACQUIRE_RETRY_INITIAL_BACKOFF * attempt as u32 + } + + fn is_retryable_lock_failure(resp: &LockResponse) -> bool { + resp.error + .as_deref() + .map(|error| { + let error = error.to_ascii_lowercase(); + error.contains("timeout") || error.contains("rpc failed") + }) + .unwrap_or(false) + } + async fn release_entries(entries: Vec<(LockId, Arc)>, context: &'static str) { let mut pending = entries; @@ -408,14 +434,57 @@ impl DistributedLock { } } + async fn acquire_lock_quorum_with_retry(&self, request: &LockRequest) -> Result { + let start = std::time::Instant::now(); + let mut attempt = 1; + let mut last_result = None; + + loop { + let elapsed = start.elapsed(); + if elapsed >= request.acquire_timeout { + break; + } + + let remaining = request.acquire_timeout - elapsed; + let mut attempt_request = request.clone(); + attempt_request.acquire_timeout = remaining; + + let result = self.acquire_lock_quorum_once(&attempt_request).await?; + if result.response.success + || !result.individual_locks.is_empty() + || result.pending_cleanup_spawned + || !Self::is_retryable_lock_failure(&result.response) + || attempt >= LOCK_ACQUIRE_RETRY_ATTEMPTS + { + return Ok(result); + } + + last_result = Some(result); + let backoff = Self::lock_acquire_retry_backoff(attempt); + if start.elapsed().saturating_add(backoff) >= request.acquire_timeout { + break; + } + + tokio::time::sleep(backoff).await; + attempt += 1; + } + + Ok(last_result.unwrap_or_else(|| LockAcquireQuorumResult { + response: LockResponse::failure("Lock acquisition timeout", request.acquire_timeout), + individual_locks: Vec::new(), + pending_cleanup_spawned: false, + })) + } + /// Quorum-based lock acquisition: success if at least the required quorum succeeds. /// Collects all individual lock_ids from successful clients and creates an aggregate lock_id. /// Returns the LockResponse with aggregate lock_id and individual lock mappings. - async fn acquire_lock_quorum(&self, request: &LockRequest) -> Result<(LockResponse, Vec<(LockId, Arc)>)> { + async fn acquire_lock_quorum_once(&self, request: &LockRequest) -> Result { let required_quorum = self.required_quorum(request.lock_type); let mut pending = self.spawn_lock_requests(request); let mut individual_locks: Vec<(LockId, Arc)> = Vec::new(); let fallback_lock_id = request.lock_id.clone(); + let mut last_failure = None; while let Some(join_result) = pending.join_next().await { match join_result { @@ -434,18 +503,22 @@ impl DistributedLock { } } else { let error = resp.error.unwrap_or_else(|| "unknown error".to_string()); - self.log_failed_lock_response(request, idx, error); + self.log_failed_lock_response(request, idx, error.clone()); + last_failure = Some(error); } } Ok((idx, Err(err))) => { tracing::warn!("Failed to acquire lock on client {}: {}", idx, err); + last_failure = Some(err.to_string()); } Err(err) => { tracing::warn!("Lock acquisition task join failed: {}", err); + last_failure = Some(err.to_string()); } } if individual_locks.len() >= required_quorum { + let pending_cleanup_spawned = !pending.is_empty(); if !pending.is_empty() { Self::spawn_pending_cleanup( pending, @@ -479,12 +552,17 @@ impl DistributedLock { }, Duration::ZERO, ); - return Ok((resp, individual_locks)); + return Ok(LockAcquireQuorumResult { + response: resp, + individual_locks, + pending_cleanup_spawned, + }); } - if individual_locks.len() + pending.len() < required_quorum { + if !individual_locks.is_empty() && individual_locks.len() + pending.len() < required_quorum { let rollback_count = individual_locks.len(); Self::spawn_release_cleanup(individual_locks.clone(), "distributed_lock_quorum_rollback"); + let pending_cleanup_spawned = !pending.is_empty(); if !pending.is_empty() { Self::spawn_pending_cleanup( pending, @@ -494,21 +572,33 @@ impl DistributedLock { ); } - let resp = LockResponse::failure( - format!("Failed to acquire quorum: {rollback_count}/{required_quorum} required"), - Duration::ZERO, - ); - return Ok((resp, individual_locks)); + let mut error = format!("Failed to acquire quorum: {rollback_count}/{required_quorum} required"); + if let Some(last_failure) = last_failure { + error.push_str("; last failure: "); + error.push_str(&last_failure); + } + let resp = LockResponse::failure(error, Duration::ZERO); + return Ok(LockAcquireQuorumResult { + response: resp, + individual_locks, + pending_cleanup_spawned, + }); } } let rollback_count = individual_locks.len(); Self::spawn_release_cleanup(individual_locks.clone(), "distributed_lock_quorum_rollback"); - let resp = LockResponse::failure( - format!("Failed to acquire quorum: {rollback_count}/{required_quorum} required"), - Duration::ZERO, - ); - Ok((resp, individual_locks)) + let mut error = format!("Failed to acquire quorum: {rollback_count}/{required_quorum} required"); + if let Some(last_failure) = last_failure { + error.push_str("; last failure: "); + error.push_str(&last_failure); + } + let resp = LockResponse::failure(error, Duration::ZERO); + Ok(LockAcquireQuorumResult { + response: resp, + individual_locks, + pending_cleanup_spawned: false, + }) } } diff --git a/crates/lock/src/namespace/tests.rs b/crates/lock/src/namespace/tests.rs index b95fbfed5..06d981e4f 100644 --- a/crates/lock/src/namespace/tests.rs +++ b/crates/lock/src/namespace/tests.rs @@ -110,6 +110,75 @@ impl crate::client::LockClient for DelayedClient { } } +#[derive(Debug)] +struct FlakyAcquireClient { + inner: LocalClient, + failed_acquires_remaining: AtomicUsize, + acquire_attempts: AtomicUsize, +} + +impl FlakyAcquireClient { + fn new(manager: Arc, failed_acquires: usize) -> Self { + Self { + inner: LocalClient::with_manager(manager), + failed_acquires_remaining: AtomicUsize::new(failed_acquires), + acquire_attempts: AtomicUsize::new(0), + } + } + + fn acquire_attempts(&self) -> usize { + self.acquire_attempts.load(Ordering::SeqCst) + } +} + +#[async_trait::async_trait] +impl crate::client::LockClient for FlakyAcquireClient { + async fn acquire_lock(&self, request: &LockRequest) -> crate::Result { + self.acquire_attempts.fetch_add(1, Ordering::SeqCst); + if self + .failed_acquires_remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| remaining.checked_sub(1)) + .is_ok() + { + return Ok(LockResponse::failure("Lock acquisition timeout", request.acquire_timeout)); + } + + self.inner.acquire_lock(request).await + } + + async fn release(&self, lock_id: &LockId) -> crate::Result { + self.inner.release(lock_id).await + } + + async fn refresh(&self, lock_id: &LockId) -> crate::Result { + self.inner.refresh(lock_id).await + } + + async fn force_release(&self, lock_id: &LockId) -> crate::Result { + self.inner.force_release(lock_id).await + } + + async fn check_status(&self, lock_id: &LockId) -> crate::Result> { + self.inner.check_status(lock_id).await + } + + async fn get_stats(&self) -> crate::Result { + self.inner.get_stats().await + } + + async fn close(&self) -> crate::Result<()> { + self.inner.close().await + } + + async fn is_online(&self) -> bool { + self.inner.is_online().await + } + + async fn is_local(&self) -> bool { + self.inner.is_local().await + } +} + #[derive(Debug)] struct FlakyReleaseClient { inner: LocalClient, @@ -672,6 +741,69 @@ async fn test_namespace_lock_distributed_unlock_retries_release_false() { ); } +#[tokio::test] +async fn test_namespace_lock_distributed_retries_transient_acquire_timeout() { + let managers = (0..3).map(|_| Arc::new(GlobalLockManager::new())).collect::>(); + let flaky_clients = managers + .iter() + .map(|manager| Arc::new(FlakyAcquireClient::new(manager.clone(), 1))) + .collect::>(); + let clients = flaky_clients + .iter() + .map(|client| client.clone() as Arc) + .collect::>(); + + let lock = NamespaceLock::with_clients("flaky-acquire".to_string(), clients); + let resource = create_test_object_key("bucket", "object-flaky-acquire"); + + let guard = lock + .get_write_lock(resource, "owner-a", Duration::from_secs(1)) + .await + .expect("transient timeout should be retried before the acquire budget expires"); + + assert!( + flaky_clients.iter().all(|client| client.acquire_attempts() >= 2), + "each simulated node should be retried after an initial timeout" + ); + + drop(guard); +} + +#[tokio::test] +async fn test_namespace_lock_distributed_waits_full_timeout_for_late_release() { + let managers = (0..3).map(|_| Arc::new(GlobalLockManager::new())).collect::>(); + let clients = managers + .iter() + .map(|manager| Arc::new(LocalClient::with_manager(manager.clone())) as Arc) + .collect::>(); + + let lock = Arc::new(NamespaceLock::with_clients("late-release".to_string(), clients)); + let resource = create_test_object_key("bucket", "object-late-release"); + + let guard_a = lock + .get_write_lock(resource.clone(), "owner-a", Duration::from_secs(1)) + .await + .expect("owner-a should acquire the initial distributed lock"); + + let lock_for_owner_b = lock.clone(); + let resource_for_owner_b = resource.clone(); + let waiter = tokio::spawn(async move { + lock_for_owner_b + .get_write_lock(resource_for_owner_b, "owner-b", Duration::from_secs(1)) + .await + }); + + tokio::time::sleep(Duration::from_millis(950)).await; + drop(guard_a); + + let guard_b = waiter + .await + .expect("owner-b wait task should complete") + .expect("owner-b should acquire after a late release within the original timeout"); + + drop(guard_b); +} + #[test] fn test_namespace_lock_distributed_drop_without_runtime_does_not_panic() { let (manager, resource, guard) = {