perf(semaphore): avoid the queue lock for positive-balance releases

CHANGELOG.md

@@ -32,6 +32,10 @@ All notable changes to this project will be documented in this file.
 * Prevent deadlocks in `singleflight::Group::work` and `try_work` when a duplicate key's destructor calls back into the same group.
 * Fix undefined behavior when sending or dropping a oneshot sender races with receiver completion or cancellation.
 
+### Improvements
+
+* Add a lock-free fast path for releasing `Semaphore` permits while permits remain available, which speeds up dropping `RwLock` read guards, releasing multi-permit `Semaphore` permits, and returning objects to a bounded pool that is not at capacity. `Mutex` releases are unchanged.
+
 ### Notable changes
 
 * Executor waker operations, including cloning, waking, and dropping, are expected not to panic; recovery from panicking waker callbacks is no longer supported.

asyncband/src/internal/semaphore.rs

@@ -27,6 +27,7 @@
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::MutexGuard;
+use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
@@ -44,6 +45,10 @@ use crate::internal::waker_batch::WakerBatch;
 pub struct Semaphore {
     /// The current number of available permits in the semaphore.
     permits: AtomicUsize,
+    /// Set while a locked release is losing races on the balance, so that other releases join
+    /// the lock instead of adding to the traffic. A hint only: a stale read in either direction
+    /// is still correct.
+    contended: AtomicBool,
     waiters: Mutex<WaitList<WaitNode>>,
 }
 
@@ -59,6 +64,7 @@ impl Semaphore {
     pub const fn new(permits: usize) -> Self {
         Self {
             permits: AtomicUsize::new(permits),
+            contended: AtomicBool::new(false),
             waiters: Mutex::new(WaitList::new()),
         }
     }
@@ -142,6 +148,33 @@ impl Semaphore {
 
     /// Adds `n` permits to the semaphore.
     pub fn release(&self, n: usize) {
+        if n == 0 {
+            return;
+        }
+
+        // A waiter is linked only after the balance has been drained to zero, and permits reach
+        // the balance again only once the queue is empty, so a positive balance means no waiter
+        // is linked and the permits can be added without the queue lock.
+        //
+        // Try the exchange once before falling back to the locked path. Other threads can
+        // still use the fast paths, so a locked balance update may also need to retry.
+        let current = self.permits.load(Ordering::Relaxed);
+        if current != 0
+            && !self.contended.load(Ordering::Relaxed)
+            && self.try_add_to_balance(current, n).is_ok()
+        {
+            return;
+        }
+
+        self.insert_permits_with_lock(n, self.waiters.lock());
+    }
+
+    /// Adds `n` permits to a semaphore whose every permit the caller holds.
+    ///
+    /// The balance is zero while one caller holds every permit, so the positive-balance path of
+    /// [`release`](Self::release) cannot succeed and mutex guards skip its probe. The locked path
+    /// is correct for any balance, so the precondition only affects speed.
+    pub fn release_all_held(&self, n: usize) {
         if n != 0 {
             self.insert_permits_with_lock(n, self.waiters.lock());
         }
@@ -175,6 +208,24 @@ impl Semaphore {
         wakers.wake_all();
     }
 
+    /// Adds `n` permits to a balance expected to hold `current`, or returns the balance observed
+    /// instead.
+    ///
+    /// The overflow check and the addition are one exchange because the positive-balance path of
+    /// [`release`](Self::release) can grow the balance even while the queue lock is held. The
+    /// exchange is the strong form because `release` tries it only once.
+    ///
+    /// ORDERING: Release publishes the work protected by the released permits to the Acquire
+    /// load or exchange that next observes this balance.
+    fn try_add_to_balance(&self, current: usize, n: usize) -> Result<(), usize> {
+        let next = current.checked_add(n).unwrap_or_else(|| {
+            panic!("number of added permits ({n}) would overflow usize::MAX (prev: {current})")
+        });
+        self.permits
+            .compare_exchange(current, next, Ordering::Release, Ordering::Relaxed)
+            .map(|_| ())
+    }
+
     fn insert_permits_with_lock<'a>(
         &'a self,
         mut rem: usize,
@@ -200,13 +251,21 @@ impl Semaphore {
             }
 
             if rem > 0 && waiters.is_empty() {
-                // The lock serializes additions; concurrent operations can only remove permits.
-                let current = self.permits.load(Ordering::Relaxed);
-                assert!(
-                    current.checked_add(rem).is_some(),
-                    "number of added permits ({rem}) would overflow usize::MAX (prev: {current})"
-                );
-                self.permits.fetch_add(rem, Ordering::Release);
+                // Fast releases can grow the balance under this lock, so the overflow check and
+                // addition must be one exchange. A zero balance cannot change under the lock, but
+                // still needs an RMW to extend the previous release sequence.
+                let mut current = self.permits.load(Ordering::Relaxed);
+                if current == 0 {
+                    self.permits.fetch_add(rem, Ordering::Release);
+                } else {
+                    while let Err(actual) = self.try_add_to_balance(current, rem) {
+                        // Losing under the lock means another thread is on the balance now:
+                        // route its release through the lock as well.
+                        self.contended.store(true, Ordering::Relaxed);
+                        current = actual;
+                    }
+                    self.contended.store(false, Ordering::Relaxed);
+                }
                 rem = 0;
             }
 
@@ -384,6 +443,8 @@ fn acquired_or_enqueue(
 
 #[cfg(test)]
 mod tests {
+    use std::panic;
+    use std::panic::AssertUnwindSafe;
     use std::sync::Arc;
     use std::sync::atomic::AtomicUsize;
     use std::sync::atomic::Ordering;
@@ -412,6 +473,89 @@ mod tests {
         }
     }
 
+    #[test]
+    fn release_with_positive_balance_does_not_take_the_queue_lock() {
+        let semaphore = Semaphore::new(1);
+
+        // Holding the queue lock deadlocks a release that needs it.
+        let queue = semaphore.waiters.lock();
+        semaphore.release(2);
+        assert_eq!(semaphore.available_permits(), 3);
+        assert!(semaphore.try_acquire(2));
+        semaphore.release(2);
+        assert_eq!(semaphore.available_permits(), 3);
+        drop(queue);
+    }
+
+    #[test]
+    fn contended_release_takes_the_locked_path_and_clears_the_flag() {
+        let semaphore = Semaphore::new(1);
+        semaphore.contended.store(true, Ordering::Relaxed);
+
+        // Only the locked path clears the flag, so a cleared flag proves the fast path was skipped.
+        semaphore.release(1);
+        assert_eq!(semaphore.available_permits(), 2);
+        assert!(!semaphore.contended.load(Ordering::Relaxed));
+    }
+
+    #[test]
+    fn release_with_zero_balance_hands_permits_to_the_queue() {
+        let semaphore = Semaphore::new(0);
+        let counter = Arc::new(WakeCounter(AtomicUsize::new(0)));
+        let waker = Waker::from(counter.clone());
+        let mut acquire = semaphore.poll_acquire(2);
+        assert!(acquire.poll_once(&waker).is_pending());
+
+        semaphore.release(3);
+        assert_eq!(counter.0.load(Ordering::Relaxed), 1);
+        assert_eq!(semaphore.available_permits(), 1);
+        assert!(acquire.poll_once(&waker).is_ready());
+
+        // Queue empty and balance positive again: the lock is not needed.
+        let queue = semaphore.waiters.lock();
+        semaphore.release(1);
+        drop(queue);
+        assert_eq!(semaphore.available_permits(), 2);
+    }
+
+    #[test]
+    fn locked_path_adds_to_a_positive_balance() {
+        let semaphore = Semaphore::new(2);
+
+        // A release that observed zero can find a positive balance once it holds the lock.
+        semaphore.insert_permits_with_lock(3, semaphore.waiters.lock());
+        assert_eq!(semaphore.available_permits(), 5);
+    }
+
+    #[test]
+    fn locked_overflow_panics_without_adding_permits() {
+        let semaphore = Semaphore::new(usize::MAX);
+        let result = panic::catch_unwind(AssertUnwindSafe(|| {
+            semaphore.insert_permits_with_lock(1, semaphore.waiters.lock());
+        }));
+        assert!(result.is_err());
+        assert_eq!(semaphore.available_permits(), usize::MAX);
+        assert!(semaphore.waiters.lock().is_empty());
+    }
+
+    #[test]
+    fn release_all_held_hands_permits_to_the_queue() {
+        let semaphore = Semaphore::new(1);
+        let counter = Arc::new(WakeCounter(AtomicUsize::new(0)));
+        let waker = Waker::from(counter.clone());
+        assert!(semaphore.try_acquire(1));
+        let mut acquire = semaphore.poll_acquire(1);
+        assert!(acquire.poll_once(&waker).is_pending());
+
+        semaphore.release_all_held(1);
+        assert_eq!(counter.0.load(Ordering::Relaxed), 1);
+        assert_eq!(semaphore.available_permits(), 0);
+        assert!(acquire.poll_once(&waker).is_ready());
+
+        semaphore.release_all_held(1);
+        assert_eq!(semaphore.available_permits(), 1);
+    }
+
     #[test]
     fn release_distributes_permits_across_multiple_batches() {
         const WAITER_COUNT: usize = WakerBatch::STACK_SIZE * 2 + 1;

asyncband/src/mutex/mod.rs

@@ -277,7 +277,7 @@ unsafe impl<T: ?Sized + Send + Sync> Sync for MutexGuard<'_, T> {}
 
 impl<T: ?Sized> Drop for MutexGuard<'_, T> {
     fn drop(&mut self) {
-        self.lock.s.release(1);
+        self.lock.s.release_all_held(1);
     }
 }
 
@@ -442,7 +442,7 @@ unsafe impl<T: ?Sized + Send + Sync> Sync for OwnedMutexGuard<T> {}
 
 impl<T: ?Sized> Drop for OwnedMutexGuard<T> {
     fn drop(&mut self) {
-        self.lock.s.release(1);
+        self.lock.s.release_all_held(1);
     }
 }
 
@@ -638,7 +638,7 @@ unsafe impl<T: ?Sized + Sync> Sync for MappedMutexGuard<'_, T> {}
 
 impl<T: ?Sized> Drop for MappedMutexGuard<'_, T> {
     fn drop(&mut self) {
-        self.s.release(1);
+        self.s.release_all_held(1);
     }
 }
 
@@ -837,7 +837,7 @@ unsafe impl<T: ?Sized + Send + Sync, U: ?Sized + Send + Sync> Sync for OwnedMapp
 
 impl<T: ?Sized, U: ?Sized> Drop for OwnedMappedMutexGuard<T, U> {
     fn drop(&mut self) {
-        self.lock.s.release(1);
+        self.lock.s.release_all_held(1);
     }
 }
 

benchmarks/benches/primitives/pool/bounded.rs

@@ -60,6 +60,19 @@ fn bounded_warm_get_and_return(bencher: Bencher) {
     });
 }
 
+#[divan::bench]
+fn bounded_warm_get_and_return_with_spare_capacity(bencher: Bencher) {
+    let pool = bounded::Pool::new(bounded::PoolConfig::new(8), Manager);
+    let mut context = bench_context();
+    drop(poll_ready(pool.get(), &mut context).unwrap());
+
+    bencher.bench_local(|| {
+        let object = poll_ready(pool.get(), &mut context).unwrap();
+        black_box(*object);
+        drop(object);
+    });
+}
+
 #[divan::bench]
 fn bounded_contended_handoff(bencher: Bencher) {
     let pool = bounded::Pool::new(bounded::PoolConfig::new(1), Manager);

benchmarks/benches/primitives/rwlock/read.rs

@@ -38,3 +38,14 @@ fn eight_reads_then_one_write(bencher: Bencher) {
         black_box(*guard)
     });
 }
+
+#[divan::bench]
+fn read_reuse(bencher: Bencher) {
+    let lock = RwLock::new(0usize);
+    let mut context = bench_context();
+
+    bencher.bench_local(|| {
+        let guard = poll_ready(lock.read(), &mut context);
+        black_box(*guard)
+    });
+}

benchmarks/benches/primitives/semaphore/acquire.rs

@@ -54,6 +54,15 @@ fn try_acquire_release(bencher: Bencher) {
     });
 }
 
+#[divan::bench]
+fn try_acquire_release_with_spare_permits(bencher: Bencher) {
+    let semaphore = Semaphore::new(8);
+
+    bencher.bench_local(|| {
+        drop(black_box(semaphore.try_acquire(black_box(2)).unwrap()));
+    });
+}
+
 #[divan::bench]
 fn owned_try_acquire_release(bencher: Bencher) {
     let semaphore = Arc::new(Semaphore::new(8));