From ed45d1be2d8bb7b20ddd141637501295a9d26d04 Mon Sep 17 00:00:00 2001 From: houseme Date: Mon, 22 Jun 2026 13:53:21 +0800 Subject: [PATCH] fix(lock): recover stale distributed object guards (#3720) fix(lock): reclaim expired stale object guards --- crates/lock/src/client/local.rs | 184 +++++++++++++++++++++-------- crates/lock/src/namespace/tests.rs | 45 +++++++ 2 files changed, 179 insertions(+), 50 deletions(-) diff --git a/crates/lock/src/client/local.rs b/crates/lock/src/client/local.rs index 04aa1e011..ea6299ca5 100644 --- a/crates/lock/src/client/local.rs +++ b/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>>>, + guard_storage: Vec>>>, /// Mask for fast shard index calculation (shard_count - 1) shard_mask: usize, /// Optional lock manager (if None, uses global singleton) manager: Option>, } +#[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>>> = + let guard_storage: Vec>>> = (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>> { + fn get_shard(&self, lock_id: &LockId) -> &Arc>> { 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 { 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 { - // For local locks, refresh is not needed as they don't expire automatically - Ok(true) + async fn refresh(&self, lock_id: &LockId) -> Result { + 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 { @@ -172,21 +251,26 @@ impl LockClient for LocalClient { async fn check_status(&self, lock_id: &LockId) -> Result> { 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, diff --git a/crates/lock/src/namespace/tests.rs b/crates/lock/src/namespace/tests.rs index a87fb863b..a492380df 100644 --- a/crates/lock/src/namespace/tests.rs +++ b/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::>(); + let flaky_clients = managers + .iter() + .map(|manager| Arc::new(FlakyReleaseClient::new(manager.clone(), 8))) + .collect::>(); + let clients = flaky_clients + .iter() + .map(|client| client.clone() as Arc) + .collect::>(); + + 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::>();