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

CHANGELOG.md

@@ -41,6 +41,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 
+* Add a lock-free fast path to `Semaphore` release for a positive balance, taken by `RwLock` read-guard drops, multi-permit `Semaphore` releases, and bounded pool returns below capacity. `Mutex` releases are unchanged.
 * Allow `watch::channel` to store non-`Clone` values for publication and change notification; only owning reads through `Receiver::get` and `Receiver::recv` require `Clone`.
 * Finish releasing buffered bounded MPSC messages even if one message destructor panics.
 * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy.

asyncband/src/internal/semaphore.rs

@@ -198,6 +198,30 @@ 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.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());
         }
@@ -238,6 +262,24 @@ impl Semaphore {
         wake_all(wakers.into_iter());
     }
 
+    /// 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(
         &self,
         mut rem: usize,
@@ -286,13 +328,17 @@ impl Semaphore {
                 if rem > 0 && waiters.is_empty() {
                     // Retire the remainder before the overflow check so unwinding cannot retry it.
                     let added = std::mem::take(&mut rem);
-                    // The lock serializes additions; concurrent operations can only remove permits.
-                    let current = self.permits.load(Ordering::Relaxed);
-                    assert!(
-                        current.checked_add(added).is_some(),
-                        "number of added permits ({added}) would overflow usize::MAX (prev: {current})"
-                    );
-                    self.permits.fetch_add(added, 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(added, Ordering::Release);
+                    } else {
+                        while let Err(actual) = self.try_add_to_balance(current, added) {
+                            current = actual;
+                        }
+                    }
                 }
 
                 // Neither wake callbacks nor destruction of the taken waker run under this lock.
@@ -502,6 +548,78 @@ 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 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_does_not_retry_during_unwinding() {
+        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_to_all_waiters() {
         const WAITER_COUNT: usize = 35;

asyncband/src/mutex/mod.rs

@@ -279,7 +279,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);
     }
 }
 
@@ -444,7 +444,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);
     }
 }
 
@@ -639,7 +639,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);
     }
 }
 
@@ -845,7 +845,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) {
         // Release the lock by calling release on the semaphore
-        self.lock.s.release(1);
+        self.lock.s.release_all_held(1);
     }
 }
 

benchmarks/asyncband/pool/bounded.rs

@@ -71,6 +71,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/asyncband/rwlock/read.rs

@@ -39,3 +39,14 @@ fn read_heavy_reuse(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/asyncband/semaphore/acquire.rs

@@ -69,6 +69,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));

tests-integration/tests/semaphore_ordering_test.rs

@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! An acquire synchronizes with the release whose release sequence it reads from, and every
+//! write to the balance is a read-modify-write, so that sequence extends through later releases
+//! and acquisitions. A plain store would end it, and Miri would report a data race here.
+//!
+//! The race shows only on schedules where the final acquire reads the latest balance rather than
+//! a stale one, so the test runs under several seeds.
+
+use std::cell::UnsafeCell;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::thread;
+
+use asyncband::semaphore::Semaphore;
+
+struct Data(UnsafeCell<u64>);
+
+// SAFETY: accesses are ordered only through the semaphore, which is the property under test.
+unsafe impl Sync for Data {}
+
+/// Sequences the steps in time without adding a happens-before edge.
+fn wait_for(step: &AtomicUsize, value: usize) {
+    while step.load(Ordering::Relaxed) < value {
+        thread::yield_now();
+    }
+}
+
+#[test]
+fn acquire_synchronizes_with_every_earlier_release() {
+    let semaphore = Arc::new(Semaphore::new(1));
+    let data = Arc::new(Data(UnsafeCell::new(0)));
+    let step = Arc::new(AtomicUsize::new(0));
+
+    // Writes, then releases onto a positive balance.
+    let writer = {
+        let (semaphore, data, step) = (semaphore.clone(), data.clone(), step.clone());
+        thread::spawn(move || {
+            // SAFETY: the reader is ordered after this write by the semaphore.
+            unsafe { *data.0.get() = 42 };
+            semaphore.release(1);
+            step.store(1, Ordering::Relaxed);
+        })
+    };
+
+    // Takes every permit, so the balance is zero.
+    let drainer = {
+        let (semaphore, step) = (semaphore.clone(), step.clone());
+        thread::spawn(move || {
+            wait_for(&step, 1);
+            loop {
+                if let Some(permit) = semaphore.try_acquire(2) {
+                    permit.forget();
+                    break;
+                }
+                thread::yield_now();
+            }
+            step.store(2, Ordering::Relaxed);
+        })
+    };
+
+    // Releases onto the zero balance without any synchronization with the writer.
+    let releaser = {
+        let (semaphore, step) = (semaphore.clone(), step.clone());
+        thread::spawn(move || {
+            wait_for(&step, 2);
+            semaphore.release(1);
+            step.store(3, Ordering::Relaxed);
+        })
+    };
+
+    // Acquires the releaser's permit and reads the data.
+    let reader = {
+        let (semaphore, data, step) = (semaphore.clone(), data.clone(), step.clone());
+        thread::spawn(move || {
+            wait_for(&step, 3);
+            let permit = loop {
+                if let Some(permit) = semaphore.try_acquire(1) {
+                    break permit;
+                }
+                thread::yield_now();
+            };
+            // SAFETY: every release before this acquire happens-before it.
+            let value = unsafe { *data.0.get() };
+            drop(permit);
+            value
+        })
+    };
+
+    writer.join().unwrap();
+    drainer.join().unwrap();
+    releaser.join().unwrap();
+    assert_eq!(reader.join().unwrap(), 42);
+}