From 7b87d4d13b90e4c7fbc32db3124b2561ce44ecaf Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Thu, 9 Jul 2026 01:07:41 +0800 Subject: [PATCH] fix(lock): stop waiter-count leak and dead pool recycle (#4520) Make OptimizedNotify::wait_for_read/wait_for_write cancellation-safe with an RAII decrement guard so the waiter counter is decremented even when the future is dropped at the await (capped-wait timeout or task abort), preventing the counter from leaking upward on a hot lock. Fix cleanup_expired_batch so recycling into the object pool actually works: Arc::try_unwrap on a clone could never succeed, so no state was ever recycled. Collect keys during retain, remove them afterward to obtain the owned Arc, and try_unwrap that before returning the state to the pool. Refs rustfs/backlog#1035 --- crates/lock/src/fast_lock/optimized_notify.rs | 80 +++++++++++- crates/lock/src/fast_lock/shard.rs | 118 +++++++++++++++--- 2 files changed, 178 insertions(+), 20 deletions(-) diff --git a/crates/lock/src/fast_lock/optimized_notify.rs b/crates/lock/src/fast_lock/optimized_notify.rs index 8aa183af1..3674933a2 100644 --- a/crates/lock/src/fast_lock/optimized_notify.rs +++ b/crates/lock/src/fast_lock/optimized_notify.rs @@ -32,6 +32,29 @@ pub struct OptimizedNotify { pub notify_pool_index: AtomicUsize, } +/// Cancellation-safe waiter-count ticket for [`OptimizedNotify`]. +/// +/// Increments the referenced counter on construction and decrements it on +/// drop. The decrement therefore runs even when the awaiting future is +/// cancelled/dropped before the `notified()` await completes, so the counter +/// cannot leak on capped-wait timeouts. +struct WaiterCountGuard<'a> { + counter: &'a AtomicU32, +} + +impl<'a> WaiterCountGuard<'a> { + fn new(counter: &'a AtomicU32) -> Self { + counter.fetch_add(1, Ordering::AcqRel); + Self { counter } + } +} + +impl Drop for WaiterCountGuard<'_> { + fn drop(&mut self) { + self.counter.fetch_sub(1, Ordering::AcqRel); + } +} + impl OptimizedNotify { pub fn new() -> Self { // Use random pool index to distribute load @@ -67,18 +90,20 @@ impl OptimizedNotify { /// Wait for reader notification pub async fn wait_for_read(&self) { - self.reader_waiters.fetch_add(1, Ordering::AcqRel); + // RAII guard decrements the counter even if this future is dropped at + // the await below (e.g. when the caller wraps it in a `timeout(...)` + // that elapses), preventing the waiter counter from leaking upward. + let _guard = WaiterCountGuard::new(&self.reader_waiters); let pool_index = self.notify_pool_index.load(Ordering::Relaxed) % NOTIFY_POOL.len(); NOTIFY_POOL[pool_index].notified().await; - self.reader_waiters.fetch_sub(1, Ordering::AcqRel); } /// Wait for writer notification pub async fn wait_for_write(&self) { - self.writer_waiters.fetch_add(1, Ordering::AcqRel); + // See `wait_for_read` for why the decrement must be RAII-guarded. + let _guard = WaiterCountGuard::new(&self.writer_waiters); let pool_index = self.notify_pool_index.load(Ordering::Relaxed) % NOTIFY_POOL.len(); NOTIFY_POOL[pool_index].notified().await; - self.writer_waiters.fetch_sub(1, Ordering::AcqRel); } /// Check if anyone is waiting @@ -132,4 +157,51 @@ mod tests { assert!(timeout(Duration::from_millis(100), handle).await.is_ok()); } + + #[tokio::test] + async fn test_reader_waiter_count_released_on_abort() { + let notify = Arc::new(OptimizedNotify::new()); + let notify_for_task = notify.clone(); + + let handle = tokio::spawn(async move { + notify_for_task.wait_for_read().await; + }); + + // Wait until the task has registered itself as a waiting reader. + timeout(Duration::from_secs(1), async { + while notify.reader_waiters.load(Ordering::Acquire) == 0 { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("reader waiter never registered"); + assert_eq!(notify.reader_waiters.load(Ordering::Acquire), 1); + + handle.abort(); + let _ = handle.await; + + assert_eq!( + notify.reader_waiters.load(Ordering::Acquire), + 0, + "reader waiter count must return to 0 after the waiting future is dropped" + ); + } + + #[tokio::test] + async fn test_writer_waiter_count_released_after_capped_timeouts() { + // Reproduces the shard slow-path pattern: the notification is never + // delivered, so each capped `timeout` elapses and drops the waiting + // future. Without an RAII decrement the counter would climb by one per + // iteration; with it, the counter must return to 0 every time. + let notify = OptimizedNotify::new(); + + for _ in 0..5 { + let _ = timeout(Duration::from_millis(10), notify.wait_for_write()).await; + assert_eq!( + notify.writer_waiters.load(Ordering::Acquire), + 0, + "writer waiter count must return to 0 after a capped-wait timeout" + ); + } + } } diff --git a/crates/lock/src/fast_lock/shard.rs b/crates/lock/src/fast_lock/shard.rs index 9de3d10a6..769075063 100644 --- a/crates/lock/src/fast_lock/shard.rs +++ b/crates/lock/src/fast_lock/shard.rs @@ -722,9 +722,13 @@ impl LockShard { let mut objects = self.objects.write(); let mut processed = 0; - // Process in batches to avoid long-held locks - let mut to_recycle = Vec::new(); - objects.retain(|_key, state| { + // Collect the keys to evict during the walk, then remove them afterward. + // `retain` only exposes a shared `&Arc`, and cloning it keeps + // `strong_count >= 2`, so `Arc::try_unwrap` on that clone can never + // succeed. Removing the key after the walk hands back the owned `Arc`, + // which can be unwrapped and recycled into the pool. + let mut to_remove = Vec::new(); + objects.retain(|key, state| { if processed >= max_batch_size { return true; // Stop processing after batch limit } @@ -735,24 +739,25 @@ impl LockShard { let idle_time = now_millis.saturating_sub(last_access_millis); if idle_time > cleanup_threshold_millis { - // Try to recycle the state back to pool if possible - if let Ok(state_box) = Arc::try_unwrap(state.clone()) { - to_recycle.push(state_box); - } + to_remove.push(key.clone()); cleaned += 1; - false // Remove - } else { - true // Keep } - } else { - true // Keep active locks } + // Removal is deferred to the loop below so the owned `Arc` can be + // recovered for recycling. + true }); - // Return recycled objects to pool - for state_box in to_recycle { - let boxed_state = Box::new(state_box); - self.object_pool.release(boxed_state); + // Evict the collected keys and recycle their states into the pool. + // Removing the key yields the owned `Arc`; `try_unwrap` succeeds only + // when this shard held the last reference, in which case the inner + // state is reboxed and returned to the pool. + for key in to_remove { + if let Some(state) = objects.remove(&key) + && let Ok(state) = Arc::try_unwrap(state) + { + self.object_pool.release(Box::new(state)); + } } self.metrics.record_cleanup(cleaned); @@ -1002,6 +1007,19 @@ mod tests { waiter_handle.abort(); let _ = waiter_handle.await; + // The OptimizedNotify writer waiter counter must not leak after the + // waiting future is dropped by the abort. + if let Some(state) = shard.objects.read().get(&key).cloned() { + assert_eq!( + state + .optimized_notify + .writer_waiters + .load(std::sync::atomic::Ordering::Acquire), + 0, + "optimized_notify writer waiter count must return to 0 after abort" + ); + } + assert!(shard.release_lock(&key, &owner1, LockMode::Exclusive)); let followup_reader = ObjectLockRequest { @@ -1067,6 +1085,19 @@ mod tests { waiter_handle.abort(); let _ = waiter_handle.await; + // The OptimizedNotify reader waiter counter must not leak after the + // waiting future is dropped by the abort. + if let Some(state) = shard.objects.read().get(&key).cloned() { + assert_eq!( + state + .optimized_notify + .reader_waiters + .load(std::sync::atomic::Ordering::Acquire), + 0, + "optimized_notify reader waiter count must return to 0 after abort" + ); + } + assert!(shard.release_lock(&key, &writer_owner, LockMode::Exclusive)); let followup_writer = ObjectLockRequest { @@ -1084,4 +1115,59 @@ mod tests { ); assert!(shard.release_lock(&key, &followup_owner, LockMode::Exclusive)); } + + #[tokio::test] + async fn test_cleanup_expired_batch_recycles_into_pool() { + let shard = LockShard::new(0); + let owner: Arc = Arc::from("owner"); + const N: usize = 8; + + // Populate the shard with idle (acquired then released) lock states. + for i in 0..N { + let key = ObjectKey::new("bucket", format!("obj-{i}")); + let request = ObjectLockRequest { + key: key.clone(), + mode: LockMode::Exclusive, + owner: owner.clone(), + acquire_timeout: Duration::from_secs(1), + lock_timeout: Duration::from_secs(30), + priority: LockPriority::Normal, + }; + assert!(shard.acquire_lock(&request).await.is_ok()); + assert!(shard.release_lock(&key, &owner, LockMode::Exclusive)); + } + + // The pool started empty, so every acquisition allocated a fresh state: + // N misses and no releases yet. + let (_hits, misses, releases_before, _pool_size) = shard.pool_stats(); + assert_eq!(misses, N as u64); + assert_eq!(releases_before, 0); + + // Make sure the idle states are old enough to be evicted, then force a + // batch cleanup with a zero idle threshold. + tokio::time::sleep(Duration::from_millis(5)).await; + let cleaned = shard.cleanup_expired_batch(N, 0); + assert_eq!(cleaned, N, "all idle states should be cleaned"); + assert_eq!(shard.lock_count(), 0, "cleaned entries must be removed from the shard"); + + // The evicted states must actually be recycled back into the pool. + let (_hits, _misses, releases_after, pool_size) = shard.pool_stats(); + assert_eq!(releases_after, N as u64, "cleanup must recycle evicted states into the pool"); + assert_eq!(pool_size, N, "recycled states must be available in the pool"); + + // A subsequent acquire should reuse a pooled state, raising the hit rate. + let key = ObjectKey::new("bucket", "reuse"); + let request = ObjectLockRequest { + key, + mode: LockMode::Exclusive, + owner: owner.clone(), + acquire_timeout: Duration::from_secs(1), + lock_timeout: Duration::from_secs(30), + priority: LockPriority::Normal, + }; + assert!(shard.acquire_lock(&request).await.is_ok()); + let (hits, _misses, _releases, _pool_size) = shard.pool_stats(); + assert!(hits >= 1, "subsequent acquire should hit the recycled pool"); + assert!(shard.pool_hit_rate() > 0.0, "pool hit rate should rise after recycling"); + } }