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 <marshawcoco@users.noreply.github.com>
Co-authored-by: houseme <housemecn@gmail.com>

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));

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<LockResponse>);
 
+struct LockAcquireQuorumResult {
+    response: LockResponse,
+    individual_locks: Vec<(LockId, Arc<dyn LockClient>)>,
+    pending_cleanup_spawned: bool,
+}
+
 impl DistributedLock {
     /// Create new distributed lock
     pub fn new(namespace: String, clients: Vec<Arc<dyn LockClient>>, 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<dyn LockClient>)>, context: &'static str) {
         let mut pending = entries;
 
@@ -408,14 +434,57 @@ impl DistributedLock {
         }
     }
 
+    async fn acquire_lock_quorum_with_retry(&self, request: &LockRequest) -> Result<LockAcquireQuorumResult> {
+        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<dyn LockClient>)>)> {
+    async fn acquire_lock_quorum_once(&self, request: &LockRequest) -> Result<LockAcquireQuorumResult> {
         let required_quorum = self.required_quorum(request.lock_type);
         let mut pending = self.spawn_lock_requests(request);
         let mut individual_locks: Vec<(LockId, Arc<dyn LockClient>)> = 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,
+        })
     }
 }
 

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<GlobalLockManager>, 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<LockResponse> {
+        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<bool> {
+        self.inner.release(lock_id).await
+    }
+
+    async fn refresh(&self, lock_id: &LockId) -> crate::Result<bool> {
+        self.inner.refresh(lock_id).await
+    }
+
+    async fn force_release(&self, lock_id: &LockId) -> crate::Result<bool> {
+        self.inner.force_release(lock_id).await
+    }
+
+    async fn check_status(&self, lock_id: &LockId) -> crate::Result<Option<LockInfo>> {
+        self.inner.check_status(lock_id).await
+    }
+
+    async fn get_stats(&self) -> crate::Result<LockStats> {
+        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::<Vec<_>>();
+    let flaky_clients = managers
+        .iter()
+        .map(|manager| Arc::new(FlakyAcquireClient::new(manager.clone(), 1)))
+        .collect::<Vec<_>>();
+    let clients = flaky_clients
+        .iter()
+        .map(|client| client.clone() as Arc<dyn LockClient>)
+        .collect::<Vec<_>>();
+
+    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::<Vec<_>>();
+    let clients = managers
+        .iter()
+        .map(|manager| Arc::new(LocalClient::with_manager(manager.clone())) as Arc<dyn LockClient>)
+        .collect::<Vec<_>>();
+
+    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) = {