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

@@ -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"
+            );
+        }
+    }
 }

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
+        // Collect the keys to evict during the walk, then remove them afterward.
-        let mut to_recycle = Vec::new();
+        // `retain` only exposes a shared `&Arc`, and cloning it keeps
-        objects.retain(|_key, state| {
+        // `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
+                    to_remove.push(key.clone());
-                    if let Ok(state_box) = Arc::try_unwrap(state.clone()) {
-                        to_recycle.push(state_box);
-                    }
                     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
+        // Evict the collected keys and recycle their states into the pool.
-        for state_box in to_recycle {
+        // Removing the key yields the owned `Arc`; `try_unwrap` succeeds only
-            let boxed_state = Box::new(state_box);
+        // when this shard held the last reference, in which case the inner
-            self.object_pool.release(boxed_state);
+        // 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<str> = 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");
+    }
 }