diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ab5099..bf413a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,10 @@ All notable changes to this project will be documented in this file. * Add bounded MPMC `reserve` and `try_reserve` methods returning a borrowed `Permit`, so callers can wait for capacity before constructing a value; sends and reservations receive capacity in wait-queue order, and unused permits release it. +### Improvements + +* Reduce `Semaphore::release` latency when it notifies waiting acquirers. + ## v0.7.3 (2026-09-29) ### New features diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 6c7cc4a..66a1ec6 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -25,6 +25,8 @@ // https://github.com/tokio-rs/tokio/blob/bb9d57017e100985f86d8ca41ac105ee9140423e/tokio/src/sync/batch_semaphore.rs use std::future::Future; +use std::panic; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::MutexGuard; use std::sync::atomic::AtomicUsize; @@ -188,62 +190,65 @@ impl Semaphore { mut rem: usize, waiters: MutexGuard<'_, WaitList>, ) { - let mut batch = WakerBatch::new(); + // A panicking wake callback must not stop the remaining permits from being distributed, so + // keep the first payload here and rethrow it once the release is complete. + let mut first_panic = None; let mut lock = Some(waiters); - - // One iterator covers the entire release. If a callback panics, `wake_all` keeps pulling - // batches during unwinding, so the remaining permits are still distributed and notified. - wake_all(std::iter::from_fn(|| { - loop { - if let Some(waker) = batch.next() { - return Some(waker); - } - if rem == 0 { - return None; - } - - let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); - while !batch.will_spill() { - match waiters.unlink_first_waiter(|node| { - if node.permits <= rem { - rem -= node.permits; - node.permits = 0; - true - } else { - node.permits -= rem; - rem = 0; - false + while rem > 0 { + let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); + let mut batch = WakerBatch::new(); + while !batch.will_spill() { + match waiters.unlink_first_waiter(|node| { + if node.permits <= rem { + rem -= node.permits; + node.permits = 0; + true + } else { + node.permits -= rem; + rem = 0; + false + } + }) { + None => break, + Some((id, waiter)) => { + let remove_now = waiter.waker.is_none(); + if let Some(waker) = waiter.waker.take() { + batch.push(waker); } - }) { - None => break, - Some((id, waiter)) => { - let remove_now = waiter.waker.is_none(); - if let Some(waker) = waiter.waker.take() { - batch.push(waker); - } - if remove_now { - waiters.remove_unlinked_waiter(id); - } + if remove_now { + waiters.remove_unlinked_waiter(id); } } } + } - 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); - } + 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); + } - // Neither wake callbacks nor destruction of the taken waker run under this lock. - drop(waiters); + drop(waiters); + + let result = panic::catch_unwind(AssertUnwindSafe(|| { + wake_all(&mut batch); + })); + if let Err(payload) = result { + if first_panic.is_none() { + first_panic = Some(payload); + } } - })); + } + + if let Some(payload) = first_panic { + panic::resume_unwind(payload); + } } }