fix(lock): recover stale distributed object guards (#3720)

fix(lock): reclaim expired stale object guards

crates/lock/src/client/local.rs

@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
+use std::time::{Duration, SystemTime};
 use tokio::sync::RwLock;
 
 use crate::{
@@ -29,13 +30,39 @@ const DEFAULT_GUARD_SHARD_COUNT: usize = 64;
 #[derive(Debug)]
 pub struct LocalClient {
     /// Sharded guard storage to reduce lock contention
-    guard_storage: Vec<Arc<RwLock<HashMap<LockId, FastLockGuard>>>>,
+    guard_storage: Vec<Arc<RwLock<HashMap<LockId, LocalGuardEntry>>>>,
     /// Mask for fast shard index calculation (shard_count - 1)
     shard_mask: usize,
     /// Optional lock manager (if None, uses global singleton)
     manager: Option<Arc<GlobalLockManager>>,
 }
 
+#[derive(Debug)]
+struct LocalGuardEntry {
+    guard: FastLockGuard,
+    expires_at: SystemTime,
+    ttl: Duration,
+}
+
+impl LocalGuardEntry {
+    fn new(guard: FastLockGuard, ttl: Duration) -> Self {
+        let now = SystemTime::now();
+        Self {
+            guard,
+            expires_at: now + ttl,
+            ttl,
+        }
+    }
+
+    fn is_expired(&self) -> bool {
+        self.expires_at <= SystemTime::now()
+    }
+
+    fn refresh(&mut self) {
+        self.expires_at = SystemTime::now() + self.ttl;
+    }
+}
+
 impl LocalClient {
     /// Create new local client with default shard count
     pub fn new() -> Self {
@@ -47,7 +74,7 @@ impl LocalClient {
     pub fn with_shard_count(shard_count: usize) -> Self {
         assert!(shard_count.is_power_of_two(), "Shard count must be power of 2");
 
-        let guard_storage: Vec<Arc<RwLock<HashMap<LockId, FastLockGuard>>>> =
+        let guard_storage: Vec<Arc<RwLock<HashMap<LockId, LocalGuardEntry>>>> =
             (0..shard_count).map(|_| Arc::new(RwLock::new(HashMap::new()))).collect();
 
         Self {
@@ -82,10 +109,40 @@ impl LocalClient {
     }
 
     /// Get the shard for a given lock ID
-    fn get_shard(&self, lock_id: &LockId) -> &Arc<RwLock<HashMap<LockId, FastLockGuard>>> {
+    fn get_shard(&self, lock_id: &LockId) -> &Arc<RwLock<HashMap<LockId, LocalGuardEntry>>> {
         let index = self.get_shard_index(lock_id);
         &self.guard_storage[index]
     }
+
+    async fn reclaim_expired_guards_for_resource(&self, resource: &crate::ObjectKey) -> usize {
+        let mut reclaimed = 0usize;
+
+        for shard in &self.guard_storage {
+            let expired_entries = {
+                let mut guards = shard.write().await;
+                let mut retained = HashMap::with_capacity(guards.len());
+                let mut expired_entries = Vec::new();
+
+                for (lock_id, entry) in std::mem::take(&mut *guards) {
+                    if &lock_id.resource == resource && entry.is_expired() {
+                        expired_entries.push(entry);
+                    } else {
+                        retained.insert(lock_id, entry);
+                    }
+                }
+
+                *guards = retained;
+                expired_entries
+            };
+
+            for mut entry in expired_entries {
+                let _ = entry.guard.release();
+                reclaimed = reclaimed.saturating_add(1);
+            }
+        }
+
+        reclaimed
+    }
 }
 
 impl Default for LocalClient {
@@ -98,51 +155,67 @@ impl Default for LocalClient {
 impl LockClient for LocalClient {
     async fn acquire_lock(&self, request: &LockRequest) -> Result<LockResponse> {
         let lock_manager = self.get_lock_manager();
+        let reclaimed_before_acquire = self.reclaim_expired_guards_for_resource(&request.resource).await;
 
-        let lock_request = match request.lock_type {
+        let build_lock_request = || match request.lock_type {
             LockType::Exclusive => crate::ObjectLockRequest::new_write(request.resource.clone(), request.owner.clone())
                 .with_acquire_timeout(request.acquire_timeout),
             LockType::Shared => crate::ObjectLockRequest::new_read(request.resource.clone(), request.owner.clone())
                 .with_acquire_timeout(request.acquire_timeout),
         };
 
-        match lock_manager.acquire_lock(lock_request).await {
-            Ok(guard) => {
-                let lock_id = request.lock_id.clone();
+        let mut retried_after_reclaim = reclaimed_before_acquire > 0;
+        loop {
+            match lock_manager.acquire_lock(build_lock_request()).await {
+                Ok(guard) => {
+                    let lock_id = request.lock_id.clone();
+                    let acquired_at = SystemTime::now();
+                    let expires_at = acquired_at + request.ttl;
 
-                {
-                    let shard = self.get_shard(&lock_id);
-                    let mut guards = shard.write().await;
-                    guards.insert(lock_id.clone(), guard);
+                    {
+                        let shard = self.get_shard(&lock_id);
+                        let mut guards = shard.write().await;
+                        guards.insert(lock_id.clone(), LocalGuardEntry::new(guard, request.ttl));
+                    }
+
+                    let lock_info = LockInfo {
+                        id: lock_id,
+                        resource: request.resource.clone(),
+                        lock_type: request.lock_type,
+                        status: crate::types::LockStatus::Acquired,
+                        owner: request.owner.clone(),
+                        acquired_at,
+                        expires_at,
+                        last_refreshed: acquired_at,
+                        metadata: request.metadata.clone(),
+                        priority: request.priority,
+                        wait_start_time: None,
+                    };
+                    return Ok(LockResponse::success(lock_info, Duration::ZERO));
+                }
+                Err(crate::fast_lock::LockResult::Timeout) => {
+                    if !retried_after_reclaim && self.reclaim_expired_guards_for_resource(&request.resource).await > 0 {
+                        retried_after_reclaim = true;
+                        continue;
+                    }
+                    return Ok(LockResponse::failure("Lock acquisition timeout", request.acquire_timeout));
+                }
+                Err(crate::fast_lock::LockResult::Conflict {
+                    current_owner,
+                    current_mode,
+                }) => {
+                    if !retried_after_reclaim && self.reclaim_expired_guards_for_resource(&request.resource).await > 0 {
+                        retried_after_reclaim = true;
+                        continue;
+                    }
+                    return Ok(LockResponse::failure(
+                        format!("Lock conflict: resource held by {current_owner} in {current_mode:?} mode"),
+                        Duration::ZERO,
+                    ));
+                }
+                Err(crate::fast_lock::LockResult::Acquired) => {
+                    unreachable!("Acquired should not be an error")
                 }
-
-                let lock_info = LockInfo {
-                    id: lock_id,
-                    resource: request.resource.clone(),
-                    lock_type: request.lock_type,
-                    status: crate::types::LockStatus::Acquired,
-                    owner: request.owner.clone(),
-                    acquired_at: std::time::SystemTime::now(),
-                    expires_at: std::time::SystemTime::now() + request.ttl,
-                    last_refreshed: std::time::SystemTime::now(),
-                    metadata: request.metadata.clone(),
-                    priority: request.priority,
-                    wait_start_time: None,
-                };
-                Ok(LockResponse::success(lock_info, std::time::Duration::ZERO))
-            }
-            Err(crate::fast_lock::LockResult::Timeout) => {
-                Ok(LockResponse::failure("Lock acquisition timeout", request.acquire_timeout))
-            }
-            Err(crate::fast_lock::LockResult::Conflict {
-                current_owner,
-                current_mode,
-            }) => Ok(LockResponse::failure(
-                format!("Lock conflict: resource held by {current_owner} in {current_mode:?} mode"),
-                std::time::Duration::ZERO,
-            )),
-            Err(crate::fast_lock::LockResult::Acquired) => {
-                unreachable!("Acquired should not be an error")
             }
         }
     }
@@ -152,7 +225,7 @@ impl LockClient for LocalClient {
         let mut guards = shard.write().await;
         if let Some(guard) = guards.remove(lock_id) {
             // Guard automatically releases the lock when dropped
-            drop(guard);
+            drop(guard.guard);
             Ok(true)
         } else {
             // Lock not found or already released
@@ -160,9 +233,15 @@ impl LockClient for LocalClient {
         }
     }
 
-    async fn refresh(&self, _lock_id: &LockId) -> Result<bool> {
-        // For local locks, refresh is not needed as they don't expire automatically
-        Ok(true)
+    async fn refresh(&self, lock_id: &LockId) -> Result<bool> {
+        let shard = self.get_shard(lock_id);
+        let mut guards = shard.write().await;
+        if let Some(entry) = guards.get_mut(lock_id) {
+            entry.refresh();
+            Ok(true)
+        } else {
+            Ok(false)
+        }
     }
 
     async fn force_release(&self, lock_id: &LockId) -> Result<bool> {
@@ -172,21 +251,26 @@ impl LockClient for LocalClient {
     async fn check_status(&self, lock_id: &LockId) -> Result<Option<LockInfo>> {
         let shard = self.get_shard(lock_id);
         let guards = shard.read().await;
-        if let Some(guard) = guards.get(lock_id) {
+        if let Some(entry) = guards.get(lock_id) {
             // We have an active guard for this lock
-            let lock_type = match guard.mode() {
+            let lock_type = match entry.guard.mode() {
                 crate::LockMode::Shared => LockType::Shared,
                 crate::LockMode::Exclusive => LockType::Exclusive,
             };
+            let status = if entry.is_expired() {
+                LockStatus::Expired
+            } else {
+                LockStatus::Acquired
+            };
             Ok(Some(LockInfo {
                 id: lock_id.clone(),
                 resource: lock_id.resource.clone(),
                 lock_type,
-                status: LockStatus::Acquired,
-                owner: guard.owner().to_string(),
-                acquired_at: std::time::SystemTime::now(),
-                expires_at: std::time::SystemTime::now() + std::time::Duration::from_secs(30),
-                last_refreshed: std::time::SystemTime::now(),
+                status,
+                owner: entry.guard.owner().to_string(),
+                acquired_at: SystemTime::now(),
+                expires_at: entry.expires_at,
+                last_refreshed: SystemTime::now(),
                 metadata: LockMetadata::default(),
                 priority: LockPriority::Normal,
                 wait_start_time: None,

crates/lock/src/namespace/tests.rs

@@ -785,6 +785,51 @@ async fn test_namespace_lock_distributed_unlock_retries_release_false() {
     );
 }
 
+#[tokio::test]
+async fn test_namespace_lock_distributed_reclaims_expired_same_resource_after_failed_unlocks() {
+    let managers = (0..3).map(|_| Arc::new(GlobalLockManager::new())).collect::<Vec<_>>();
+    let flaky_clients = managers
+        .iter()
+        .map(|manager| Arc::new(FlakyReleaseClient::new(manager.clone(), 8)))
+        .collect::<Vec<_>>();
+    let clients = flaky_clients
+        .iter()
+        .map(|client| client.clone() as Arc<dyn LockClient>)
+        .collect::<Vec<_>>();
+
+    let lock = NamespaceLock::with_clients("expired-unlock-recovery".to_string(), clients);
+    let resource = create_test_object_key("bucket", "object-expired-unlock-recovery");
+    let request = LockRequest::new(resource.clone(), LockType::Exclusive, "owner-a")
+        .with_acquire_timeout(Duration::from_millis(100))
+        .with_ttl(Duration::from_millis(50));
+
+    let mut guard = lock
+        .acquire_guard(&request)
+        .await
+        .expect("initial owner should acquire the distributed lock")
+        .expect("distributed acquire should return a guard");
+
+    assert!(guard.release(), "distributed guard should enqueue unlock retries");
+    tokio::time::sleep(Duration::from_millis(80)).await;
+
+    let recovered_request = LockRequest::new(resource.clone(), LockType::Exclusive, "owner-b")
+        .with_acquire_timeout(Duration::from_millis(200))
+        .with_ttl(Duration::from_millis(50));
+
+    let recovered_guard = lock
+        .acquire_guard(&recovered_request)
+        .await
+        .expect("expired stale distributed guards should be reclaimed")
+        .expect("recovered acquire should return a guard");
+
+    drop(recovered_guard);
+
+    assert!(
+        flaky_clients.iter().all(|client| client.release_attempts() >= 1),
+        "unlock retries should have been attempted before TTL-based reclamation"
+    );
+}
+
 #[tokio::test]
 async fn test_namespace_lock_distributed_retries_transient_acquire_timeout() {
     let managers = (0..3).map(|_| Arc::new(GlobalLockManager::new())).collect::<Vec<_>>();