diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b4a7b7..73470d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file. ### New features +* Add `admission::PriorityShare` to combine shared-capacity priority thresholds with fair sharing among owners at the same priority. * Add an opt-in `asyncband::blocking::FutureExt` bridge with `block_on` and `wait_timeout` methods for waiting on runtime-agnostic futures from synchronous code. ### Breaking changes diff --git a/README.md b/README.md index fe8df56..1650de2 100644 --- a/README.md +++ b/README.md @@ -32,25 +32,26 @@ Asyncband is a runtime-agnostic library providing essential synchronization prim The crate enables no primitives by default. Categories describe each primitive's primary purpose and do not add another module level, so public paths remain concise, such as `asyncband::mutex` and `asyncband::once::OnceCell`. -| Category | Primitive | Feature | Purpose | -| ----------------------- | ------------------------------------------------------------------------------------ | -------------- | ----------------------------------------------------------------------- | -| Shared state | [`Mutex`](https://docs.rs/asyncband/*/asyncband/mutex/struct.Mutex.html) | `mutex` | Protect shared data with asynchronous mutual exclusion. | -| | [`RwLock`](https://docs.rs/asyncband/*/asyncband/rwlock/struct.RwLock.html) | `rwlock` | Allow multiple readers or one writer. | -| | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | -| One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | -| | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | -| | [`OnceMap`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceMap.html) | `once-map` | Initialize and store one value per key. | -| Task coordination | [`Barrier`](https://docs.rs/asyncband/*/asyncband/barrier/struct.Barrier.html) | `barrier` | Wait until all participants reach a synchronization point. | -| | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | -| | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | -| | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | -| Channels | [`oneshot::channel`](https://docs.rs/asyncband/*/asyncband/oneshot/fn.channel.html) | `oneshot` | Send one value between two tasks. | -| | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | -| | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | -| | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | -| Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | -| | [`FairShare`](https://docs.rs/asyncband/*/asyncband/admission/struct.FairShare.html) | `admission` | Fairly share bounded concurrency across keys. | -| | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | +| Category | Primitive | Feature | Purpose | +| ----------------------- | -------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------------------------------- | +| Shared state | [`Mutex`](https://docs.rs/asyncband/*/asyncband/mutex/struct.Mutex.html) | `mutex` | Protect shared data with asynchronous mutual exclusion. | +| | [`RwLock`](https://docs.rs/asyncband/*/asyncband/rwlock/struct.RwLock.html) | `rwlock` | Allow multiple readers or one writer. | +| | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | +| One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | +| | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | +| | [`OnceMap`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceMap.html) | `once-map` | Initialize and store one value per key. | +| Task coordination | [`Barrier`](https://docs.rs/asyncband/*/asyncband/barrier/struct.Barrier.html) | `barrier` | Wait until all participants reach a synchronization point. | +| | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | +| | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | +| | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | +| Channels | [`oneshot::channel`](https://docs.rs/asyncband/*/asyncband/oneshot/fn.channel.html) | `oneshot` | Send one value between two tasks. | +| | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | +| | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | +| | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | +| Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | +| | [`FairShare`](https://docs.rs/asyncband/*/asyncband/admission/struct.FairShare.html) | `admission` | Fairly share bounded concurrency across keys. | +| | [`PriorityShare`](https://docs.rs/asyncband/*/asyncband/admission/struct.PriorityShare.html) | `admission` | Reserve shared capacity by priority while preserving per-key fairness. | +| | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | ## Installation diff --git a/asyncband/src/admission/fair_share.rs b/asyncband/src/admission/fair_share.rs index 5458c04..8112af5 100644 --- a/asyncband/src/admission/fair_share.rs +++ b/asyncband/src/admission/fair_share.rs @@ -15,21 +15,14 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; -use std::collections::VecDeque; -use std::future::Future; use std::hash::BuildHasher; use std::hash::Hash; use std::hash::RandomState; -use std::pin::Pin; use std::sync::Arc; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; -use crate::internal::arena::Arena; -use crate::internal::arena::ArenaKey; -use crate::internal::mutex::Mutex; +use super::OwnedPrioritySharePermit; +use super::PriorityShare; +use super::PrioritySharePermit; /// An admission controller that fairly shares a fixed number of permits across keys. /// @@ -45,7 +38,7 @@ where K: Eq + Hash, S: BuildHasher, { - state: Mutex>, + admission: Arc>, } impl FairShare @@ -84,8 +77,11 @@ where /// Panics if `permits` is zero. pub fn with_hasher(permits: usize, hash_builder: S) -> Self { assert!(permits > 0, "FairShare requires at least one permit"); + let admission = PriorityShare::with_hasher([permits], hash_builder) + .pop() + .expect("one priority was configured"); Self { - state: Mutex::new(State::new(permits, hash_builder)), + admission: Arc::new(admission), } } @@ -94,7 +90,7 @@ where /// A permit already assigned to a queued acquisition counts as held by its /// key, even if that acquisition has not yet been polled again. pub fn available_permits(&self) -> usize { - self.state.lock().available_permits + self.admission.available_permits() } /// Returns the number of acquisitions currently waiting for a permit. @@ -102,19 +98,16 @@ where /// An acquisition is no longer counted once it has been assigned a permit, /// even if its future has not yet been polled again. pub fn num_waiters(&self) -> usize { - self.state.lock().num_waiters + self.admission.num_waiters() } /// Attempts to acquire one permit for `key` without waiting. /// /// This method does not bypass queued acquisitions. pub fn try_acquire(&self, key: K) -> Option> { - let key = Arc::new(key); - let admitted = self.state.lock().try_admit(key.clone()); - admitted.then(|| FairSharePermit { - admission: self, - key, - }) + self.admission + .try_acquire(key) + .map(|permit| FairSharePermit { permit }) } /// Acquires one permit for `key`. @@ -125,11 +118,8 @@ where /// a permit has already been assigned, cancellation releases it for another /// queued acquisition. pub async fn acquire(&self, key: K) -> FairSharePermit<'_, K, S> { - let key = Arc::new(key); - Acquire::new(self, key.clone()).await; FairSharePermit { - admission: self, - key, + permit: self.admission.acquire(key).await, } } @@ -138,11 +128,10 @@ where /// The admission controller must be wrapped in an [`Arc`] to call this /// method. pub fn try_acquire_owned(self: Arc, key: K) -> Option> { - let key = Arc::new(key); - let admitted = self.state.lock().try_admit(key.clone()); - admitted.then(|| OwnedFairSharePermit { - admission: self, - key, + let permit = self.admission.clone().try_acquire_owned(key)?; + Some(OwnedFairSharePermit { + permit, + _admission: self, }) } @@ -155,323 +144,10 @@ where /// /// This method has the same cancellation behavior as [`Self::acquire`]. pub async fn acquire_owned(self: Arc, key: K) -> OwnedFairSharePermit { - let key = Arc::new(key); - Acquire::new(&self, key.clone()).await; + let permit = self.admission.clone().acquire_owned(key).await; OwnedFairSharePermit { - admission: self, - key, - } - } - - fn release(&self, key: &K) { - let waker = { - let mut state = self.state.lock(); - state.release(key); - state.admit_one() - }; - if let Some(waker) = waker { - waker.wake(); - } - } -} - -#[derive(Debug)] -struct State -where - K: Eq + Hash, - S: BuildHasher, -{ - total_permits: usize, - available_permits: usize, - num_waiters: usize, - next_sequence: u64, - groups: HashMap, GroupState, S>, - waiters: Arena, -} - -impl State -where - K: Eq + Hash, - S: BuildHasher, -{ - fn new(permits: usize, hash_builder: S) -> Self { - Self { - total_permits: permits, - available_permits: permits, - num_waiters: 0, - next_sequence: 0, - groups: HashMap::with_hasher(hash_builder), - waiters: Arena::new(), - } - } - - fn try_admit(&mut self, key: Arc) -> bool { - if self.available_permits == 0 || self.num_waiters != 0 { - return false; - } - - self.available_permits -= 1; - self.groups.entry(key).or_default().held_permits += 1; - true - } - - fn enqueue(&mut self, key: Arc, waker: &Waker) -> ArenaKey { - let sequence = self.next_sequence; - self.next_sequence += 1; - - let waiter = self.waiters.insert(Waiter { - waker: Some(waker.clone()), - }); - self.groups - .entry(key) - .or_default() - .queue - .push_back(QueuedWaiter { - id: waiter, - sequence, - }); - self.num_waiters += 1; - waiter - } - - fn poll_waiter(&mut self, waiter: ArenaKey, waker: &Waker) -> Poll<()> { - let state = &mut self.waiters[waiter]; - - if state.waker.is_none() { - self.waiters.remove(waiter); - Poll::Ready(()) - } else { - if state - .waker - .as_ref() - .is_none_or(|current| !current.will_wake(waker)) - { - state.waker = Some(waker.clone()); - } - Poll::Pending - } - } - - fn cancel(&mut self, waiter_id: ArenaKey, key: &K) { - let waiter = self.waiters.remove(waiter_id); - if waiter.waker.is_none() { - self.release(key); - return; - } - - let remove_group = { - let group = self - .groups - .get_mut(key) - .expect("FairShare waiter group is missing"); - if group - .queue - .front() - .is_some_and(|waiter| waiter.id == waiter_id) - { - group.queue.pop_front(); - } else { - let position = group - .queue - .iter() - .position(|candidate| candidate.id == waiter_id) - .expect("FairShare waiter is missing from its group"); - group.queue.remove(position); - } - group.held_permits == 0 && group.queue.is_empty() - }; - - self.num_waiters -= 1; - if remove_group { - self.groups.remove(key); - } - } - - fn admit_waiters(&mut self, wakers: &mut Vec) { - while let Some(waker) = self.admit_one() { - wakers.push(waker); - } - } - - fn admit_one(&mut self) -> Option { - if self.available_permits == 0 || self.num_waiters == 0 { - return None; - } - - let key = self - .next_group() - .expect("FairShare has pending acquisitions without a group"); - let waiter = { - let group = self - .groups - .get_mut(&key) - .expect("FairShare pending group is missing"); - let waiter = group - .queue - .pop_front() - .expect("FairShare pending group has no waiters") - .id; - group.held_permits += 1; - waiter - }; - - self.available_permits -= 1; - self.num_waiters -= 1; - - Some( - self.waiters[waiter] - .waker - .take() - .expect("pending FairShare waiter must have a waker"), - ) - } - - fn next_group(&self) -> Option> { - self.groups - .iter() - .filter_map(|(key, group)| { - let waiter = group.queue.front()?; - Some((group.held_permits, waiter.sequence, key)) - }) - .min_by_key(|(held_permits, sequence, _)| (*held_permits, *sequence)) - .map(|(_, _, key)| key.clone()) - } - - fn release(&mut self, key: &K) { - let remove_group = { - let group = self - .groups - .get_mut(key) - .expect("FairShare released a permit for an unknown key"); - debug_assert!(group.held_permits > 0); - group.held_permits -= 1; - group.held_permits == 0 && group.queue.is_empty() - }; - - if remove_group { - self.groups.remove(key); - } - - self.available_permits += 1; - debug_assert!(self.available_permits <= self.total_permits); - } -} - -#[derive(Debug, Default)] -struct GroupState { - held_permits: usize, - queue: VecDeque, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -struct QueuedWaiter { - id: ArenaKey, - sequence: u64, -} - -#[derive(Debug)] -struct Waiter { - /// `None` once this waiter has been admitted but has not yet observed completion. - waker: Option, -} - -#[derive(Debug)] -struct Acquire<'a, K, S> -where - K: Eq + Hash, - S: BuildHasher, -{ - admission: &'a FairShare, - key: Arc, - waiter: Option, - completed: bool, -} - -impl<'a, K, S> Acquire<'a, K, S> -where - K: Eq + Hash, - S: BuildHasher, -{ - fn new(admission: &'a FairShare, key: Arc) -> Self { - Self { - admission, - key, - waiter: None, - completed: false, - } - } -} - -impl Drop for Acquire<'_, K, S> -where - K: Eq + Hash, - S: BuildHasher, -{ - fn drop(&mut self) { - let Some(waiter) = self.waiter.take() else { - return; - }; - - let waker = { - let mut state = self.admission.state.lock(); - state.cancel(waiter, &self.key); - state.admit_one() - }; - if let Some(waker) = waker { - waker.wake(); - } - } -} - -impl Future for Acquire<'_, K, S> -where - K: Eq + Hash, - S: BuildHasher, -{ - type Output = (); - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - if this.completed { - return Poll::Ready(()); - } - - if let Some(waiter) = this.waiter { - if this - .admission - .state - .lock() - .poll_waiter(waiter, cx.waker()) - .is_ready() - { - this.waiter = None; - this.completed = true; - return Poll::Ready(()); - } - return Poll::Pending; - } - - let mut wakers = Vec::new(); - let ready = { - let mut state = this.admission.state.lock(); - if state.try_admit(this.key.clone()) { - this.completed = true; - return Poll::Ready(()); - } - - let waiter = state.enqueue(this.key.clone(), cx.waker()); - this.waiter = Some(waiter); - state.admit_waiters(&mut wakers); - state.poll_waiter(waiter, cx.waker()).is_ready() - }; - wake_all(wakers); - - if ready { - this.waiter = None; - this.completed = true; - Poll::Ready(()) - } else { - Poll::Pending + permit, + _admission: self, } } } @@ -491,8 +167,7 @@ where K: Eq + Hash, S: BuildHasher, { - admission: &'a FairShare, - key: Arc, + permit: PrioritySharePermit<'a, K, S>, } impl FairSharePermit<'_, K, S> @@ -502,17 +177,7 @@ where { /// Returns the key associated with this permit. pub fn key(&self) -> &K { - &self.key - } -} - -impl Drop for FairSharePermit<'_, K, S> -where - K: Eq + Hash, - S: BuildHasher, -{ - fn drop(&mut self) { - self.admission.release(&self.key); + self.permit.key() } } @@ -532,8 +197,8 @@ where K: Eq + Hash, S: BuildHasher, { - admission: Arc>, - key: Arc, + permit: OwnedPrioritySharePermit, + _admission: Arc>, } impl OwnedFairSharePermit @@ -543,22 +208,6 @@ where { /// Returns the key associated with this permit. pub fn key(&self) -> &K { - &self.key - } -} - -impl Drop for OwnedFairSharePermit -where - K: Eq + Hash, - S: BuildHasher, -{ - fn drop(&mut self) { - self.admission.release(&self.key); - } -} - -fn wake_all(wakers: Vec) { - for waker in wakers { - waker.wake(); + self.permit.key() } } diff --git a/asyncband/src/admission/mod.rs b/asyncband/src/admission/mod.rs index ab87e88..2dd59d0 100644 --- a/asyncband/src/admission/mod.rs +++ b/asyncband/src/admission/mod.rs @@ -18,15 +18,26 @@ //! Admission control policies for bounded asynchronous work. //! //! This module provides [`FairShare`], a work-conserving admission policy for -//! workloads partitioned by key. It maintains a fixed number of permits and, -//! when contended, admits work for the key with the fewest permits currently -//! held. Ties are resolved by queue order. +//! workloads partitioned by key, and [`PriorityShare`], which adds strict +//! priorities and reserved headroom within one shared capacity. Both policies +//! admit work for the key with the fewest permits currently held. For +//! [`PriorityShare`], that count spans all priorities and is considered after +//! selecting the highest eligible priority. Ties are resolved by queue order. //! -//! Fairness applies to the number of permits held by contending keys. It does -//! not reserve permits for idle keys or account for differences in execution -//! time or work cost. +//! Fairness applies to the number of permits held by contending keys. Neither +//! policy accounts for differences in execution time or work cost. +//! [`FairShare`] does not reserve permits for idle keys. [`PriorityShare`] uses +//! admission thresholds to reserve headroom for higher priorities, so capacity +//! can remain unused while lower-priority work waits. Its constructor returns +//! one priority-bound handle per configured threshold; all of those handles +//! share the same scheduler. mod fair_share; +mod priority_share; + pub use fair_share::FairShare; pub use fair_share::FairSharePermit; pub use fair_share::OwnedFairSharePermit; +pub use priority_share::OwnedPrioritySharePermit; +pub use priority_share::PriorityShare; +pub use priority_share::PrioritySharePermit; diff --git a/asyncband/src/admission/priority_share/mod.rs b/asyncband/src/admission/priority_share/mod.rs new file mode 100644 index 0000000..acf5d65 --- /dev/null +++ b/asyncband/src/admission/priority_share/mod.rs @@ -0,0 +1,344 @@ +// 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. + +mod scheduler; + +use std::hash::BuildHasher; +use std::hash::Hash; +use std::hash::RandomState; +use std::sync::Arc; + +use scheduler::Scheduler; + +/// A priority-bound handle in a priority-share admission family. +/// +/// [`PriorityShare::new`] creates one handle per configured priority, ordered +/// from lowest to highest. Those handles share one capacity and one scheduler. +/// A handle keeps its priority when it is moved or cloned, so acquisition +/// methods do not take a priority argument. Separate constructor calls create +/// independent admission families. +/// +/// Configuration entries add capacity from the lowest to the highest priority. +/// The admission limit at priority `p` is the sum of entries up to and +/// including `p`, and every assigned permit counts toward that limit regardless +/// of its priority. For example, `[4, 1]` lets priority 0 enter while fewer than +/// four permits are assigned and priority 1 while fewer than five are assigned. +/// A higher priority can therefore use all shared capacity, including headroom +/// unavailable to lower priorities. +/// +/// When contended, the highest eligible priority is admitted first. Within +/// that priority, the owner with the fewest permits currently held across all +/// priorities is admitted first. Ties are resolved by queue order. The same +/// owner may acquire and wait at multiple priorities; all of its permits count +/// toward one fair-share identity. +/// +/// Priority only affects admission. An acquisition that has already been +/// assigned a permit is never revoked for later higher-priority work. Sustained +/// higher-priority demand can therefore starve lower priorities, and reserved +/// headroom can remain unused while lower-priority work waits. Priorities with +/// equal admission limits still differ under contention because the higher one +/// is admitted first. +/// +/// Available-permit observations are scoped to the handle's priority. Values +/// from different handles overlap and must not be added together. +#[derive(Debug)] +pub struct PriorityShare +where + K: Eq + Hash, + S: BuildHasher, +{ + scheduler: Arc>, + priority: usize, +} + +impl Clone for PriorityShare +where + K: Eq + Hash, + S: BuildHasher, +{ + fn clone(&self) -> Self { + Self { + scheduler: self.scheduler.clone(), + priority: self.priority, + } + } +} + +impl PriorityShare +where + K: Eq + Hash, +{ + /// Creates a priority-share admission family. + /// + /// `capacity_increments` lists the additional shared capacity unlocked at + /// each priority from lowest to highest. The returned vector contains one + /// priority-bound handle for every entry in the same order. Individual + /// entries may be zero as long as the total capacity is nonzero. + /// + /// # Panics + /// + /// Panics if `capacity_increments` is empty, its total is zero, or its total + /// overflows `usize`. + /// + /// # Examples + /// + /// ``` + /// use asyncband::admission::PriorityShare; + /// + /// let priorities = PriorityShare::::new([2, 1]); + /// let low = &priorities[0]; + /// let high = &priorities[1]; + /// + /// assert_eq!(low.priority(), 0); + /// assert_eq!(high.priority(), 1); + /// assert_eq!(low.available_permits(), 2); + /// assert_eq!(high.available_permits(), 3); + /// assert_eq!(low.try_acquire("low".to_owned()).unwrap().priority(), 0); + /// assert_eq!(high.try_acquire("high".to_owned()).unwrap().priority(), 1); + /// ``` + pub fn new(capacity_increments: C) -> Vec + where + C: AsRef<[usize]>, + { + Self::with_hasher(capacity_increments, RandomState::new()) + } +} + +impl PriorityShare +where + K: Eq + Hash, + S: BuildHasher, +{ + /// Creates a priority-share admission family with the given hash builder. + /// + /// `capacity_increments` has the same meaning as in [`Self::new`]. + /// + /// # Panics + /// + /// Panics under the same conditions as [`Self::new`]. + pub fn with_hasher(capacity_increments: C, hash_builder: S) -> Vec + where + C: AsRef<[usize]>, + { + let capacity_increments = capacity_increments.as_ref(); + assert!( + !capacity_increments.is_empty(), + "PriorityShare requires at least one priority" + ); + + let mut total_permits = 0usize; + let mut admission_limits = Vec::with_capacity(capacity_increments.len()); + for increment in capacity_increments { + total_permits = total_permits + .checked_add(*increment) + .expect("PriorityShare capacity overflow"); + admission_limits.push(total_permits); + } + assert!( + total_permits > 0, + "PriorityShare requires at least one permit" + ); + + let scheduler = Arc::new(Scheduler::new( + admission_limits.into_boxed_slice(), + hash_builder, + )); + (0..capacity_increments.len()) + .map(|priority| Self { + scheduler: scheduler.clone(), + priority, + }) + .collect() + } + + /// Returns the priority bound to this handle. + /// + /// Priorities are dense zero-based values, and a larger value means a + /// higher priority. + pub fn priority(&self) -> usize { + self.priority + } + + /// Returns the number of permits currently available to this priority. + /// + /// This is the number of additional acquisitions that fit below this + /// priority's shared admission limit. A permit assigned at any priority can + /// reduce this value. Values from different handles overlap and must not be + /// added together. The highest-priority handle reports the total number of + /// unassigned permits in the admission family. + /// + /// A permit assigned to a queued acquisition is no longer available even + /// if that acquisition has not been polled again. This method returns an + /// instantaneous observation; use [`Self::try_acquire`] to atomically test + /// and acquire capacity. + pub fn available_permits(&self) -> usize { + self.scheduler.available_permits(self.priority) + } + + pub(super) fn num_waiters(&self) -> usize { + self.scheduler.num_waiters() + } + + /// Attempts to acquire one permit for `owner` at this handle's priority. + /// + /// This method may bypass queued work at lower priorities. It does not + /// bypass queued work that should be admitted first at the same or a higher + /// priority. + pub fn try_acquire(&self, owner: K) -> Option> { + let owner = Arc::new(owner); + self.scheduler + .try_acquire(owner.clone(), self.priority) + .then(|| PrioritySharePermit { + admission: self, + owner, + }) + } + + /// Acquires one permit for `owner` at this handle's priority. + /// + /// # Cancel safety + /// + /// Cancelling this method loses the acquisition's place in the queue. If a + /// permit has already been assigned, cancellation releases it for another + /// queued acquisition. + pub async fn acquire(&self, owner: K) -> PrioritySharePermit<'_, K, S> { + let owner = Arc::new(owner); + self.scheduler.acquire(owner.clone(), self.priority).await; + PrioritySharePermit { + admission: self, + owner, + } + } + + /// Attempts to acquire one owned permit for `owner` at this handle's + /// priority. + /// + /// The handle must be wrapped in an [`Arc`] to call this method. + pub fn try_acquire_owned(self: Arc, owner: K) -> Option> { + let owner = Arc::new(owner); + if !self.scheduler.try_acquire(owner.clone(), self.priority) { + return None; + } + Some(OwnedPrioritySharePermit { + admission: self, + owner, + }) + } + + /// Acquires one owned permit for `owner` at this handle's priority. + /// + /// The handle must be wrapped in an [`Arc`] to call this method. + /// + /// # Cancel safety + /// + /// This method has the same cancellation behavior as [`Self::acquire`]. + pub async fn acquire_owned(self: Arc, owner: K) -> OwnedPrioritySharePermit { + let owner = Arc::new(owner); + self.scheduler.acquire(owner.clone(), self.priority).await; + OwnedPrioritySharePermit { + admission: self, + owner, + } + } + + fn release(&self, owner: &K) { + self.scheduler.release(owner); + } +} + +/// A borrowed permit from a [`PriorityShare`] handle. +/// +/// Dropping this permit returns it to the shared admission family and may admit +/// a queued acquisition. +#[must_use = "permits are released immediately when dropped"] +#[derive(Debug)] +pub struct PrioritySharePermit<'a, K, S = RandomState> +where + K: Eq + Hash, + S: BuildHasher, +{ + admission: &'a PriorityShare, + owner: Arc, +} + +impl PrioritySharePermit<'_, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + /// Returns the owner associated with this permit. + pub fn key(&self) -> &K { + &self.owner + } + + /// Returns the priority associated with this permit. + pub fn priority(&self) -> usize { + self.admission.priority + } +} + +impl Drop for PrioritySharePermit<'_, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn drop(&mut self) { + self.admission.release(&self.owner); + } +} + +/// An owned permit from a [`PriorityShare`] handle. +/// +/// Unlike [`PrioritySharePermit`], this type owns an [`Arc`] to its +/// priority-bound handle and has no lifetime parameter. Dropping it returns the +/// permit and may admit a queued acquisition. +#[must_use = "permits are released immediately when dropped"] +#[derive(Debug)] +pub struct OwnedPrioritySharePermit +where + K: Eq + Hash, + S: BuildHasher, +{ + admission: Arc>, + owner: Arc, +} + +impl OwnedPrioritySharePermit +where + K: Eq + Hash, + S: BuildHasher, +{ + /// Returns the owner associated with this permit. + pub fn key(&self) -> &K { + &self.owner + } + + /// Returns the priority associated with this permit. + pub fn priority(&self) -> usize { + self.admission.priority + } +} + +impl Drop for OwnedPrioritySharePermit +where + K: Eq + Hash, + S: BuildHasher, +{ + fn drop(&mut self) { + self.admission.release(&self.owner); + } +} diff --git a/asyncband/src/admission/priority_share/scheduler.rs b/asyncband/src/admission/priority_share/scheduler.rs new file mode 100644 index 0000000..fd14632 --- /dev/null +++ b/asyncband/src/admission/priority_share/scheduler.rs @@ -0,0 +1,432 @@ +// 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::cmp::Reverse; +use std::collections::HashMap; +use std::collections::VecDeque; +use std::future::Future; +use std::hash::BuildHasher; +use std::hash::Hash; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use crate::internal::arena::Arena; +use crate::internal::arena::ArenaKey; +use crate::internal::mutex::Mutex; + +#[derive(Debug)] +pub(super) struct Scheduler +where + K: Eq + Hash, + S: BuildHasher, +{ + state: Mutex>, +} + +impl Scheduler +where + K: Eq + Hash, + S: BuildHasher, +{ + pub(super) fn new(admission_limits: Box<[usize]>, hash_builder: S) -> Self { + Self { + state: Mutex::new(State::new(admission_limits, hash_builder)), + } + } + + pub(super) fn available_permits(&self, priority: usize) -> usize { + self.state.lock().available_for(priority) + } + + pub(super) fn num_waiters(&self) -> usize { + self.state.lock().total_num_waiters + } + + pub(super) fn try_acquire(&self, owner: Arc, priority: usize) -> bool { + self.state.lock().try_admit(owner, priority) + } + + pub(super) fn acquire(&self, owner: Arc, priority: usize) -> Acquire<'_, K, S> { + Acquire::new(self, owner, priority) + } + + pub(super) fn release(&self, owner: &K) { + let mut wakers = Vec::new(); + { + let mut state = self.state.lock(); + state.release(owner); + state.admit_waiters(&mut wakers); + } + wake_all(wakers); + } +} + +#[derive(Debug)] +struct State +where + K: Eq + Hash, + S: BuildHasher, +{ + // Entry `p` is the shared-capacity admission limit for priority `p`. + admission_limits: Box<[usize]>, + total_available_permits: usize, + total_num_waiters: usize, + next_sequence: u64, + owners: HashMap, OwnerState, S>, + waiters: Arena, +} + +impl State +where + K: Eq + Hash, + S: BuildHasher, +{ + fn new(admission_limits: Box<[usize]>, hash_builder: S) -> Self { + let total_permits = *admission_limits + .last() + .expect("priority-share requires at least one admission limit"); + debug_assert!(total_permits > 0); + Self { + admission_limits, + total_available_permits: total_permits, + total_num_waiters: 0, + next_sequence: 0, + owners: HashMap::with_hasher(hash_builder), + waiters: Arena::new(), + } + } + + fn try_admit(&mut self, owner: Arc, priority: usize) -> bool { + if !self.can_admit(priority) { + return false; + } + + self.admit(owner); + true + } + + fn enqueue(&mut self, owner: Arc, priority: usize, waker: &Waker) -> ArenaKey { + let sequence = self.next_sequence; + self.next_sequence += 1; + + let waiter = self.waiters.insert(Waiter { + sequence, + waker: Some(waker.clone()), + admitted: false, + }); + let owner = self.owners.entry(owner).or_insert_with(OwnerState::new); + if let Some(queue) = owner + .queues + .iter_mut() + .find(|queue| queue.priority == priority) + { + queue.waiters.push_back(waiter); + } else { + owner.queues.push(PriorityQueue::new(priority, waiter)); + } + self.total_num_waiters += 1; + waiter + } + + fn poll_waiter(&mut self, waiter: ArenaKey, waker: &Waker) -> Poll<()> { + let state = self + .waiters + .get_mut(waiter) + .expect("priority-share waiter is missing"); + + if state.admitted { + self.waiters.remove(waiter); + Poll::Ready(()) + } else { + if state + .waker + .as_ref() + .is_none_or(|current| !current.will_wake(waker)) + { + state.waker = Some(waker.clone()); + } + Poll::Pending + } + } + + fn cancel(&mut self, waiter_id: ArenaKey, owner: &K, priority: usize) { + let waiter = self.waiters.remove(waiter_id); + if waiter.admitted { + self.release(owner); + return; + } + + let remove_owner = { + let owner = self + .owners + .get_mut(owner) + .expect("priority-share waiter owner is missing"); + let queue = owner + .queues + .iter() + .position(|queue| queue.priority == priority) + .expect("priority-share waiter queue is missing"); + let waiter = owner.queues[queue] + .waiters + .iter() + .position(|candidate| *candidate == waiter_id) + .expect("priority-share waiter is missing from its queue"); + owner.queues[queue].waiters.remove(waiter); + if owner.queues[queue].waiters.is_empty() { + owner.queues.swap_remove(queue); + } + owner.held_permits == 0 && owner.queues.is_empty() + }; + + self.total_num_waiters -= 1; + if remove_owner { + self.owners.remove(owner); + } + } + + fn admit_waiters(&mut self, wakers: &mut Vec) { + while self.total_available_permits > 0 && self.total_num_waiters > 0 { + let Some((owner, priority)) = self.next_owner() else { + return; + }; + let owner_state = self + .owners + .get_mut(&owner) + .expect("pending priority-share owner is missing"); + let queue = owner_state + .queues + .iter() + .position(|queue| queue.priority == priority) + .expect("pending priority-share queue is missing"); + let waiter = owner_state.queues[queue] + .waiters + .pop_front() + .expect("pending priority-share queue has no waiters"); + if owner_state.queues[queue].waiters.is_empty() { + owner_state.queues.swap_remove(queue); + } + owner_state.held_permits += 1; + self.total_available_permits -= 1; + self.total_num_waiters -= 1; + + let waiter = &mut self.waiters[waiter]; + waiter.admitted = true; + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + } + + fn next_owner(&self) -> Option<(Arc, usize)> { + let mut next = None; + for (owner, state) in &self.owners { + for queue in &state.queues { + if !self.can_admit(queue.priority) { + continue; + } + let waiter = queue + .waiters + .front() + .expect("pending priority-share queue has no waiters"); + let order = ( + Reverse(queue.priority), + state.held_permits, + self.waiters[*waiter].sequence, + ); + if next.as_ref().is_none_or(|(best, _, _)| order < *best) { + next = Some((order, owner.clone(), queue.priority)); + } + } + } + next.map(|(_, owner, priority)| (owner, priority)) + } + + fn can_admit(&self, priority: usize) -> bool { + self.available_for(priority) > 0 + } + + fn available_for(&self, priority: usize) -> usize { + debug_assert!(priority < self.admission_limits.len()); + let total_permits = self.admission_limits[self.admission_limits.len() - 1]; + let held_permits = total_permits - self.total_available_permits; + self.admission_limits[priority].saturating_sub(held_permits) + } + + fn admit(&mut self, owner: Arc) { + self.owners + .entry(owner) + .or_insert_with(OwnerState::new) + .held_permits += 1; + self.total_available_permits -= 1; + } + + fn release(&mut self, owner: &K) { + let remove_owner = { + let owner = self + .owners + .get_mut(owner) + .expect("priority-share released a permit for an unknown owner"); + debug_assert!(owner.held_permits > 0); + owner.held_permits -= 1; + owner.held_permits == 0 && owner.queues.is_empty() + }; + + if remove_owner { + self.owners.remove(owner); + } + + self.total_available_permits += 1; + debug_assert!( + self.total_available_permits <= self.admission_limits[self.admission_limits.len() - 1] + ); + } +} + +#[derive(Debug)] +struct OwnerState { + held_permits: usize, + queues: Vec, +} + +impl OwnerState { + fn new() -> Self { + Self { + held_permits: 0, + queues: Vec::new(), + } + } +} + +#[derive(Debug)] +struct PriorityQueue { + priority: usize, + waiters: VecDeque, +} + +impl PriorityQueue { + fn new(priority: usize, waiter: ArenaKey) -> Self { + Self { + priority, + waiters: VecDeque::from([waiter]), + } + } +} + +#[derive(Debug)] +struct Waiter { + sequence: u64, + waker: Option, + admitted: bool, +} + +#[derive(Debug)] +pub(super) struct Acquire<'a, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + scheduler: &'a Scheduler, + owner: Arc, + priority: usize, + waiter: Option, + completed: bool, +} + +impl<'a, K, S> Acquire<'a, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn new(scheduler: &'a Scheduler, owner: Arc, priority: usize) -> Self { + Self { + scheduler, + owner, + priority, + waiter: None, + completed: false, + } + } +} + +impl Drop for Acquire<'_, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + fn drop(&mut self) { + let Some(waiter) = self.waiter.take() else { + return; + }; + + let mut wakers = Vec::new(); + { + let mut state = self.scheduler.state.lock(); + state.cancel(waiter, &self.owner, self.priority); + state.admit_waiters(&mut wakers); + } + wake_all(wakers); + } +} + +impl Future for Acquire<'_, K, S> +where + K: Eq + Hash, + S: BuildHasher, +{ + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + if this.completed { + return Poll::Ready(()); + } + + if let Some(waiter) = this.waiter { + if this + .scheduler + .state + .lock() + .poll_waiter(waiter, cx.waker()) + .is_ready() + { + this.waiter = None; + this.completed = true; + return Poll::Ready(()); + } + return Poll::Pending; + } + + let mut state = this.scheduler.state.lock(); + if state.try_admit(this.owner.clone(), this.priority) { + this.completed = true; + return Poll::Ready(()); + } + + let waiter = state.enqueue(this.owner.clone(), this.priority, cx.waker()); + this.waiter = Some(waiter); + Poll::Pending + } +} + +fn wake_all(wakers: Vec) { + for waker in wakers { + waker.wake(); + } +} diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 505f400..7ea1ef2 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -51,14 +51,14 @@ //! //! # API guide //! -//! | Use case | APIs | Cargo features | -//! | -------------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------- | -//! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | -//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | -//! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | -//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | -//! | Control workloads | [`semaphore::Semaphore`], [`admission::FairShare`], [`singleflight::Group`] | `semaphore`, `admission`, `singleflight` | -//! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | +//! | Use case | APIs | Cargo features | +//! | -------------------------- | --------------------------------------------------------------------------------------------------------- | ------------------------------------------- | +//! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | +//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | +//! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | +//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | +//! | Control workloads | [`semaphore::Semaphore`], [`admission::FairShare`], [`admission::PriorityShare`], [`singleflight::Group`] | `semaphore`, `admission`, `singleflight` | +//! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | //! //! # Runtime and blocking model //! diff --git a/tests-integration/tests/admission_test.rs b/tests-integration/tests/admission_test.rs index eaebbc5..1f98384 100644 --- a/tests-integration/tests/admission_test.rs +++ b/tests-integration/tests/admission_test.rs @@ -50,6 +50,22 @@ fn tracks_available_permits() { assert_eq!(admission.available_permits(), 2); } +#[test] +fn tracks_waiters() { + let admission = FairShare::new(1); + let held = admission.try_acquire("held").unwrap(); + + { + let acquire = admission.acquire("waiting"); + let mut acquire = pin!(acquire); + assert!(poll_once(acquire.as_mut()).is_pending()); + assert_eq!(admission.num_waiters(), 1); + } + + assert_eq!(admission.num_waiters(), 0); + drop(held); +} + #[test] fn uses_all_permits_without_reservations() { let admission = FairShare::new(3); diff --git a/tests-integration/tests/priority_share_test.rs b/tests-integration/tests/priority_share_test.rs new file mode 100644 index 0000000..91e7820 --- /dev/null +++ b/tests-integration/tests/priority_share_test.rs @@ -0,0 +1,532 @@ +// 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::collections::hash_map::DefaultHasher; +use std::hash::BuildHasherDefault; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use asyncband::admission::FairShare; +use asyncband::admission::PriorityShare; +use tests_integration::poll_once; + +#[test] +#[should_panic(expected = "PriorityShare requires at least one priority")] +fn empty_capacity_increments_panic() { + PriorityShare::::new([]); +} + +#[test] +#[should_panic(expected = "PriorityShare requires at least one permit")] +fn zero_total_capacity_panics() { + PriorityShare::::new([0, 0]); +} + +#[test] +#[should_panic(expected = "PriorityShare capacity overflow")] +fn capacity_overflow_panics() { + PriorityShare::::new([usize::MAX, 1]); +} + +#[test] +fn returns_priority_bound_handles_that_share_state() { + let priorities = PriorityShare::new([1, 0, 2]); + + assert_eq!(priorities.len(), 3); + for (priority, admission) in priorities.iter().enumerate() { + assert_eq!(admission.priority(), priority); + } + + let high = priorities[2].clone(); + let permits = ["high-0", "high-1", "high-2"].map(|owner| high.try_acquire(owner).unwrap()); + assert!(priorities[0].try_acquire("low").is_none()); + assert!(priorities[1].try_acquire("middle").is_none()); + assert!(high.try_acquire("high-3").is_none()); + + drop(permits); + let low_permit = priorities[0].try_acquire("low").unwrap(); + assert_eq!(low_permit.priority(), 0); + assert!(priorities[1].try_acquire("middle").is_none()); + assert!(high.try_acquire("high").is_some()); +} + +#[test] +fn available_permits_are_scoped_to_each_priority() { + let priorities = PriorityShare::new([4, 1]); + let low = &priorities[0]; + let high = &priorities[1]; + + assert_eq!(low.available_permits(), 4); + assert_eq!(high.available_permits(), 5); + let permits = ["high-0", "high-1", "high-2"].map(|owner| high.try_acquire(owner).unwrap()); + assert_eq!(low.available_permits(), 1); + assert_eq!(high.available_permits(), 2); + + drop(permits); + assert_eq!(low.available_permits(), 4); + assert_eq!(high.available_permits(), 5); +} + +#[test] +fn all_priorities_count_toward_shared_admission_thresholds() { + let priorities = PriorityShare::new([4, 1]); + let low = &priorities[0]; + let high = &priorities[1]; + let low0 = low.try_acquire("low-0").unwrap(); + let low1 = low.try_acquire("low-1").unwrap(); + let low2 = low.try_acquire("low-2").unwrap(); + let high0 = high.try_acquire("high-0").unwrap(); + + assert!(low.try_acquire("low-3").is_none()); + + let high1 = high.try_acquire("high-1").unwrap(); + assert_eq!(high1.priority(), 1); + assert!(high.try_acquire("high-2").is_none()); + + drop((low0, low1, low2, high0, high1)); + assert!(low.try_acquire("low-after-release").is_some()); +} + +#[test] +fn higher_priority_can_use_entire_shared_capacity() { + let priorities = PriorityShare::new([2, 1]); + let high = &priorities[1]; + let permits = [ + high.try_acquire("high-0").unwrap(), + high.try_acquire("high-1").unwrap(), + high.try_acquire("high-2").unwrap(), + ]; + + assert!(high.try_acquire("high-3").is_none()); + + drop(permits); + assert!(high.try_acquire("high-after-release").is_some()); +} + +#[test] +fn higher_priority_bypasses_queued_lower_priority() { + let priorities = PriorityShare::new([1, 1]); + let low = &priorities[0]; + let high = &priorities[1]; + let low_held = low.try_acquire("low-held").unwrap(); + + let low_waiter = low.acquire("low-waiter"); + let mut low_waiter = pin!(low_waiter); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + + let high_permit = high.try_acquire("high").unwrap(); + assert_eq!(high_permit.key(), &"high"); + assert_eq!(high_permit.priority(), 1); + + drop(low_held); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + + drop(high_permit); + let low_permit = match poll_once(low_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => { + panic!("the lower-priority waiter should enter below its shared threshold") + } + }; + + drop(low_permit); +} + +#[test] +fn released_capacity_goes_to_the_highest_priority_waiter() { + let priorities = PriorityShare::new([1, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held = low.try_acquire("held").unwrap(); + + let low_waiter = low.acquire("low"); + let mut low_waiter = pin!(low_waiter); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + + let high_waiter = high.acquire("high"); + let mut high_waiter = pin!(high_waiter); + assert!(poll_once(high_waiter.as_mut()).is_pending()); + + drop(held); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + let high_permit = match poll_once(high_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("the higher-priority waiter should be admitted first"), + }; + + drop(high_permit); + let low_permit = match poll_once(low_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("the lower-priority waiter should eventually be admitted"), + }; + drop(low_permit); +} + +#[test] +fn assigned_permit_is_not_revoked_for_a_higher_priority_waiter() { + let priorities = PriorityShare::new([1, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held = low.try_acquire("held").unwrap(); + + let low_waiter = low.acquire("low"); + let mut low_waiter = pin!(low_waiter); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + drop(held); + + let high_waiter = high.acquire("high"); + let mut high_waiter = pin!(high_waiter); + assert!(poll_once(high_waiter.as_mut()).is_pending()); + + let low_permit = match poll_once(low_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("an assigned lower-priority permit must not be revoked"), + }; + assert!(poll_once(high_waiter.as_mut()).is_pending()); + + drop(low_permit); + let high_permit = match poll_once(high_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("the higher-priority waiter should receive the released permit"), + }; + drop(high_permit); +} + +#[test] +fn cancelling_an_assigned_permit_reassigns_it_by_priority() { + let priorities = PriorityShare::new([1, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held = low.try_acquire("held").unwrap(); + + let mut low_waiter = Box::pin(low.acquire("low")); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + drop(held); + + let high_waiter = high.acquire("high"); + let mut high_waiter = pin!(high_waiter); + assert!(poll_once(high_waiter.as_mut()).is_pending()); + + drop(low_waiter); + let high_permit = match poll_once(high_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("cancellation should reassign the permit by priority"), + }; + drop(high_permit); +} + +#[test] +fn fair_shares_within_one_priority() { + let priorities = PriorityShare::new([0, 3]); + let high = &priorities[1]; + let mut held_by_a = vec![ + high.try_acquire("a").unwrap(), + high.try_acquire("a").unwrap(), + high.try_acquire("a").unwrap(), + ]; + + let acquire_a = high.acquire("a"); + let mut acquire_a = pin!(acquire_a); + assert!(poll_once(acquire_a.as_mut()).is_pending()); + + let acquire_b = high.acquire("b"); + let mut acquire_b = pin!(acquire_b); + assert!(poll_once(acquire_b.as_mut()).is_pending()); + + let acquire_c = high.acquire("c"); + let mut acquire_c = pin!(acquire_c); + assert!(poll_once(acquire_c.as_mut()).is_pending()); + + drop(held_by_a.pop().unwrap()); + let permit_b = match poll_once(acquire_b.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("owner b should receive the first released permit"), + }; + + drop(held_by_a.pop().unwrap()); + let permit_c = match poll_once(acquire_c.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("owner c should receive the second released permit"), + }; + + drop(held_by_a.pop().unwrap()); + let permit_a = match poll_once(acquire_a.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("owner a should receive the third released permit"), + }; + drop((permit_a, permit_b, permit_c)); +} + +#[test] +fn same_owner_can_acquire_multiple_priorities() { + let priorities = PriorityShare::new([1, 1]); + let low = &priorities[0]; + let high = &priorities[1]; + + let low_permit = low.try_acquire("owner").unwrap(); + let high_permit = high.try_acquire("owner").unwrap(); + + assert_eq!(low_permit.priority(), 0); + assert_eq!(high_permit.priority(), 1); + + drop((low_permit, high_permit)); + assert!(low.try_acquire("owner-after-release").is_some()); +} + +#[test] +fn fairness_counts_an_owners_permits_across_priorities() { + let priorities = PriorityShare::new([2, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held_a0 = low.try_acquire("a").unwrap(); + let held_a1 = low.try_acquire("a").unwrap(); + + let acquire_a = high.acquire("a"); + let mut acquire_a = pin!(acquire_a); + assert!(poll_once(acquire_a.as_mut()).is_pending()); + + let acquire_b = high.acquire("b"); + let mut acquire_b = pin!(acquire_b); + assert!(poll_once(acquire_b.as_mut()).is_pending()); + + drop(held_a0); + assert!(poll_once(acquire_a.as_mut()).is_pending()); + let permit_b = match poll_once(acquire_b.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("owner b has the smaller global share"), + }; + + drop(permit_b); + let permit_a = match poll_once(acquire_a.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("owner a should receive the next permit"), + }; + drop((held_a1, permit_a)); +} + +#[test] +fn same_owner_can_wait_at_multiple_priorities() { + let priorities = PriorityShare::new([1, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held = low.try_acquire("held").unwrap(); + + let low_waiter = low.acquire("owner"); + let mut low_waiter = pin!(low_waiter); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + + let high_waiter = high.acquire("owner"); + let mut high_waiter = pin!(high_waiter); + assert!(poll_once(high_waiter.as_mut()).is_pending()); + + drop(held); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + let high_permit = match poll_once(high_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("the owner's higher-priority queue should be selected"), + }; + + drop(high_permit); + let low_permit = match poll_once(low_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("the owner's lower-priority queue should remain attached"), + }; + drop(low_permit); +} + +#[test] +fn cancelling_one_priority_preserves_the_owners_other_queue() { + let priorities = PriorityShare::new([1, 0]); + let low = &priorities[0]; + let high = &priorities[1]; + let held = low.try_acquire("held").unwrap(); + + let low_waiter = low.acquire("owner"); + let mut low_waiter = pin!(low_waiter); + assert!(poll_once(low_waiter.as_mut()).is_pending()); + + let mut high_waiter = Box::pin(high.acquire("owner")); + assert!(poll_once(high_waiter.as_mut()).is_pending()); + drop(high_waiter); + + drop(held); + let low_permit = match poll_once(low_waiter.as_mut()) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("cancelling one queue must not detach the other"), + }; + drop(low_permit); +} + +#[test] +fn single_priority_matches_fair_share() { + let fair = FairShare::new(2); + let mut priorities = PriorityShare::new([2]); + let priority = priorities.pop().unwrap(); + + let fair_a0 = fair.try_acquire("a").unwrap(); + let fair_a1 = fair.try_acquire("a").unwrap(); + let priority_a0 = priority.try_acquire("a").unwrap(); + let priority_a1 = priority.try_acquire("a").unwrap(); + + let fair_a = fair.acquire("a"); + let mut fair_a = pin!(fair_a); + let priority_a = priority.acquire("a"); + let mut priority_a = pin!(priority_a); + assert_eq!( + poll_once(fair_a.as_mut()).is_pending(), + poll_once(priority_a.as_mut()).is_pending() + ); + + let fair_b = fair.acquire("b"); + let mut fair_b = pin!(fair_b); + let priority_b = priority.acquire("b"); + let mut priority_b = pin!(priority_b); + assert_eq!( + poll_once(fair_b.as_mut()).is_pending(), + poll_once(priority_b.as_mut()).is_pending() + ); + + drop((fair_a0, priority_a0)); + assert!(poll_once(fair_a.as_mut()).is_pending()); + assert!(poll_once(priority_a.as_mut()).is_pending()); + let fair_b_permit = poll_once(fair_b.as_mut()); + let priority_b_permit = poll_once(priority_b.as_mut()); + assert!(fair_b_permit.is_ready()); + assert!(priority_b_permit.is_ready()); + + drop((fair_b_permit, priority_b_permit)); + drop((fair_a1, priority_a1)); + assert!(poll_once(fair_a.as_mut()).is_ready()); + assert!(poll_once(priority_a.as_mut()).is_ready()); +} + +#[test] +fn supports_a_custom_hash_builder() { + let mut priorities = PriorityShare::>::with_hasher( + [1], + BuildHasherDefault::default(), + ); + let admission = priorities.pop().unwrap(); + let permit = admission.try_acquire("tenant".to_owned()).unwrap(); + assert_eq!(permit.key(), "tenant"); +} + +#[test] +fn owned_permit_keeps_the_priority_handle_alive() { + let mut priorities = PriorityShare::new([1]); + let admission = Arc::new(priorities.pop().unwrap()); + let permit = admission + .clone() + .try_acquire_owned("tenant") + .expect("a permit should be available"); + + drop(admission); + assert_eq!(permit.key(), &"tenant"); + assert_eq!(permit.priority(), 0); + drop(permit); +} + +#[test] +fn acquire_futures_are_send() { + fn assert_send(_: T) {} + + let mut priorities = PriorityShare::::new([1]); + let admission = priorities.pop().unwrap(); + assert_send(admission.acquire("tenant".to_owned())); + + let admission = Arc::new(admission); + assert_send(admission.acquire_owned("tenant".to_owned())); +} + +#[test] +fn deterministic_events_match_capacity_model() { + for capacity_increments in [&[2, 1, 1][..], &[0, 2, 0, 1], &[1, 0, 2, 0]] { + let priorities = PriorityShare::new(capacity_increments); + let mut held = Vec::new(); + let mut seed = 0x4d59_5df4_d0f3_3173u64; + + for owner in 0..2_000usize { + seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + + if !held.is_empty() && seed % 3 == 0 { + let index = seed as usize % held.len(); + drop(held.swap_remove(index)); + } else { + let priority = seed as usize % capacity_increments.len(); + let expected = can_admit(capacity_increments, held.len(), priority); + let permit = priorities[priority].try_acquire(owner); + assert_eq!(permit.is_some(), expected); + if let Some(permit) = permit { + held.push(permit); + } + } + } + } +} + +fn can_admit(capacity_increments: &[usize], held: usize, priority: usize) -> bool { + held < capacity_increments[..=priority].iter().sum() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn stress_test_preserves_shared_capacity_limit() { + const CAPACITY_INCREMENTS: [usize; 3] = [2, 2, 1]; + + let priorities: Vec<_> = PriorityShare::::new(CAPACITY_INCREMENTS) + .into_iter() + .map(Arc::new) + .collect(); + let total_capacity = CAPACITY_INCREMENTS.iter().sum::(); + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::new(); + + for admission in &priorities { + for owner in 0..8usize { + for _ in 0..16usize { + let admission = admission.clone(); + let active = active.clone(); + let max_active = max_active.clone(); + handles.push(tokio::spawn(async move { + let _permit = admission.acquire_owned(owner).await; + let active_count = active.fetch_add(1, Ordering::SeqCst) + 1; + assert!(active_count <= total_capacity); + max_active.fetch_max(active_count, Ordering::SeqCst); + + tokio::task::yield_now().await; + active.fetch_sub(1, Ordering::SeqCst); + })); + } + } + } + + for handle in handles { + handle.await.unwrap(); + } + + assert_eq!(active.load(Ordering::SeqCst), 0); + assert!(max_active.load(Ordering::SeqCst) <= total_capacity); + let high = &priorities[2]; + let permits: Vec<_> = (0..total_capacity) + .map(|owner| high.try_acquire(owner).unwrap()) + .collect(); + assert!(high.try_acquire(total_capacity).is_none()); + drop(permits); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 149583c..409e6cd 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -18,6 +18,9 @@ use asyncband::admission::FairShare; use asyncband::admission::FairSharePermit; use asyncband::admission::OwnedFairSharePermit; +use asyncband::admission::OwnedPrioritySharePermit; +use asyncband::admission::PriorityShare; +use asyncband::admission::PrioritySharePermit; use asyncband::barrier::Barrier; use asyncband::broadcast; use asyncband::condvar::Condvar; @@ -48,6 +51,9 @@ fn public_types_are_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::(); assert_send_and_sync::(); @@ -95,6 +101,9 @@ fn public_types_are_unpin() { assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); assert_unpin::(); assert_unpin::(); assert_unpin::();