From 376871715ac54180569580336cccd645b6d7d772 Mon Sep 17 00:00:00 2001 From: Orthur Date: Wed, 9 Sep 2026 12:25:47 -0400 Subject: [PATCH] perf(semaphore): avoid the queue lock for positive-balance releases --- CHANGELOG.md | 1 + asyncband/src/internal/semaphore.rs | 123 ++++++++++++++++-- asyncband/src/mutex/mod.rs | 8 +- benchmarks/asyncband/pool/bounded.rs | 13 ++ benchmarks/asyncband/rwlock/read.rs | 11 ++ benchmarks/asyncband/semaphore/acquire.rs | 9 ++ .../tests/semaphore_ordering_test.rs | 110 ++++++++++++++++ tests-integration/tests/semaphore_test.rs | 109 ++++++++++++++++ xtask/src/main.rs | 12 ++ 9 files changed, 384 insertions(+), 12 deletions(-) create mode 100644 tests-integration/tests/semaphore_ordering_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index ef8a4f69..5fa79fd2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,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. * 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. * Make completed and abandoned `Completion` waits lock-free while preserving cancellable pending registration. diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 8c2de0c1..e11c48c7 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -200,6 +200,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. + // + // The exchange is tried once; on failure the locked path takes over, so contending + // releasers queue on the lock instead of spinning on the balance. + 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()); } @@ -240,6 +264,24 @@ impl Semaphore { crate::internal::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, @@ -276,14 +318,18 @@ impl Semaphore { } if rem > 0 && waiters.is_empty() { - // Holding `waiters` serializes all permit additions. Concurrent operations can - // only remove permits, so the count cannot grow between this check and fetch_add. - 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); + // The positive-balance path of `release` can grow the balance while the lock is + // held, so the overflow check and the addition are one exchange. A zero balance + // cannot change under the lock, but the addition stays a read-modify-write so + // that it extends the release sequence of the previous release. + 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) { + current = actual; + } + } rem = 0; } @@ -491,6 +537,67 @@ 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 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_drains_more_than_one_wake_batch() { const WAITER_COUNT: usize = WAKE_BATCH_SIZE + 3; diff --git a/asyncband/src/mutex/mod.rs b/asyncband/src/mutex/mod.rs index ed049641..34aa79b7 100644 --- a/asyncband/src/mutex/mod.rs +++ b/asyncband/src/mutex/mod.rs @@ -279,7 +279,7 @@ unsafe impl Sync for MutexGuard<'_, T> {} impl 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 Sync for OwnedMutexGuard {} impl Drop for OwnedMutexGuard { fn drop(&mut self) { - self.lock.s.release(1); + self.lock.s.release_all_held(1); } } @@ -639,7 +639,7 @@ unsafe impl Sync for MappedMutexGuard<'_, T> {} impl Drop for MappedMutexGuard<'_, T> { fn drop(&mut self) { - self.s.release(1); + self.s.release_all_held(1); } } @@ -845,7 +845,7 @@ unsafe impl Sync for OwnedMapp impl Drop for OwnedMappedMutexGuard { 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); } } diff --git a/benchmarks/asyncband/pool/bounded.rs b/benchmarks/asyncband/pool/bounded.rs index 7441c260..e6148490 100644 --- a/benchmarks/asyncband/pool/bounded.rs +++ b/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); diff --git a/benchmarks/asyncband/rwlock/read.rs b/benchmarks/asyncband/rwlock/read.rs index 692dbafb..c4f1c151 100644 --- a/benchmarks/asyncband/rwlock/read.rs +++ b/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) + }); +} diff --git a/benchmarks/asyncband/semaphore/acquire.rs b/benchmarks/asyncband/semaphore/acquire.rs index 9d3f31d2..447ad1dc 100644 --- a/benchmarks/asyncband/semaphore/acquire.rs +++ b/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)); diff --git a/tests-integration/tests/semaphore_ordering_test.rs b/tests-integration/tests/semaphore_ordering_test.rs new file mode 100644 index 00000000..70d7a2fc --- /dev/null +++ b/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); + +// 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); +} diff --git a/tests-integration/tests/semaphore_test.rs b/tests-integration/tests/semaphore_test.rs index a3f4f675..f416a417 100644 --- a/tests-integration/tests/semaphore_test.rs +++ b/tests-integration/tests/semaphore_test.rs @@ -17,11 +17,16 @@ use std::pin::pin; use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use std::task::Wake; use std::task::Waker; +use std::thread; +use asyncband::blocking::FutureExt; use asyncband::semaphore::Semaphore; use tests_integration::PanicWake; use tests_integration::WakeCounter; @@ -272,3 +277,107 @@ fn reduce_permits_takes_priority_over_pending_acquires() { drop(permit); assert_eq!(s.available_permits(), 1); } + +/// Releases race acquisitions that drain the balance to zero and link waiters. The permit count +/// must be conserved, and no more permits than the capacity may ever be held at once. +#[test] +fn concurrent_releases_conserve_permits() { + const PERMITS: usize = 3; + const THREADS: usize = 8; + const ITERATIONS: usize = 4_000; + + fn next(state: &mut u64) -> usize { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (*state >> 33) as usize + } + + let semaphore = Arc::new(Semaphore::new(PERMITS)); + let held = Arc::new(AtomicUsize::new(0)); + let start = Arc::new(Barrier::new(THREADS)); + let workers = (0..THREADS) + .map(|seed| { + let semaphore = semaphore.clone(); + let held = held.clone(); + let start = start.clone(); + thread::spawn(move || { + let mut state = seed as u64 + 1; + start.wait(); + let hold = |n: usize, state: &mut u64| { + let now = held.fetch_add(n, Ordering::AcqRel) + n; + assert!(now <= PERMITS, "{now} permits held at once"); + for _ in 0..next(state) % 8 { + std::hint::spin_loop(); + } + held.fetch_sub(n, Ordering::AcqRel); + }; + for _ in 0..ITERATIONS { + let n = next(&mut state) % PERMITS + 1; + match next(&mut state) % 5 { + 0 => { + if let Some(permit) = semaphore.try_acquire(n) { + hold(n, &mut state); + drop(permit); + } + } + 1 => { + let permit = FutureExt::block_on(semaphore.acquire(n)); + hold(n, &mut state); + permit.forget(); + semaphore.release(n); + } + 2 => { + semaphore.reduce_permits(n); + semaphore.release(n); + } + 3 => { + let mut acquire = pin!(semaphore.acquire(n)); + let poll = poll_with(acquire.as_mut(), Waker::noop()); + if let Poll::Ready(permit) = poll { + hold(n, &mut state); + drop(permit); + } + } + _ => { + let permit = FutureExt::block_on(semaphore.acquire(n)); + hold(n, &mut state); + drop(permit); + } + } + } + }) + }) + .collect::>(); + + for worker in workers { + worker.join().unwrap(); + } + assert_eq!(semaphore.available_permits(), PERMITS); + assert!(semaphore.try_acquire(PERMITS).is_some()); +} + +/// Two releases race on a balance one below `usize::MAX`: exactly one may succeed, and the other +/// must panic before adding anything, whichever path each of them takes. +#[test] +fn concurrent_releases_at_the_limit_panic_exactly_once() { + let semaphore = Arc::new(Semaphore::new(usize::MAX - 1)); + let start = Arc::new(Barrier::new(2)); + let workers = (0..2) + .map(|_| { + let semaphore = semaphore.clone(); + let start = start.clone(); + thread::spawn(move || { + start.wait(); + std::panic::catch_unwind(|| semaphore.release(1)).is_ok() + }) + }) + .collect::>(); + let succeeded = workers + .into_iter() + .map(|worker| worker.join().unwrap()) + .filter(|succeeded| *succeeded) + .count(); + assert_eq!(succeeded, 1); + assert_eq!(semaphore.available_permits(), usize::MAX); +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 51fb9866..dce50bd2 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -125,6 +125,11 @@ impl CommandMiri { "tests-integration", &["--test", "phaser_test"], )); + run_command(make_miri_cmd_with_seeds( + "tests-integration", + &["--test", "semaphore_ordering_test"], + "0..8", + )); } } @@ -425,6 +430,13 @@ fn make_miri_cmd(package: &str, target: &[&str]) -> StdCommand { cmd } +/// Runs a Miri target under a range of seeds, for tests whose failure depends on the schedule. +fn make_miri_cmd_with_seeds(package: &str, target: &[&str], seeds: &str) -> StdCommand { + let mut cmd = make_miri_cmd(package, target); + cmd.env("MIRIFLAGS", format!("-Zmiri-many-seeds={seeds}")); + cmd +} + fn make_format_cmd(fix: bool) -> StdCommand { let mut cmd = find_command("cargo"); cmd.args(["+nightly", "fmt", "--all"]);