From 88f50bd9a8399c260fa9e7bd540d35a611da60df Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 5 Sep 2026 13:04:09 -0400 Subject: [PATCH] feat(capacity-limiter): explore a borrower-aware capacity limiter --- asyncband/Cargo.toml | 1 + asyncband/src/capacity_limiter/mod.rs | 490 ++++++++++++++++++ asyncband/src/capacity_limiter/regression.rs | 370 +++++++++++++ asyncband/src/capacity_limiter/tests.rs | 241 +++++++++ asyncband/src/internal/mod.rs | 3 + asyncband/src/lib.rs | 2 + benchmarks/Cargo.toml | 1 + .../asyncband/capacity_limiter/contended.rs | 137 +++++ benchmarks/asyncband/capacity_limiter/mod.rs | 22 + .../asyncband/capacity_limiter/uncontended.rs | 117 +++++ benchmarks/asyncband/main.rs | 1 + tests-integration/Cargo.toml | 1 + .../tests/capacity_limiter_test.rs | 324 ++++++++++++ tests-integration/tests/traits_test.rs | 11 + 14 files changed, 1721 insertions(+) create mode 100644 asyncband/src/capacity_limiter/mod.rs create mode 100644 asyncband/src/capacity_limiter/regression.rs create mode 100644 asyncband/src/capacity_limiter/tests.rs create mode 100644 benchmarks/asyncband/capacity_limiter/contended.rs create mode 100644 benchmarks/asyncband/capacity_limiter/mod.rs create mode 100644 benchmarks/asyncband/capacity_limiter/uncontended.rs create mode 100644 tests-integration/tests/capacity_limiter_test.rs diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index d2d74f45..9fa37870 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -47,6 +47,7 @@ default = [] barrier = [] blocking = [] broadcast = [] +capacity-limiter = [] completion = [] condvar = ["mutex"] event = [] diff --git a/asyncband/src/capacity_limiter/mod.rs b/asyncband/src/capacity_limiter/mod.rs new file mode 100644 index 00000000..3e91dc6f --- /dev/null +++ b/asyncband/src/capacity_limiter/mod.rs @@ -0,0 +1,490 @@ +// 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. + +//! Limits concurrent work to one permit per caller-supplied borrower identity. +//! +//! [`CapacityLimiter`] keeps the identity registry, available capacity, resize deficit, and FIFO +//! waiter queue under one mutex. [`Permit`] returns capacity and removes its identity when dropped. +//! Async acquisitions register on first poll; a duplicate pending or held identity is rejected. +//! +//! This is a prototype for issue #224. A concrete workload is still needed to justify adding it +//! to the crate. The existing semaphore remains sufficient when borrower checks are unnecessary. +//! +//! # Identity and callbacks +//! +//! Identities are explicit `Eq + Hash + Clone` values, not runtime task identities. Key `Hash` and +//! `Eq` implementations must be stable, non-panicking, and must not re-enter the limiter: they run +//! under its state lock. Key cloning and destruction, and waker callbacks, run outside that lock. +//! +//! # Diagnostics +//! +//! `borrowed()` counts permits delivered to callers. `waiting()` counts registered acquisitions +//! not yet delivered, including committed grants awaiting another poll. Each accessor is an +//! individual snapshot; calling several accessors does not produce a combined atomic snapshot. +//! +//! Debug output for the limiter contains counts only. A permit or acquisition prints its own +//! borrower, without traversing the limiter or revealing other borrower identities. +//! +//! # Example +//! +//! ``` +//! use asyncband::capacity_limiter::CapacityLimiter; +//! use asyncband::capacity_limiter::TryAcquireError; +//! +//! let limiter = CapacityLimiter::new(2); +//! let permit = limiter.try_acquire("request-1").unwrap(); +//! assert_eq!( +//! limiter.try_acquire("request-1").unwrap_err(), +//! TryAcquireError::AlreadyBorrowed +//! ); +//! limiter.set_total(0); // The existing permit remains valid. +//! drop(permit); // Repays the deficit. +//! assert_eq!(limiter.available(), 0); +//! limiter.set_total(1); +//! assert!(limiter.try_acquire("request-1").is_ok()); +//! ``` + +use std::collections::HashSet; +use std::fmt; +use std::future::Future; +use std::hash::Hash; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; +use crate::internal::wake_all; + +#[derive(Debug)] +struct WaitNode { + /// Set by the releasing side once this waiter has been handed a token. + granted: bool, + waker: Option, +} + +#[derive(Debug)] +struct State { + permits: usize, + total: usize, + /// Tokens owed to a shrink that could not be satisfied from available capacity. + deficit: usize, + borrowed: usize, + borrowers: HashSet, + waiters: WaitList, +} + +impl State { + /// Returns one token to the limiter and reports the waiter to wake, if any. + /// + /// The caller must wake outside the lock. + fn release_one(&mut self) -> Option { + if self.deficit > 0 { + self.deficit -= 1; + return None; + } + + match self.waiters.unlink_first_waiter(|node| { + node.granted = true; + true + }) { + // The node stays addressable until the waiter polls or is dropped. + Some((_, node)) => node.waker.take(), + None => { + self.permits += 1; + None + } + } + } +} + +/// A capacity limiter whose registry, capacity, and waiter queue share one lock. +pub struct CapacityLimiter { + state: Mutex>, +} + +impl fmt::Debug for CapacityLimiter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let (total, available, borrowed, waiting) = { + let state = self.state.lock(); + ( + state.total, + state.permits, + state.borrowed, + state.borrowers.len() - state.borrowed, + ) + }; + f.debug_struct("CapacityLimiter") + .field("total", &total) + .field("available", &available) + .field("borrowed", &borrowed) + .field("waiting", &waiting) + .finish_non_exhaustive() + } +} + +impl CapacityLimiter { + /// Creates a limiter that admits `total` concurrent borrowers. + pub fn new(total: usize) -> Self { + Self { + state: Mutex::new(State { + permits: total, + total, + deficit: 0, + borrowed: 0, + borrowers: HashSet::new(), + waiters: WaitList::new(), + }), + } + } + + /// Returns the configured total capacity. + /// + /// This is a snapshot taken under the same lock as permit accounting. + pub fn total(&self) -> usize { + self.state.lock().total + } + + /// Returns the number of permits delivered to callers and not yet dropped. + /// + /// Committed grants awaiting another poll are not counted here. + pub fn borrowed(&self) -> usize { + self.state.lock().borrowed + } + + /// Returns the number of tokens available without waiting. + pub fn available(&self) -> usize { + self.state.lock().permits + } + + /// Returns the number of registered acquisitions not yet delivered to callers. + /// + /// This includes committed grants awaiting another poll. The count is exact at the instant + /// this method holds the lock; it is not the length of the ungranted FIFO queue. + pub fn waiting(&self) -> usize { + let state = self.state.lock(); + state.borrowers.len() - state.borrowed + } + + /// Sets the total capacity without revoking existing grants or delivered permits. + /// + /// A shrink first consumes available capacity, then records a deficit repaid by releases. + /// Growth first repays that deficit, then grants queued borrowers in FIFO order. + /// Zero and `usize::MAX` are valid totals. + /// + /// Concurrent calls commit under the state lock. Wake callbacks run after unlocking and may + /// resize reentrantly. The adjustment is committed before waking; if one callback panics, + /// remaining notifications are still attempted before the panic continues. + pub fn set_total(&self, total: usize) { + let mut wakers = Vec::new(); + { + let mut state = self.state.lock(); + let previous = state.total; + state.total = total; + + if total > previous { + let mut added = total - previous; + let repaid = added.min(state.deficit); + state.deficit -= repaid; + added -= repaid; + while added > 0 && !state.waiters.is_empty() { + wakers.extend(state.release_one()); + added -= 1; + } + state.permits += added; + } else { + let shortfall = previous - total; + let taken = shortfall.min(state.permits); + state.permits -= taken; + state.deficit += shortfall - taken; + } + } + + wake_all(wakers.into_iter()); + } +} + +impl CapacityLimiter { + /// Acquires a token for `borrower`, waiting until capacity is available. + /// + /// Registration and duplicate checking happen on first poll. A borrower already queued or + /// holding a permit is rejected with [`AlreadyBorrowed`] without waiting. + /// + /// # Cancel safety + /// + /// Dropping a pending acquisition removes its identity and queue position. If it was already + /// granted a token, the token is returned or passed to the next queued borrower. + pub fn acquire(&self, borrower: B) -> Acquire<'_, B> { + Acquire { + limiter: self, + borrower: Some(borrower), + id: None, + registered: false, + } + } + + /// Attempts to acquire a token for `borrower` without waiting. + pub fn try_acquire(&self, borrower: B) -> Result, TryAcquireError> { + let key = borrower.clone(); + let mut state = self.state.lock(); + + // Check before insertion so a duplicate cannot replace and drop a stored key under the + // lock. This intentionally pays for separate lookup and insertion on the success path. + if state.borrowers.contains(&borrower) { + return Err(TryAcquireError::AlreadyBorrowed); + } + if state.permits == 0 { + return Err(TryAcquireError::NoCapacity); + } + state.borrowers.insert(key); + + state.permits -= 1; + state.borrowed += 1; + drop(state); + + Ok(Permit { + limiter: self, + borrower: Some(borrower), + }) + } +} + +/// The future returned by [`CapacityLimiter::acquire`]. +#[must_use = "futures do nothing unless polled"] +pub struct Acquire<'a, B: Eq + Hash + Clone> { + limiter: &'a CapacityLimiter, + /// Taken once the token is handed to a permit, so `Drop` knows whether it still owns cleanup. + borrower: Option, + id: Option, + registered: bool, +} + +impl fmt::Debug for Acquire<'_, B> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Acquire") + .field("borrower", &self.borrower) + .field("registered", &self.registered) + .finish_non_exhaustive() + } +} + +// The future holds no self-references; its waiter node lives in the limiter's arena. +impl Unpin for Acquire<'_, B> {} + +impl<'a, B: Eq + Hash + Clone> Future for Acquire<'a, B> { + type Output = Result, AlreadyBorrowed>; + + fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let limiter = this.limiter; + let mut key = if this.registered { + None + } else { + Some(this.borrower.as_ref().expect("pending borrower").clone()) + }; + let mut prepared_waker = None; + loop { + let mut state = limiter.state.lock(); + if !this.registered { + let borrower = this.borrower.as_ref().expect("pending borrower"); + // As in try_acquire, avoid replacing and dropping a duplicate key under the lock. + if state.borrowers.contains(borrower) { + return Poll::Ready(Err(AlreadyBorrowed)); + } + state + .borrowers + .insert(key.take().expect("unregistered borrower key")); + this.registered = true; + if state.permits > 0 { + state.permits -= 1; + state.borrowed += 1; + drop(state); + return Poll::Ready(Ok(this.take_permit())); + } + // Register before cloning: a reentrant release can commit this waiter while the + // lock is released for the callback. The next iteration rechecks the grant. + this.id = Some(state.waiters.push_back(WaitNode { + granted: false, + waker: None, + })); + } + + let id = this.id.expect("pending waiter id"); + if state.waiters.waiter_mut(id).granted { + let node = state.waiters.remove_unlinked_waiter(id); + state.borrowed += 1; + this.id = None; + drop(state); + let permit = this.take_permit(); + drop(node); + return Poll::Ready(Ok(permit)); + } + + let waker = &mut state.waiters.waiter_mut(id).waker; + if let Some(prepared) = prepared_waker.take() { + let old = waker.replace(prepared); + drop(state); + drop(old); + return Poll::Pending; + } + if waker + .as_ref() + .is_some_and(|held| held.will_wake(context.waker())) + { + return Poll::Pending; + } + drop(state); + prepared_waker = Some(context.waker().clone()); + } + } +} + +impl<'a, B: Eq + Hash + Clone> Acquire<'a, B> { + fn take_permit(&mut self) -> Permit<'a, B> { + Permit { + limiter: self.limiter, + borrower: Some( + self.borrower + .take() + .expect("a granted acquire still owns its borrower"), + ), + } + } +} + +impl Drop for Acquire<'_, B> { + fn drop(&mut self) { + // `None` means the token was handed to a permit, which owns the cleanup from here. + let Some(borrower) = self.borrower.take() else { + return; + }; + // Never registered: either never polled, or rejected as a duplicate of someone else's + // entry, which must not be removed here. + if !self.registered { + return; + } + + let (waker, node, key) = { + let mut state = self.limiter.state.lock(); + let key = state.borrowers.take(&borrower); + let node = self.id.take().map(|id| { + state.waiters.unlink_waiter(id, |_| true); + state.waiters.remove_unlinked_waiter(id) + }); + let waker = if node.as_ref().is_some_and(|node| node.granted) { + state.release_one() + } else { + None + }; + (waker, node, key) + }; + if let Some(waker) = waker { + waker.wake(); + } + drop(node); + drop(key); + } +} + +/// A token borrowed from a [`CapacityLimiter`]. +#[must_use = "tokens are returned immediately when dropped"] +pub struct Permit<'a, B: Eq + Hash + Clone> { + limiter: &'a CapacityLimiter, + borrower: Option, +} + +impl fmt::Debug for Permit<'_, B> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Permit") + .field("borrower", &self.borrower) + .finish_non_exhaustive() + } +} + +impl Permit<'_, B> { + /// Returns the borrower this token was acquired for. + pub fn borrower(&self) -> &B { + self.borrower + .as_ref() + .expect("a permit holds its borrower until it is dropped") + } +} + +impl Drop for Permit<'_, B> { + fn drop(&mut self) { + let Some(borrower) = self.borrower.take() else { + return; + }; + + let (waker, key) = { + let mut state = self.limiter.state.lock(); + let key = state.borrowers.take(&borrower); + state.borrowed -= 1; + (state.release_one(), key) + }; + if let Some(waker) = waker { + waker.wake(); + } + drop(key); + } +} + +/// The error returned when a borrower attempts to hold two tokens at once. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AlreadyBorrowed; + +impl fmt::Display for AlreadyBorrowed { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "borrower already holds a token from this limiter") + } +} + +impl std::error::Error for AlreadyBorrowed {} + +/// The error returned by [`CapacityLimiter::try_acquire`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TryAcquireError { + /// No capacity is available without waiting. + NoCapacity, + /// The borrower already holds a token from this limiter. + AlreadyBorrowed, +} + +impl From for TryAcquireError { + fn from(_: AlreadyBorrowed) -> Self { + TryAcquireError::AlreadyBorrowed + } +} + +impl fmt::Display for TryAcquireError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TryAcquireError::NoCapacity => write!(f, "no capacity available"), + TryAcquireError::AlreadyBorrowed => AlreadyBorrowed.fmt(f), + } + } +} + +impl std::error::Error for TryAcquireError {} + +#[cfg(test)] +mod tests; + +#[cfg(test)] +mod regression; diff --git a/asyncband/src/capacity_limiter/regression.rs b/asyncband/src/capacity_limiter/regression.rs new file mode 100644 index 00000000..de03d366 --- /dev/null +++ b/asyncband/src/capacity_limiter/regression.rs @@ -0,0 +1,370 @@ +// 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. + +use std::future::Future; +use std::mem::ManuallyDrop; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; +use std::task::Wake; +use std::task::Waker; +use std::thread; + +use super::CapacityLimiter; + +struct Callback(Box); + +impl Wake for Callback { + fn wake(self: Arc) { + (self.0)(); + } +} + +#[test] +fn committed_grants_remain_pending_until_delivery() { + let limiter = CapacityLimiter::new(0); + let mut acquire = Box::pin(limiter.acquire(1)); + assert!( + acquire + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + limiter.set_total(1); + assert_eq!(limiter.borrowed(), 0); + assert_eq!(limiter.waiting(), 1); + assert_eq!(limiter.available(), 0); + limiter.set_total(0); + let Poll::Ready(Ok(permit)) = acquire + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + else { + panic!("shrink must not revoke the committed grant"); + }; + assert_eq!(limiter.borrowed(), 1); + assert_eq!(limiter.waiting(), 0); + drop(permit); + assert_eq!(limiter.available(), 0); +} + +#[test] +fn large_resizes_and_deficit_do_not_loop_over_capacity() { + let limiter = CapacityLimiter::new(0); + limiter.set_total(usize::MAX); + let permit = limiter.try_acquire(1).unwrap(); + limiter.set_total(0); + limiter.set_total(usize::MAX); + assert_eq!(limiter.available(), usize::MAX - 1); + drop(permit); + assert_eq!(limiter.available(), usize::MAX); +} + +#[test] +fn wake_panic_still_attempts_all_notifications() { + let limiter = CapacityLimiter::new(0); + let calls = Arc::new(AtomicUsize::new(0)); + let mut pending = Vec::new(); + for id in 0..64 { + let calls = calls.clone(); + let waker = Waker::from(Arc::new(Callback(Box::new(move || { + calls.fetch_add(1, Ordering::Relaxed); + assert!(id != 0, "first wake panics"); + })))); + let mut future = Box::pin(limiter.acquire(id)); + assert!( + future + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + pending.push(future); + } + assert!(catch_unwind(AssertUnwindSafe(|| limiter.set_total(64))).is_err()); + assert_eq!(calls.load(Ordering::Relaxed), 64); + assert_eq!(limiter.waiting(), 64); + for mut future in pending { + assert!(matches!( + future + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok(_)) + )); + } + assert_eq!(limiter.available(), 64); + assert_eq!(limiter.waiting(), 0); +} + +// Arc's Wake adapter has no clone callback. This fixture exercises the full RawWaker contract. +// Each raw pointer owns one Arc, and all callbacks are Send + Sync. +struct CloneCallback(Box); + +unsafe fn clone_raw(data: *const ()) -> RawWaker { + // SAFETY: The pointer came from Arc::into_raw for CloneCallback. Cloning borrows that Arc. + let callback = ManuallyDrop::new(unsafe { Arc::from_raw(data.cast::()) }); + (callback.0)(); + let cloned = Arc::clone(&callback); + RawWaker::new(Arc::into_raw(cloned).cast(), &VTABLE) +} + +unsafe fn drop_raw(data: *const ()) { + // SAFETY: Consuming wake or drop releases exactly this pointer's owned Arc. + drop(unsafe { Arc::from_raw(data.cast::()) }); +} + +unsafe fn wake_by_ref_raw(_: *const ()) {} + +static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_raw, drop_raw, wake_by_ref_raw, drop_raw); + +fn clone_callback(callback: impl Fn() + Send + Sync + 'static) -> Waker { + let data = Arc::into_raw(Arc::new(CloneCallback(Box::new(callback)))).cast(); + // SAFETY: The vtable preserves Arc ownership and its callbacks are thread safe. + unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) } +} + +#[test] +fn clone_can_grant_before_registration_finishes() { + let limiter = Arc::new(CapacityLimiter::new(0)); + let observer = limiter.clone(); + let waker = clone_callback(move || { + assert_eq!(observer.waiting(), 1); + observer.set_total(1); + }); + let mut future = Box::pin(limiter.acquire(1)); + assert!(matches!( + future.as_mut().poll(&mut Context::from_waker(&waker)), + Poll::Ready(Ok(_)) + )); + assert_eq!(limiter.available(), 1); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn clone_panic_can_be_cancelled_without_leaking_identity() { + let limiter = CapacityLimiter::new(0); + let waker = clone_callback(|| panic!("clone panic")); + let mut future = Box::pin(limiter.acquire(1)); + assert!( + catch_unwind(AssertUnwindSafe(|| future + .as_mut() + .poll(&mut Context::from_waker(&waker)))) + .is_err() + ); + drop(future); + assert_eq!(limiter.waiting(), 0); + limiter.set_total(1); + assert!(limiter.try_acquire(1).is_ok()); +} + +struct DropCallback(Box); + +// This fixture needs an owned destructor; Waker::noop() cannot exercise drop callbacks. +#[allow(clippy::manual_noop_waker)] +impl Wake for DropCallback { + fn wake(self: Arc) {} +} + +impl Drop for DropCallback { + fn drop(&mut self) { + (self.0)(); + } +} + +#[test] +fn replacing_a_waker_allows_its_destructor_to_grant() { + let limiter = Arc::new(CapacityLimiter::new(0)); + let observer = limiter.clone(); + let old = Waker::from(Arc::new(DropCallback(Box::new(move || { + observer.set_total(1) + })))); + let mut future = Box::pin(limiter.acquire(1)); + assert!( + future + .as_mut() + .poll(&mut Context::from_waker(&old)) + .is_pending() + ); + drop(old); + assert!( + future + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + assert!(matches!( + future + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok(_)) + )); + assert_eq!(limiter.available(), 1); +} + +#[test] +fn cancellation_removes_identity_before_dropping_waker() { + let limiter = Arc::new(CapacityLimiter::new(0)); + let observer = limiter.clone(); + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let old = Waker::from(Arc::new(DropCallback(Box::new(move || { + assert_eq!(observer.waiting(), 0); + observer.set_total(1); + assert!(observer.try_acquire(1).is_ok()); + counter.fetch_add(1, Ordering::Relaxed); + })))); + let mut future = Box::pin(limiter.acquire(1)); + assert!( + future + .as_mut() + .poll(&mut Context::from_waker(&old)) + .is_pending() + ); + drop(old); + drop(future); + assert_eq!(calls.load(Ordering::Relaxed), 1); +} + +#[test] +fn concurrent_extreme_resizes_preserve_capacity() { + let limiter = CapacityLimiter::::new(usize::MAX); + thread::scope(|scope| { + for offset in 0..2 { + let limiter = &limiter; + scope.spawn(move || { + for round in 0..5000 { + let target = if (round + offset) % 2 == 0 { + 0 + } else { + usize::MAX + }; + limiter.set_total(target); + } + }); + } + }); + assert_eq!(limiter.available(), limiter.total()); + limiter.set_total(0); + assert_eq!(limiter.available(), 0); +} + +#[test] +fn wake_callback_can_resize_without_revoking_its_grant() { + let limiter = Arc::new(CapacityLimiter::new(0)); + let observer = limiter.clone(); + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let waker = Waker::from(Arc::new(Callback(Box::new(move || { + observer.set_total(0); + counter.fetch_add(1, Ordering::Relaxed); + })))); + let mut acquire = Box::pin(limiter.acquire(1)); + assert!( + acquire + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + limiter.set_total(1); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(limiter.total(), 0); + assert!(matches!( + acquire + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok(_)) + )); + assert_eq!(limiter.available(), 0); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn cancelling_a_grant_removes_identity_before_waking_the_next_borrower() { + let limiter = Arc::new(CapacityLimiter::new(0)); + let mut first = Box::pin(limiter.acquire("x")); + assert!( + first + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + + let observer = limiter.clone(); + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let waker = Waker::from(Arc::new(Callback(Box::new(move || { + assert_eq!(observer.waiting(), 1); + assert_eq!(observer.borrowed(), 0); + // x is no longer registered; y owns the committed capacity. This must not be a duplicate. + assert_eq!( + observer.try_acquire("x").unwrap_err(), + super::TryAcquireError::NoCapacity + ); + counter.fetch_add(1, Ordering::Relaxed); + })))); + let mut second = Box::pin(limiter.acquire("y")); + assert!( + second + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + limiter.set_total(1); // Only x receives a grant; y must be woken by cancellation below. + assert_eq!(calls.load(Ordering::Relaxed), 0); + drop(first); + assert_eq!(calls.load(Ordering::Relaxed), 1); + let Poll::Ready(Ok(permit)) = second + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + else { + panic!("cancelling x must pass its grant to y"); + }; + assert_eq!(permit.borrower(), &"y"); + drop(permit); + assert!(limiter.try_acquire("x").is_ok()); +} + +#[test] +fn debug_output_does_not_traverse_other_borrowers() { + let limiter = CapacityLimiter::new(2); + let first = limiter.try_acquire("private-first-key").unwrap(); + let second = limiter.try_acquire("private-second-key").unwrap(); + let mut pending = Box::pin(limiter.acquire("private-waiting-key")); + assert!( + pending + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + + let limiter_debug = format!("{limiter:?}"); + assert!(!limiter_debug.contains("private-")); + let permit_debug = format!("{first:?}"); + assert!(permit_debug.contains("private-first-key")); + assert!(!permit_debug.contains("private-second-key")); + assert!(!permit_debug.contains("private-waiting-key")); + let acquire_debug = format!("{pending:?}"); + assert!(acquire_debug.contains("private-waiting-key")); + assert!(!acquire_debug.contains("private-first-key")); + assert!(!acquire_debug.contains("private-second-key")); + drop(second); +} diff --git a/asyncband/src/capacity_limiter/tests.rs b/asyncband/src/capacity_limiter/tests.rs new file mode 100644 index 00000000..a58fccf9 --- /dev/null +++ b/asyncband/src/capacity_limiter/tests.rs @@ -0,0 +1,241 @@ +// 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. + +use std::pin::pin; + +use super::*; +use crate::capacity_limiter::TryAcquireError; + +// FIFO handoff, cancellation, and deficit accounting share the same state. + +fn poll_acquire<'a, B: Eq + Hash + Clone>( + acquire: Pin<&mut Acquire<'a, B>>, +) -> Poll, AlreadyBorrowed>> { + acquire.poll(&mut Context::from_waker(Waker::noop())) +} + +#[test] +fn capacity_bounds_concurrent_borrowers() { + let limiter = CapacityLimiter::new(2); + let a = limiter.try_acquire("a").unwrap(); + let b = limiter.try_acquire("b").unwrap(); + + assert_eq!(limiter.borrowed(), 2); + assert_eq!( + limiter.try_acquire("c").unwrap_err(), + TryAcquireError::NoCapacity + ); + + drop(a); + drop(b); + assert_eq!(limiter.available(), 2); + assert_eq!(limiter.borrowed(), 0); +} + +#[test] +fn one_token_per_borrower() { + let limiter = CapacityLimiter::new(4); + let held = limiter.try_acquire("a").unwrap(); + + assert!(limiter.available() > 0); + assert_eq!( + limiter.try_acquire("a").unwrap_err(), + TryAcquireError::AlreadyBorrowed + ); + + drop(held); + assert!(limiter.try_acquire("a").is_ok()); +} + +#[test] +fn queued_borrowers_are_granted_in_order() { + let limiter = CapacityLimiter::new(0); + + let mut first = pin!(limiter.acquire("a")); + assert!(poll_acquire(first.as_mut()).is_pending()); + let mut second = pin!(limiter.acquire("b")); + assert!(poll_acquire(second.as_mut()).is_pending()); + assert_eq!(limiter.waiting(), 2); + + limiter.set_total(1); + + // The token goes to the borrower that queued first. + assert!(poll_acquire(second.as_mut()).is_pending()); + let permit = match poll_acquire(first.as_mut()) { + Poll::Ready(result) => result.unwrap(), + Poll::Pending => panic!("the first waiter must be granted the token"), + }; + assert_eq!(permit.borrower(), &"a"); + + drop(permit); + assert!(matches!(poll_acquire(second.as_mut()), Poll::Ready(Ok(_)))); +} + +#[test] +fn duplicate_acquire_is_rejected_on_first_poll() { + let limiter = CapacityLimiter::new(1); + let _held = limiter.try_acquire("a").unwrap(); + + let mut duplicate = pin!(limiter.acquire("a")); + assert!(matches!( + poll_acquire(duplicate.as_mut()), + Poll::Ready(Err(AlreadyBorrowed)) + )); + + // The rejected acquire must not have disturbed the live registration. + assert_eq!(limiter.borrowed(), 1); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn cancelled_acquire_releases_the_identity() { + let limiter = CapacityLimiter::new(0); + + { + let mut queued = pin!(limiter.acquire("a")); + assert!(poll_acquire(queued.as_mut()).is_pending()); + assert_eq!(limiter.waiting(), 1); + } + + assert_eq!(limiter.waiting(), 0); + limiter.set_total(1); + assert!(limiter.try_acquire("a").is_ok()); +} + +#[test] +fn cancelling_after_a_grant_passes_the_token_on() { + let limiter = CapacityLimiter::new(1); + let held = limiter.try_acquire("a").unwrap(); + let mut next = pin!(limiter.acquire("c")); + + { + // Queued ahead of "c", so the released token is offered here first. + let mut queued = pin!(limiter.acquire("b")); + assert!(poll_acquire(queued.as_mut()).is_pending()); + assert!(poll_acquire(next.as_mut()).is_pending()); + + // "b" is handed the token, then cancelled at the end of this scope before taking it. + drop(held); + } + + // The token must reach "c" rather than vanish with the cancelled future. The permit has to + // be bound: dropping it as a temporary would return the token before it can be observed. + let permit = match poll_acquire(next.as_mut()) { + Poll::Ready(result) => result.unwrap(), + Poll::Pending => panic!("the cancelled grant must be passed to the next waiter"), + }; + assert_eq!(permit.borrower(), &"c"); + assert_eq!(limiter.borrowed(), 1); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn cancelling_a_lone_grant_returns_the_token() { + let limiter = CapacityLimiter::new(1); + let held = limiter.try_acquire("a").unwrap(); + + { + let mut queued = pin!(limiter.acquire("b")); + assert!(poll_acquire(queued.as_mut()).is_pending()); + drop(held); + } + + assert_eq!(limiter.available(), 1); + assert_eq!(limiter.borrowed(), 0); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn shrinking_never_revokes_borrowed_tokens() { + let limiter = CapacityLimiter::new(3); + let a = limiter.try_acquire("a").unwrap(); + let b = limiter.try_acquire("b").unwrap(); + let c = limiter.try_acquire("c").unwrap(); + + limiter.set_total(1); + assert_eq!(limiter.total(), 1); + assert_eq!(limiter.borrowed(), 3); + + drop(a); + assert_eq!(limiter.available(), 0); + drop(b); + assert_eq!(limiter.available(), 0); + + drop(c); + assert_eq!(limiter.available(), 1); +} + +#[test] +fn a_deficit_is_repaid_before_queued_borrowers() { + let limiter = CapacityLimiter::new(1); + let held = limiter.try_acquire("a").unwrap(); + + let mut queued = pin!(limiter.acquire("b")); + assert!(poll_acquire(queued.as_mut()).is_pending()); + + limiter.set_total(0); + drop(held); + + // The released token repays the shrink instead of admitting the waiter. + assert!(poll_acquire(queued.as_mut()).is_pending()); + assert_eq!(limiter.available(), 0); + + limiter.set_total(1); + assert!(matches!(poll_acquire(queued.as_mut()), Poll::Ready(Ok(_)))); +} +#[test] +fn cancelled_acquire_leaves_no_registration() { + let limiter = CapacityLimiter::new(0); + + { + let mut acquire = pin!(limiter.acquire("a")); + let mut context = Context::from_waker(Waker::noop()); + assert!(acquire.as_mut().poll(&mut context).is_pending()); + assert_eq!(limiter.state.lock().borrowers.len(), 1); + } + + assert!(limiter.state.lock().borrowers.is_empty()); + assert_eq!(limiter.state.lock().borrowed, 0); +} + +#[test] +fn rejected_try_acquire_leaves_no_registration() { + let limiter = CapacityLimiter::new(1); + let held = limiter.try_acquire("a").unwrap(); + + assert!(limiter.try_acquire("b").is_err()); + assert_eq!(limiter.state.lock().borrowers.len(), 1); + + drop(held); + assert!(limiter.state.lock().borrowers.is_empty()); +} + +#[test] +fn granted_permit_registers_exactly_once() { + let limiter = CapacityLimiter::new(2); + let first = limiter.try_acquire("a").unwrap(); + let second = limiter.try_acquire("b").unwrap(); + + assert_eq!(limiter.state.lock().borrowers.len(), 2); + assert_eq!(limiter.state.lock().borrowed, 2); + + drop(first); + drop(second); + + assert!(limiter.state.lock().borrowers.is_empty()); + assert_eq!(limiter.state.lock().borrowed, 0); +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index c6363791..76545a15 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -53,6 +53,7 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { pub(crate) mod atomic_waker; #[cfg(any( + feature = "capacity-limiter", feature = "barrier", feature = "broadcast", feature = "event", @@ -80,6 +81,7 @@ pub(crate) mod countdown; pub(crate) mod value_cell; #[cfg(any( + feature = "capacity-limiter", feature = "barrier", feature = "broadcast", feature = "event", @@ -107,6 +109,7 @@ pub(crate) mod mutex; pub(crate) mod semaphore; #[cfg(any( + feature = "capacity-limiter", feature = "event", feature = "mpsc", feature = "mutex", diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 93ce485e..c3c08a58 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -124,6 +124,8 @@ pub mod barrier; pub mod blocking; #[cfg(feature = "broadcast")] pub mod broadcast; +#[cfg(feature = "capacity-limiter")] +pub mod capacity_limiter; #[cfg(feature = "completion")] pub mod completion; #[cfg(feature = "condvar")] diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 2104a826..a6e40058 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -29,6 +29,7 @@ asyncband = { workspace = true, features = [ "barrier", "blocking", "broadcast", + "capacity-limiter", "completion", "condvar", "event", diff --git a/benchmarks/asyncband/capacity_limiter/contended.rs b/benchmarks/asyncband/capacity_limiter/contended.rs new file mode 100644 index 00000000..0debbe4d --- /dev/null +++ b/benchmarks/asyncband/capacity_limiter/contended.rs @@ -0,0 +1,137 @@ +// 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. + +// Successful try-acquire contention with a small set of distinct borrower keys. +// This measures implementation cost, not whether a core primitive is needed. + +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::thread; +use std::thread::JoinHandle; + +use asyncband::capacity_limiter::CapacityLimiter; +use asyncband::semaphore::Semaphore; +use divan::Bencher; +use divan::black_box; + +const WORKER_COUNTS: &[usize] = &[1, 2, 4, 8]; + +// Held constant across worker counts so a round always performs the same amount of work and the +// reported times compare directly. Every count in `WORKER_COUNTS` divides it exactly. +const OPS_PER_ROUND: usize = 1024; + +/// A fixed set of worker threads that hammer a shared primitive one round at a time. +/// +/// The threads are created once and parked on a barrier, so thread startup stays out of the +/// measured region and each timed iteration is exactly one round of `OPS_PER_ROUND` operations. +struct Contention { + start: Arc, + done: Arc, + stop: Arc, + workers: Vec>, +} + +impl Contention { + fn new(worker_count: usize, shared: Arc, op: fn(&T, usize)) -> Self + where + T: Send + Sync + 'static, + { + let start = Arc::new(Barrier::new(worker_count + 1)); + let done = Arc::new(Barrier::new(worker_count + 1)); + let stop = Arc::new(AtomicBool::new(false)); + let ops_per_worker = OPS_PER_ROUND / worker_count; + let mut workers = Vec::with_capacity(worker_count); + + for index in 0..worker_count { + let shared = shared.clone(); + let start = start.clone(); + let done = done.clone(); + let stop = stop.clone(); + + workers.push(thread::spawn(move || { + loop { + start.wait(); + if stop.load(Ordering::Acquire) { + return; + } + for _ in 0..ops_per_worker { + op(&shared, index); + } + done.wait(); + } + })); + } + + Self { + start, + done, + stop, + workers, + } + } + + fn round(&self) { + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for Contention { + fn drop(&mut self) { + // The workers are parked on `start` between rounds, so one more release lets them observe + // the stop flag and exit without reaching `done`. + self.stop.store(true, Ordering::Release); + self.start.wait(); + for worker in self.workers.drain(..) { + worker.join().expect("worker threads must not panic"); + } + } +} + +// Capacity equals the worker count so every attempt succeeds. The benchmark measures the cost of a +// successful acquire and release under contention, not the cost of being turned away. + +#[divan::bench(args = WORKER_COUNTS)] +fn semaphore_try_acquire_release(bencher: Bencher, workers: usize) { + let contention = Contention::new( + workers, + Arc::new(Semaphore::new(workers)), + |semaphore, _| { + drop(black_box( + semaphore.try_acquire(1).expect("capacity is available"), + )); + }, + ); + + bencher.bench_local(|| contention.round()); +} + +#[divan::bench(args = WORKER_COUNTS)] +fn limiter_try_acquire_release(bencher: Bencher, workers: usize) { + let limiter: Arc> = Arc::new(CapacityLimiter::new(workers)); + let contention = Contention::new(workers, limiter, |limiter, index| { + // Each worker owns a distinct borrower, so nothing is rejected for identity reasons and the + // registry stays at `workers` live entries. + drop(black_box( + limiter.try_acquire(index).expect("capacity is available"), + )); + }); + + bencher.bench_local(|| contention.round()); +} diff --git a/benchmarks/asyncband/capacity_limiter/mod.rs b/benchmarks/asyncband/capacity_limiter/mod.rs new file mode 100644 index 00000000..06b14edf --- /dev/null +++ b/benchmarks/asyncband/capacity_limiter/mod.rs @@ -0,0 +1,22 @@ +// 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. + +// Self-benchmarks with a bare semaphore as a cost reference. Identity tracking and permit +// accounting differ, so the gap does not isolate the cost of hashing or any particular lock. + +mod contended; +mod uncontended; diff --git a/benchmarks/asyncband/capacity_limiter/uncontended.rs b/benchmarks/asyncband/capacity_limiter/uncontended.rs new file mode 100644 index 00000000..2e32159b --- /dev/null +++ b/benchmarks/asyncband/capacity_limiter/uncontended.rs @@ -0,0 +1,117 @@ +// 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. + +use std::pin::pin; +use std::sync::Arc; + +use asyncband::capacity_limiter::CapacityLimiter; +use asyncband::semaphore::Semaphore; +use divan::Bencher; +use divan::black_box; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; + +// Instances are reused; construction is outside the timed operation. +// The semaphore is a cost reference with no borrower-identity contract. + +#[divan::bench] +fn semaphore_try_acquire_release(bencher: Bencher) { + let semaphore = Semaphore::new(1); + + bencher.bench_local(|| { + drop(black_box( + semaphore + .try_acquire(black_box(1)) + .expect("capacity is available"), + )); + }); +} + +#[divan::bench] +fn limiter_try_acquire_release(bencher: Bencher) { + let limiter = CapacityLimiter::new(1); + + bencher.bench_local(|| { + drop(black_box( + limiter + .try_acquire(black_box(7usize)) + .expect("capacity is available"), + )); + }); +} + +#[divan::bench] +fn limiter_try_acquire_release_arc_str(bencher: Bencher) { + let limiter = CapacityLimiter::new(1); + let borrower: Arc = Arc::from("tenant-0000000001"); + + bencher.bench_local(|| { + drop(black_box( + limiter + .try_acquire(black_box(borrower.clone())) + .expect("capacity is available"), + )); + }); +} + +#[divan::bench] +fn semaphore_handoff_reused(bencher: Bencher) { + let limiter = Semaphore::new(1); + let mut context = bench_context(); + bencher.bench_local(|| { + let held = limiter.try_acquire(black_box(1)).unwrap(); + let mut acquire = pin!(limiter.acquire(black_box(1))); + poll_pending(acquire.as_mut(), &mut context); + drop(held); + drop(poll_pinned_ready(acquire.as_mut(), &mut context)); + }); +} + +#[divan::bench] +fn limiter_handoff_reused(bencher: Bencher) { + let limiter = CapacityLimiter::new(1); + let mut context = bench_context(); + bencher.bench_local(|| { + let held = limiter.try_acquire(black_box(7usize)).unwrap(); + let mut acquire = pin!(limiter.acquire(black_box(8usize))); + poll_pending(acquire.as_mut(), &mut context); + drop(held); + drop(poll_pinned_ready(acquire.as_mut(), &mut context).unwrap()); + }); +} + +#[divan::bench] +fn semaphore_cancel_pending_reused(bencher: Bencher) { + let limiter = Semaphore::new(0); + let mut context = bench_context(); + bencher.bench_local(|| { + let mut acquire = pin!(limiter.acquire(black_box(1))); + poll_pending(acquire.as_mut(), &mut context); + }); +} + +#[divan::bench] +fn limiter_cancel_pending_reused(bencher: Bencher) { + let limiter = CapacityLimiter::new(0); + let mut context = bench_context(); + bencher.bench_local(|| { + let mut acquire = pin!(limiter.acquire(black_box(7usize))); + poll_pending(acquire.as_mut(), &mut context); + }); +} diff --git a/benchmarks/asyncband/main.rs b/benchmarks/asyncband/main.rs index 53ae706a..b5ada353 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -18,6 +18,7 @@ mod barrier; mod blocking; mod broadcast; +mod capacity_limiter; mod completion; mod condvar; mod event; diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 04df5894..292c5ba6 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -30,6 +30,7 @@ asyncband = { workspace = true, features = [ "barrier", "blocking", "broadcast", + "capacity-limiter", "completion", "condvar", "event", diff --git a/tests-integration/tests/capacity_limiter_test.rs b/tests-integration/tests/capacity_limiter_test.rs new file mode 100644 index 00000000..3d50a47e --- /dev/null +++ b/tests-integration/tests/capacity_limiter_test.rs @@ -0,0 +1,324 @@ +// 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. + +use std::pin::pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::task::Poll; + +use asyncband::capacity_limiter::AlreadyBorrowed; +use asyncband::capacity_limiter::CapacityLimiter; +use asyncband::capacity_limiter::TryAcquireError; +use tests_integration::poll_once; + +#[test] +fn zero_capacity_admits_nobody() { + let limiter = CapacityLimiter::new(0); + + assert_eq!(limiter.total(), 0); + assert_eq!( + limiter.try_acquire("a").unwrap_err(), + TryAcquireError::NoCapacity + ); +} + +#[test] +fn capacity_bounds_concurrent_borrowers() { + let limiter = CapacityLimiter::new(2); + + let a = limiter.try_acquire("a").unwrap(); + let b = limiter.try_acquire("b").unwrap(); + assert_eq!(limiter.borrowed(), 2); + assert_eq!(limiter.available(), 0); + assert_eq!( + limiter.try_acquire("c").unwrap_err(), + TryAcquireError::NoCapacity + ); + + drop(a); + let c = limiter.try_acquire("c").unwrap(); + assert_eq!(c.borrower(), &"c"); + + drop(b); + drop(c); + assert_eq!(limiter.borrowed(), 0); + assert_eq!(limiter.available(), 2); +} + +#[test] +fn one_token_per_borrower() { + let limiter = CapacityLimiter::new(4); + let held = limiter.try_acquire("a").unwrap(); + + // Capacity is available, so the rejection is about identity rather than capacity. + assert!(limiter.available() > 0); + assert_eq!( + limiter.try_acquire("a").unwrap_err(), + TryAcquireError::AlreadyBorrowed + ); + + drop(held); + assert!(limiter.try_acquire("a").is_ok()); +} + +#[tokio::test] +async fn duplicate_acquire_fails_without_waiting() { + let limiter = CapacityLimiter::new(1); + let _held = limiter.acquire("a").await.unwrap(); + + // No capacity is left, yet this resolves immediately instead of queueing, because the identity + // conflict is detected before the borrower enters the queue. + assert_eq!(limiter.acquire("a").await.unwrap_err(), AlreadyBorrowed); + assert_eq!(limiter.waiting(), 0); +} + +#[tokio::test] +async fn queued_borrower_proceeds_after_release() { + let limiter = Arc::new(CapacityLimiter::new(1)); + let held = limiter.try_acquire(1u32).unwrap(); + + let waiter = { + let limiter = limiter.clone(); + tokio::spawn(async move { + let permit = limiter.acquire(2u32).await.unwrap(); + *permit.borrower() + }) + }; + + // Let the spawned task reach the queue before capacity is returned. + while limiter.waiting() == 0 { + tokio::task::yield_now().await; + } + assert_eq!(limiter.borrowed(), 1); + + drop(held); + assert_eq!(waiter.await.unwrap(), 2); + assert_eq!(limiter.borrowed(), 0); +} + +#[test] +fn cancelled_acquire_releases_the_identity() { + let limiter = CapacityLimiter::new(0); + + { + let mut acquire = pin!(limiter.acquire("a")); + assert!(poll_once(acquire.as_mut()).is_pending()); + assert_eq!(limiter.waiting(), 1); + + // A second acquisition by the same borrower is rejected while the first one is queued. + assert_eq!( + limiter.try_acquire("a").unwrap_err(), + TryAcquireError::AlreadyBorrowed + ); + } + + assert_eq!(limiter.waiting(), 0); + + // After cancellation the borrower is free to try again. + limiter.set_total(1); + assert!(limiter.try_acquire("a").is_ok()); +} + +#[test] +fn cancelled_acquire_returns_capacity_to_the_next_borrower() { + let limiter = CapacityLimiter::new(1); + let held = limiter.try_acquire("a").unwrap(); + + { + let mut queued = pin!(limiter.acquire("b")); + assert!(poll_once(queued.as_mut()).is_pending()); + + // "b" is handed the released token while it is still suspended, and is then cancelled. + drop(held); + } + + // The token must return to the limiter rather than vanish with the cancelled future. + assert_eq!(limiter.available(), 1); + assert_eq!(limiter.borrowed(), 0); + assert_eq!(limiter.waiting(), 0); + assert!(limiter.try_acquire("c").is_ok()); +} + +#[test] +fn growing_the_total_admits_more_borrowers() { + let limiter = CapacityLimiter::new(1); + let _a = limiter.try_acquire("a").unwrap(); + assert_eq!( + limiter.try_acquire("b").unwrap_err(), + TryAcquireError::NoCapacity + ); + + limiter.set_total(3); + assert_eq!(limiter.total(), 3); + + let _b = limiter.try_acquire("b").unwrap(); + let _c = limiter.try_acquire("c").unwrap(); + assert_eq!(limiter.borrowed(), 3); + assert_eq!(limiter.available(), 0); +} + +#[test] +fn growing_the_total_wakes_a_queued_borrower() { + let limiter = CapacityLimiter::new(0); + let mut queued = pin!(limiter.acquire("a")); + assert!(poll_once(queued.as_mut()).is_pending()); + + limiter.set_total(1); + + let permit = match poll_once(queued.as_mut()) { + Poll::Ready(result) => result.expect("the queued borrower holds no other token"), + Poll::Pending => panic!("raising the total must admit the queued borrower"), + }; + assert_eq!(permit.borrower(), &"a"); + assert_eq!(limiter.borrowed(), 1); + assert_eq!(limiter.waiting(), 0); +} + +#[test] +fn shrinking_the_total_never_revokes_borrowed_tokens() { + let limiter = CapacityLimiter::new(3); + let a = limiter.try_acquire("a").unwrap(); + let b = limiter.try_acquire("b").unwrap(); + let c = limiter.try_acquire("c").unwrap(); + + limiter.set_total(1); + assert_eq!(limiter.total(), 1); + assert_eq!(limiter.borrowed(), 3); + assert_eq!(limiter.available(), 0); + + // The first two releases repay the deficit instead of freeing capacity. + drop(a); + assert_eq!(limiter.available(), 0); + drop(b); + assert_eq!(limiter.available(), 0); + + drop(c); + assert_eq!(limiter.available(), 1); + assert_eq!(limiter.borrowed(), 0); +} + +#[test] +fn shrinking_to_zero_blocks_further_admission() { + let limiter = CapacityLimiter::new(1); + limiter.set_total(0); + + assert_eq!(limiter.total(), 0); + assert_eq!( + limiter.try_acquire("a").unwrap_err(), + TryAcquireError::NoCapacity + ); +} + +#[tokio::test] +async fn queued_borrowers_are_served_in_order() { + let limiter = Arc::new(CapacityLimiter::new(1)); + let held = limiter.try_acquire(0usize).unwrap(); + + let order = Arc::new(Mutex::new(Vec::new())); + let mut handles = Vec::new(); + for borrower in 1..=4usize { + let task_limiter = limiter.clone(); + let order = order.clone(); + + handles.push(tokio::spawn(async move { + let permit = task_limiter.acquire(borrower).await.unwrap(); + order.lock().unwrap().push(*permit.borrower()); + })); + + // Queue the borrowers one at a time so the expected order is well defined. + while limiter.waiting() < borrower { + tokio::task::yield_now().await; + } + } + + drop(held); + for handle in handles { + handle.await.unwrap(); + } + + assert_eq!(*order.lock().unwrap(), vec![1, 2, 3, 4]); + assert_eq!(limiter.waiting(), 0); + assert_eq!(limiter.borrowed(), 0); +} + +#[test] +fn borrower_may_be_any_hashable_identity() { + let limiter = CapacityLimiter::new(2); + + let name: Arc = Arc::from("tenant-a"); + let first = limiter.try_acquire(name.clone()).unwrap(); + assert_eq!(first.borrower().as_ref(), "tenant-a"); + + // A distinct allocation with equal contents is the same borrower. + let same: Arc = Arc::from("tenant-a"); + assert_eq!( + limiter.try_acquire(same).unwrap_err(), + TryAcquireError::AlreadyBorrowed + ); + + let other: Arc = Arc::from("tenant-b"); + assert!(limiter.try_acquire(other).is_ok()); +} + +#[test] +fn errors_describe_their_cause() { + assert_eq!( + AlreadyBorrowed.to_string(), + "borrower already holds a token from this limiter" + ); + assert_eq!( + TryAcquireError::NoCapacity.to_string(), + "no capacity available" + ); + assert_eq!( + TryAcquireError::AlreadyBorrowed.to_string(), + "borrower already holds a token from this limiter" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn queued_contention_preserves_the_capacity_bound() { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + let limiter = Arc::new(CapacityLimiter::new(2)); + let active = Arc::new(AtomicUsize::new(0)); + let mut tasks = tokio::task::JoinSet::new(); + for borrower in 0..8 { + let limiter = limiter.clone(); + let active = active.clone(); + tasks.spawn(async move { + for _ in 0..100 { + let permit = limiter.acquire(borrower).await.unwrap(); + assert!(active.fetch_add(1, Ordering::SeqCst) < 2); + tokio::task::yield_now().await; + active.fetch_sub(1, Ordering::SeqCst); + drop(permit); + } + }); + } + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while let Some(result) = tasks.join_next().await { + result.unwrap(); + } + }) + .await + .expect("queued work must finish"); + assert_eq!(limiter.available(), 2); + assert_eq!(limiter.borrowed(), 0); + assert_eq!(limiter.waiting(), 0); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index bf10a0e7..63bc9cc9 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -19,6 +19,7 @@ use std::cell::Cell; use asyncband::barrier::Barrier; use asyncband::broadcast; +use asyncband::capacity_limiter; use asyncband::completion; use asyncband::condvar::Condvar; use asyncband::event::ManualResetEvent; @@ -70,6 +71,11 @@ impl ManageObject for PoolManager { fn public_types_are_send_and_sync() { fn assert_send_and_sync() {} + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::(); + assert_send_and_sync::(); assert_send_and_sync::(); assert_send_and_sync::(); assert_send_and_sync::(); @@ -138,6 +144,11 @@ fn movable_public_types_are_send() { fn public_types_are_unpin() { fn assert_unpin() {} + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::(); + assert_unpin::(); assert_unpin::(); assert_unpin::(); assert_unpin::();