use crate::doorbell::Doorbell;
use crossbeam_utils::CachePadded;
use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, Wake, Waker};
#[cfg(loom)]
use loom::sync::atomic::{AtomicU64, Ordering};
#[cfg(not(loom))]
use std::sync::atomic::{AtomicU64, Ordering};
const WORD_BITS: usize = 64;
const MAX_WORDS: usize = 1_024;
const SUMMARY_WORDS: usize = MAX_WORDS.div_ceil(WORD_BITS);
const _: () = assert!(
MAX_WORDS <= SUMMARY_WORDS * WORD_BITS,
"SUMMARY_WORDS is too small to address MAX_WORDS"
);
const _: () = assert!(MAX_CAPACITY >= 65_536, "MAX_WORDS was lowered below the design target");
pub const MAX_CAPACITY: usize = WORD_BITS * MAX_WORDS;
#[inline]
const fn low_mask(bits: usize) -> u64 {
if bits == 0 {
0
} else if bits >= 64 {
u64::MAX
} else {
(1u64 << bits) - 1
}
}
type ReadyWord = CachePadded<AtomicU64>;
struct ReadySet {
words: Box<[ReadyWord]>,
summary: [CachePadded<AtomicU64>; SUMMARY_WORDS],
summaries: usize,
slots: usize,
}
impl ReadySet {
fn new(slots: usize) -> Self {
assert!(slots <= MAX_CAPACITY, "Outstanding capacity exceeds {MAX_CAPACITY}");
let words = slots.div_ceil(WORD_BITS);
Self {
words: (0..words).map(|_| CachePadded::new(AtomicU64::new(0))).collect(),
summary: std::array::from_fn(|_| CachePadded::new(AtomicU64::new(0))),
summaries: if words > 1 { words.div_ceil(WORD_BITS) } else { 0 },
slots,
}
}
#[inline]
fn words(&self) -> usize {
self.words.len()
}
#[inline]
fn summaries(&self) -> usize {
self.summaries
}
#[inline]
fn mark(&self, slot: usize) -> bool {
debug_assert!(slot < self.slots);
let word = slot / WORD_BITS;
let bit = 1u64 << (slot % WORD_BITS);
let previous = self.words[word].fetch_or(bit, Ordering::Release);
if previous == 0 {
self.announce(word);
}
previous & bit == 0
}
#[inline]
fn announce(&self, word: usize) {
if self.summaries == 0 {
return;
}
self.summary[word / WORD_BITS].fetch_or(1u64 << (word % WORD_BITS), Ordering::Release);
}
#[inline]
fn take_single(&self) -> u64 {
debug_assert_eq!(self.words.len(), 1);
self.take_word(0)
}
#[inline]
fn take_summary(&self, summary: usize) -> u64 {
self.summary[summary].swap(0, Ordering::Acquire) & self.summary_mask(summary)
}
#[inline]
fn take_word(&self, word: usize) -> u64 {
self.words[word].swap(0, Ordering::Acquire) & self.valid_mask(word)
}
#[inline]
fn restore_word(&self, word: usize, bits: u64) {
let bits = bits & self.valid_mask(word);
if bits == 0 {
return;
}
let previous = self.words[word].fetch_or(bits, Ordering::Release);
if previous == 0 {
self.announce(word);
}
}
#[inline]
fn restore_summary(&self, summary: usize, words: u64) {
let words = words & self.summary_mask(summary);
if words != 0 {
self.summary[summary].fetch_or(words, Ordering::Release);
}
}
#[inline]
fn has_ready(&self) -> bool {
match self.words.len() {
0 => false,
1 => self.words[0].load(Ordering::Acquire) & self.valid_mask(0) != 0,
_ => (0..self.summaries).any(|summary| {
self.summary[summary].load(Ordering::Acquire) & self.summary_mask(summary) != 0
}),
}
}
#[inline]
fn summary_mask(&self, summary: usize) -> u64 {
let remainder = self.words.len() % WORD_BITS;
if summary + 1 == self.summaries && remainder != 0 { low_mask(remainder) } else { u64::MAX }
}
#[inline]
fn valid_mask(&self, word: usize) -> u64 {
let remainder = self.slots % WORD_BITS;
if word + 1 == self.words.len() && remainder != 0 { low_mask(remainder) } else { u64::MAX }
}
}
struct Shared {
ready: ReadySet,
bell: Doorbell,
}
impl Shared {
#[inline]
fn notify(&self, slot: usize) {
if self.bell.is_closed() {
return;
}
if self.ready.mark(slot) {
self.bell.ring();
}
}
}
struct SlotWaker {
shared: Arc<Shared>,
slot: usize,
}
impl Wake for SlotWaker {
#[inline]
fn wake(self: Arc<Self>) {
self.shared.notify(self.slot);
}
#[inline]
fn wake_by_ref(self: &Arc<Self>) {
self.shared.notify(self.slot);
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Harvest {
pub polled: usize,
pub finished: usize,
pub more_ready: bool,
}
pub struct PushError<F> {
future: F,
}
#[allow(dead_code)]
impl<F> PushError<F> {
pub fn into_future(self) -> F {
self.future
}
pub fn future(&self) -> &F {
&self.future
}
}
impl<F> fmt::Debug for PushError<F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("PushError { full: true, .. }")
}
}
impl<F> fmt::Display for PushError<F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("the outstanding set is full")
}
}
impl<F> std::error::Error for PushError<F> {}
trait Storage<F: Future> {
fn with_capacity(capacity: usize) -> Self
where
Self: Sized;
fn capacity(&self) -> usize;
fn is_occupied(&self, index: usize) -> bool;
fn insert(&mut self, index: usize, future: F);
fn poll(&mut self, index: usize, cx: &mut Context<'_>) -> Poll<F::Output>;
fn remove(&mut self, index: usize);
}
struct BoxedStorage<F> {
slots: Box<[Pin<Box<Option<F>>>]>,
}
impl<F: Future> Storage<F> for BoxedStorage<F> {
fn with_capacity(capacity: usize) -> Self {
Self { slots: (0..capacity).map(|_| Box::pin(None::<F>)).collect() }
}
#[inline]
fn capacity(&self) -> usize {
self.slots.len()
}
#[inline]
fn is_occupied(&self, index: usize) -> bool {
self.slots[index].as_ref().get_ref().is_some()
}
#[inline]
fn insert(&mut self, index: usize, future: F) {
debug_assert!(!self.is_occupied(index));
self.slots[index].as_mut().set(Some(future));
}
#[inline]
fn poll(&mut self, index: usize, cx: &mut Context<'_>) -> Poll<F::Output> {
self.slots[index]
.as_mut()
.as_pin_mut()
.expect("occupied boxed slot contains a future")
.poll(cx)
}
#[inline]
fn remove(&mut self, index: usize) {
debug_assert!(self.is_occupied(index));
self.slots[index].as_mut().set(None);
}
}
struct RestoreGuard<'a> {
shared: &'a Shared,
summary_index: usize,
summary: u64,
deferred_summary_index: usize,
deferred_summary: u64,
current_word: usize,
current_bits: u64,
deferred_word: usize,
deferred_bits: u64,
armed: bool,
}
impl<'a> RestoreGuard<'a> {
fn new(shared: &'a Shared) -> Self {
Self {
shared,
summary_index: 0,
summary: 0,
deferred_summary_index: 0,
deferred_summary: 0,
current_word: 0,
current_bits: 0,
deferred_word: 0,
deferred_bits: 0,
armed: true,
}
}
#[inline]
fn set_summary(&mut self, index: usize, words: u64) {
debug_assert_eq!(self.summary, 0);
self.summary_index = index;
self.summary = words;
}
#[inline]
fn defer_summary(&mut self, index: usize, words: u64) {
debug_assert_eq!(self.deferred_summary, 0);
self.deferred_summary_index = index;
self.deferred_summary = words;
}
#[inline]
fn activate_deferred_summary(&mut self) {
debug_assert_eq!(self.summary, 0);
self.summary_index = self.deferred_summary_index;
self.summary = std::mem::take(&mut self.deferred_summary);
}
#[inline]
fn next_announced(&mut self) -> Option<usize> {
if self.summary == 0 {
return None;
}
let bit = self.summary.trailing_zeros() as usize;
self.summary &= self.summary - 1;
Some(self.summary_index * WORD_BITS + bit)
}
#[inline]
fn set_current(&mut self, word: usize, bits: u64) {
debug_assert_eq!(self.current_bits, 0);
self.current_word = word;
self.current_bits = bits;
}
#[inline]
fn set_deferred(&mut self, word: usize, bits: u64) {
debug_assert_eq!(self.deferred_bits, 0);
self.deferred_word = word;
self.deferred_bits = bits;
}
#[inline]
fn activate_deferred(&mut self) {
debug_assert_eq!(self.current_bits, 0);
self.current_word = self.deferred_word;
self.current_bits = std::mem::take(&mut self.deferred_bits);
}
fn restore(&mut self, wake_owner: bool) {
if !self.armed {
return;
}
let had_work = self.summary != 0
|| self.deferred_summary != 0
|| self.current_bits != 0
|| self.deferred_bits != 0;
if self.current_bits != 0 {
self.shared.ready.restore_word(self.current_word, self.current_bits);
}
if self.deferred_bits != 0 {
self.shared.ready.restore_word(self.deferred_word, self.deferred_bits);
}
if self.summary != 0 {
self.shared.ready.restore_summary(self.summary_index, self.summary);
}
if self.deferred_summary != 0 {
self.shared.ready.restore_summary(self.deferred_summary_index, self.deferred_summary);
}
self.summary = 0;
self.deferred_summary = 0;
self.current_bits = 0;
self.deferred_bits = 0;
self.armed = false;
if wake_owner && had_work {
self.shared.bell.ring();
}
}
#[inline]
fn disarm(&mut self) {
debug_assert_eq!(self.summary, 0);
debug_assert_eq!(self.deferred_summary, 0);
debug_assert_eq!(self.current_bits, 0);
debug_assert_eq!(self.deferred_bits, 0);
self.armed = false;
}
}
impl Drop for RestoreGuard<'_> {
fn drop(&mut self) {
self.restore(true);
}
}
struct PollState<'a, F, S, G>
where
F: Future,
S: Storage<F>,
G: FnMut(F::Output),
{
storage: &'a mut S,
wakers: &'a [Waker],
free: &'a mut Vec<u32>,
live: &'a mut usize,
cursor: &'a mut usize,
capacity: usize,
cap: usize,
report: Harvest,
out: &'a mut G,
_future: PhantomData<fn() -> F>,
}
impl<F, S, G> PollState<'_, F, S, G>
where
F: Future,
S: Storage<F>,
G: FnMut(F::Output),
{
fn poll_current(&mut self, guard: &mut RestoreGuard<'_>) -> bool {
while guard.current_bits != 0 {
let bit = guard.current_bits.trailing_zeros() as usize;
let mask = 1u64 << bit;
let index = guard.current_word * WORD_BITS + bit;
if index >= self.capacity {
guard.current_bits &= !mask;
continue;
}
if !self.storage.is_occupied(index) {
guard.current_bits &= !mask;
*self.cursor = if index + 1 == self.capacity { 0 } else { index + 1 };
continue;
}
if self.report.polled == self.cap {
return true;
}
self.report.polled += 1;
let mut cx = Context::from_waker(&self.wakers[index]);
if let Poll::Ready(output) = self.storage.poll(index, &mut cx) {
self.storage.remove(index);
self.free.push(index as u32);
*self.live -= 1;
(self.out)(output);
self.report.finished += 1;
}
guard.current_bits &= !mask;
*self.cursor = if index + 1 == self.capacity { 0 } else { index + 1 };
}
false
}
}
struct Inner<F: Future, S: Storage<F>> {
storage: S,
wakers: Vec<Waker>,
free: Vec<u32>,
shared: Arc<Shared>,
live: usize,
cursor: usize,
_future: PhantomData<fn() -> F>,
}
#[allow(dead_code)]
impl<F: Future, S: Storage<F>> Inner<F, S> {
fn with_capacity(capacity: usize) -> Self {
assert!(capacity <= MAX_CAPACITY, "Outstanding capacity exceeds {MAX_CAPACITY}");
assert!(capacity <= u32::MAX as usize, "Outstanding capacity exceeds u32 indexing");
let storage = S::with_capacity(capacity);
debug_assert_eq!(storage.capacity(), capacity);
let shared = Arc::new(Shared { ready: ReadySet::new(capacity), bell: Doorbell::new() });
let wakers = (0..capacity)
.map(|slot| Waker::from(Arc::new(SlotWaker { shared: shared.clone(), slot })))
.collect();
Self {
storage,
wakers,
free: (0..capacity as u32).rev().collect(),
shared,
live: 0,
cursor: 0,
_future: PhantomData,
}
}
#[inline]
fn capacity(&self) -> usize {
self.storage.capacity()
}
#[inline]
fn len(&self) -> usize {
self.live
}
#[inline]
fn is_empty(&self) -> bool {
self.live == 0
}
#[inline]
fn available(&self) -> usize {
self.free.len()
}
#[inline]
fn register_owner(&self, waker: &Waker) {
self.shared.bell.register(waker);
}
#[inline]
fn try_push(&mut self, future: F) -> Result<(), PushError<F>> {
let Some(index) = self.free.pop() else {
return Err(PushError { future });
};
let index = index as usize;
debug_assert!(!self.storage.is_occupied(index));
self.storage.insert(index, future);
self.live += 1;
let _ = self.shared.ready.mark(index);
Ok(())
}
#[track_caller]
fn push(&mut self, future: F) {
if self.try_push(future).is_err() {
panic!("outstanding set overflow: capacity must cover all in-flight budgets");
}
}
fn harvest<G>(&mut self, cap: usize, mut out: G) -> Harvest
where
G: FnMut(F::Output),
{
let capacity = self.capacity();
if cap == 0 || capacity == 0 {
return Harvest { polled: 0, finished: 0, more_ready: self.shared.ready.has_ready() };
}
let Self { storage, wakers, free, shared, live, cursor, .. } = self;
let words = shared.ready.words();
let start_word = *cursor / WORD_BITS;
let start_bit = *cursor % WORD_BITS;
let mut polling = PollState::<F, S, G> {
storage,
wakers,
free,
live,
cursor,
capacity,
cap,
report: Harvest::default(),
out: &mut out,
_future: PhantomData,
};
if words == 1 {
let bits = shared.ready.take_single();
let mut guard = RestoreGuard::new(shared);
let before_cursor = bits & low_mask(start_bit);
let from_cursor = bits & !low_mask(start_bit);
guard.set_deferred(0, before_cursor);
guard.set_current(0, from_cursor);
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
guard.activate_deferred();
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
guard.disarm();
polling.report.more_ready = shared.ready.has_ready();
return polling.report;
}
let summaries = shared.ready.summaries();
let start_summary = start_word / WORD_BITS;
let mut guard = RestoreGuard::new(shared);
let announced = shared.ready.take_summary(start_summary);
let start_word_bit = start_word % WORD_BITS;
guard.defer_summary(start_summary, announced & low_mask(start_word_bit));
guard.set_summary(start_summary, announced & !low_mask(start_word_bit));
if guard.summary & (1u64 << start_word_bit) != 0 {
guard.summary &= !(1u64 << start_word_bit);
let bits = shared.ready.take_word(start_word);
guard.set_deferred(start_word, bits & low_mask(start_bit));
guard.set_current(start_word, bits & !low_mask(start_bit));
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
}
for step in 0..summaries {
if step > 0 {
debug_assert_eq!(guard.summary, 0);
let summary = (start_summary + step) % summaries;
guard.set_summary(summary, shared.ready.take_summary(summary));
}
while let Some(word) = guard.next_announced() {
guard.set_current(word, shared.ready.take_word(word));
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
}
}
guard.activate_deferred_summary();
while let Some(word) = guard.next_announced() {
guard.set_current(word, shared.ready.take_word(word));
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
}
guard.activate_deferred();
if polling.poll_current(&mut guard) {
polling.report.more_ready = true;
guard.restore(false);
return polling.report;
}
debug_assert_eq!(guard.summary, 0);
guard.disarm();
polling.report.more_ready = shared.ready.has_ready();
polling.report
}
#[cfg(all(test, not(loom)))]
fn check_invariants(&self) -> Result<(), &'static str> {
if self.storage.capacity() != self.wakers.len() {
return Err("storage and waker capacities differ");
}
if self.live + self.free.len() != self.capacity() {
return Err("live plus free does not equal capacity");
}
if self.capacity() != 0 && self.cursor >= self.capacity() {
return Err("scan cursor is outside capacity");
}
let mut free_seen = vec![false; self.capacity()];
for index in &self.free {
let index = *index as usize;
if index >= self.capacity() || free_seen[index] {
return Err("free list is out of range or duplicated");
}
if self.storage.is_occupied(index) {
return Err("free list contains an occupied slot");
}
free_seen[index] = true;
}
let occupied =
(0..self.capacity()).filter(|index| self.storage.is_occupied(*index)).count();
if occupied != self.live {
return Err("live counter disagrees with storage");
}
for (index, free) in free_seen.into_iter().enumerate() {
if free == self.storage.is_occupied(index) {
return Err("a slot is neither exclusively free nor occupied");
}
}
Ok(())
}
}
impl<F: Future, S: Storage<F>> Drop for Inner<F, S> {
fn drop(&mut self) {
self.shared.bell.close();
}
}
pub struct BoxedOutstanding<F: Future> {
inner: Inner<F, BoxedStorage<F>>,
}
macro_rules! impl_outstanding {
($name:ident, $storage:ident) => {
#[allow(dead_code)]
impl<F: Future> $name<F> {
pub fn with_capacity(capacity: usize) -> Self {
Self { inner: Inner::<F, $storage<F>>::with_capacity(capacity) }
}
#[inline]
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
#[inline]
pub fn len(&self) -> usize {
self.inner.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[inline]
pub fn available(&self) -> usize {
self.inner.available()
}
#[inline]
pub fn register_owner(&self, waker: &Waker) {
self.inner.register_owner(waker);
}
#[inline]
pub fn try_push(&mut self, future: F) -> Result<(), PushError<F>> {
self.inner.try_push(future)
}
#[track_caller]
pub fn push(&mut self, future: F) {
self.inner.push(future);
}
#[inline]
pub fn harvest(&mut self, cap: usize, out: impl FnMut(F::Output)) -> Harvest {
self.inner.harvest(cap, out)
}
#[inline]
pub fn poll_harvest(
&mut self,
cx: &mut Context<'_>,
cap: usize,
out: impl FnMut(F::Output),
) -> Harvest {
self.inner.register_owner(cx.waker());
self.inner.harvest(cap, out)
}
}
};
}
impl_outstanding!(BoxedOutstanding, BoxedStorage);
pub type Outstanding<F> = BoxedOutstanding<F>;
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use futures::task::noop_waker;
use std::collections::HashSet;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
struct Countdown {
remaining: u8,
payload: u64,
}
impl Future for Countdown {
type Output = u64;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.remaining == 0 {
Poll::Ready(self.payload)
} else {
self.remaining -= 1;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
fn countdown(remaining: u8, payload: u64) -> Countdown {
Countdown { remaining, payload }
}
fn drain<F, S>(set: &mut Inner<F, S>, cap: usize) -> Vec<F::Output>
where
F: Future,
S: Storage<F>,
{
let mut output = Vec::new();
let mut passes = 0usize;
while !set.is_empty() {
set.harvest(cap, |value| output.push(value));
passes += 1;
assert!(passes < 1_000_000, "ready work was stranded");
}
output
}
fn exercise_recycling<S: Storage<Countdown>>() {
let mut set = Inner::<Countdown, S>::with_capacity(8);
for round in 0..2_000u64 {
set.try_push(countdown((round % 4) as u8, round)).unwrap();
set.try_push(countdown(((round + 1) % 4) as u8, round + 10_000)).unwrap();
let mut done = drain(&mut set, 3);
done.sort_unstable();
assert_eq!(done, [round, round + 10_000]);
assert_eq!(set.check_invariants(), Ok(()));
}
}
#[test]
fn slots_recycle_and_every_future_completes() {
exercise_recycling::<BoxedStorage<Countdown>>();
}
#[test]
fn compiler_generated_not_unpin_futures_are_polled_in_place() {
async fn job(value: u64) -> u64 {
countdown(2, value).await
}
let mut boxed = BoxedOutstanding::with_capacity(2);
boxed.push(job(1));
assert_eq!(drain(&mut boxed.inner, 1), [1]);
let mut slab = BoxedOutstanding::with_capacity(2);
slab.push(job(2));
assert_eq!(drain(&mut slab.inner, 1), [2]);
}
fn exercise_true_multiword_rotation<S: Storage<Countdown>>() {
let mut set = Inner::<Countdown, S>::with_capacity(65);
set.push(countdown(0, 999));
assert_eq!(drain(&mut set, 1), [999]);
assert_eq!(set.cursor, 1);
for payload in 0..65u64 {
set.push(countdown(0, payload));
}
let mut first = Vec::new();
let report = set.harvest(64, |value| first.push(value));
first.sort_unstable();
assert_eq!(first, (1..65).collect::<Vec<_>>());
assert_eq!(report.polled, 64);
assert!(report.more_ready);
let mut last = Vec::new();
set.harvest(1, |value| last.push(value));
assert_eq!(last, [0]);
assert!(set.is_empty());
assert_eq!(set.check_invariants(), Ok(()));
}
fn exercise_rotation_across_a_summary_word<S: Storage<Countdown>>() {
const CAPACITY: usize = 4_160;
let mut set = Inner::<Countdown, S>::with_capacity(CAPACITY);
for payload in 0..CAPACITY as u64 {
set.push(countdown(0, payload));
}
let mut warmed = Vec::new();
set.harvest(4_100, |value| warmed.push(value));
let resume = set.cursor;
assert!(resume > 4_096, "the cursor should sit in the second summary word, not {resume}");
for payload in warmed {
set.push(countdown(0, payload));
}
let mut served = Vec::new();
set.harvest(CAPACITY, |value| served.push(value));
served.sort_unstable();
assert_eq!(
served,
(0..CAPACITY as u64).collect::<Vec<_>>(),
"a rotation beginning at {resume} missed slots"
);
assert!(set.is_empty());
assert_eq!(set.check_invariants(), Ok(()));
}
#[test]
#[cfg_attr(miri, ignore = "4,160 slots through a full rotation is beyond Miri")]
fn capped_rotation_is_fair_across_a_summary_word_boundary() {
exercise_rotation_across_a_summary_word::<BoxedStorage<Countdown>>();
}
#[test]
fn capped_rotation_is_fair_across_the_64_slot_boundary() {
exercise_true_multiword_rotation::<BoxedStorage<Countdown>>();
}
fn exercise_boundaries<S: Storage<Countdown>>() {
let ceilings: &[usize] = if cfg!(miri) {
&[0, 1, 63, 64, 65, 127, 128, 129]
} else {
&[0, 1, 63, 64, 65, 127, 128, 129, 4_095, 4_096, 4_097, 8_191, 8_192, 8_193, 16_384]
};
for &capacity in ceilings {
let mut set = Inner::<Countdown, S>::with_capacity(capacity);
for token in 0..capacity as u64 {
set.push(countdown((token % 3) as u8, token));
}
let mut done = drain(&mut set, 17);
done.sort_unstable();
assert_eq!(done, (0..capacity as u64).collect::<Vec<_>>(), "capacity {capacity}");
assert_eq!(set.check_invariants(), Ok(()));
}
}
#[test]
fn capacity_and_word_boundaries_hold() {
exercise_boundaries::<BoxedStorage<Countdown>>();
}
struct DropReady {
drops: Arc<AtomicUsize>,
}
impl Future for DropReady {
type Output = ();
fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<()> {
Poll::Ready(())
}
}
impl Drop for DropReady {
fn drop(&mut self) {
self.drops.fetch_add(1, StdOrdering::Relaxed);
}
}
fn exercise_immediate_drop<S: Storage<DropReady>>() {
let drops = Arc::new(AtomicUsize::new(0));
let mut set = Inner::<DropReady, S>::with_capacity(1);
set.push(DropReady { drops: drops.clone() });
let report = set.harvest(1, |_| {
assert_eq!(drops.load(StdOrdering::Relaxed), 1, "F must be dropped before callback");
});
assert_eq!(report.finished, 1);
assert_eq!(drops.load(StdOrdering::Relaxed), 1);
}
#[test]
fn a_completed_future_is_dropped_immediately() {
exercise_immediate_drop::<BoxedStorage<DropReady>>();
}
fn exercise_callback_unwind<S: Storage<Countdown>>() {
let mut set = Inner::<Countdown, S>::with_capacity(70);
for token in 0..70 {
set.push(countdown(0, token));
}
let panic = catch_unwind(AssertUnwindSafe(|| {
set.harvest(usize::MAX, |_| panic!("callback failure"));
}));
assert!(panic.is_err());
assert_eq!(set.len(), 69, "the completed current future was released");
let mut remaining = drain(&mut set, usize::MAX);
remaining.sort_unstable();
assert_eq!(remaining.len(), 69, "unvisited ready bits were restored");
assert_eq!(set.check_invariants(), Ok(()));
}
#[test]
fn callback_unwind_restores_same_and_later_words() {
exercise_callback_unwind::<BoxedStorage<Countdown>>();
}
fn exercise_stale_wakes<S: Storage<Countdown>>() {
let mut set = Inner::<Countdown, S>::with_capacity(1);
set.push(countdown(0, 1));
let stale = set.wakers[0].clone();
assert_eq!(drain(&mut set, 1), [1]);
stale.wake_by_ref();
let empty = set.harvest(1, |_| panic!("a stale wake invented output"));
assert_eq!(empty.finished, 0);
set.push(countdown(1, 2));
stale.wake_by_ref();
let first = set.harvest(1, |_| panic!("replacement should pend once"));
assert_eq!(first.polled, 1);
assert_eq!(drain(&mut set, 1), [2]);
}
#[test]
fn stale_wakes_are_safe_before_and_after_refill() {
exercise_stale_wakes::<BoxedStorage<Countdown>>();
}
#[test]
fn zero_cap_is_a_real_noop() {
let mut set = BoxedOutstanding::with_capacity(1);
set.push(countdown(0, 1));
let report = set.harvest(0, |_| panic!("zero cap polled a future"));
assert_eq!(report, Harvest { polled: 0, finished: 0, more_ready: true });
assert_eq!(drain(&mut set.inner, 1), [1]);
}
#[test]
fn full_try_push_returns_the_future() {
let mut set = BoxedOutstanding::with_capacity(1);
set.try_push(countdown(1, 1)).unwrap();
let error = set.try_push(countdown(2, 2)).unwrap_err();
assert_eq!(error.into_future().payload, 2);
}
struct CountWake(AtomicUsize);
impl Wake for CountWake {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, StdOrdering::Relaxed);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.fetch_add(1, StdOrdering::Relaxed);
}
}
#[test]
fn duplicate_slot_wakes_coalesce_until_the_bit_is_taken() {
let owner = Arc::new(CountWake(AtomicUsize::new(0)));
let owner_waker = Waker::from(owner.clone());
let mut set = BoxedOutstanding::<Countdown>::with_capacity(1);
set.register_owner(&owner_waker);
set.inner.wakers[0].wake_by_ref();
set.inner.wakers[0].wake_by_ref();
assert_eq!(owner.0.load(StdOrdering::Relaxed), 1);
assert_eq!(set.harvest(1, |_| unreachable!()).polled, 0);
set.register_owner(&owner_waker);
set.inner.wakers[0].wake_by_ref();
assert_eq!(owner.0.load(StdOrdering::Relaxed), 2);
}
#[test]
fn dropping_the_set_closes_stale_slot_wakers_and_releases_owner() {
let owner = Arc::new(CountWake(AtomicUsize::new(0)));
let owner_waker = Waker::from(owner.clone());
let stale = {
let set = BoxedOutstanding::<Countdown>::with_capacity(1);
set.register_owner(&owner_waker);
assert_eq!(Arc::strong_count(&owner), 3, "Arc, Waker, and AtomicWaker");
set.inner.wakers[0].clone()
};
assert_eq!(Arc::strong_count(&owner), 2, "drop must clear AtomicWaker");
stale.wake_by_ref();
assert_eq!(owner.0.load(StdOrdering::Relaxed), 0, "closed stale wake scheduled owner");
}
fn battle<S: Storage<Countdown>>() {
const CAPACITY: usize = 129;
const TOTAL: u64 = 20_000;
let mut set = Inner::<Countdown, S>::with_capacity(CAPACITY);
let mut next = 0u64;
let mut complete = HashSet::with_capacity(TOTAL as usize);
let mut rng = 0x1234_5678_9abc_def0u64;
let mut random = || {
rng ^= rng << 13;
rng ^= rng >> 7;
rng ^= rng << 17;
rng
};
while next < CAPACITY as u64 {
set.push(countdown((random() % 5) as u8, next));
next += 1;
}
let mut passes = 0usize;
while complete.len() < TOTAL as usize {
let cap = (random() % 31 + 1) as usize;
let mut finished = Vec::new();
set.harvest(cap, |token| finished.push(token));
for token in finished {
assert!(complete.insert(token), "future completed twice: {token}");
if next < TOTAL {
set.push(countdown((random() % 5) as u8, next));
next += 1;
}
}
passes += 1;
assert!(passes < 1_000_000, "battle run starved ready work");
assert_eq!(set.check_invariants(), Ok(()));
}
assert!(set.is_empty());
}
#[test]
fn randomized_capped_refill_matches_an_independent_model() {
battle::<BoxedStorage<Countdown>>();
}
#[test]
fn poll_harvest_registers_before_checking() {
let owner = noop_waker();
let mut cx = Context::from_waker(&owner);
let mut set = BoxedOutstanding::with_capacity(1);
set.push(countdown(0, 7));
let mut output = Vec::new();
let report = set.poll_harvest(&mut cx, 1, |value| output.push(value));
assert_eq!(report.finished, 1);
assert_eq!(output, [7]);
}
#[test]
#[should_panic(expected = "Outstanding capacity exceeds")]
fn capacity_beyond_the_configured_ceiling_is_rejected() {
let _ = BoxedOutstanding::<Countdown>::with_capacity(MAX_CAPACITY + 1);
}
}
#[cfg(all(test, loom))]
mod loom_tests {
use super::{ReadySet, low_mask};
use loom::sync::Arc;
use loom::sync::atomic::{AtomicBool, Ordering};
use loom::thread;
fn collect(ready: &ReadySet) -> Vec<u64> {
match ready.words() {
0 => Vec::new(),
1 => vec![ready.take_single()],
words => {
let summary = ready.take_summary(0);
(0..words)
.map(
|word| {
if summary & (1u64 << word) != 0 { ready.take_word(word) } else { 0 }
},
)
.collect()
}
}
}
#[test]
fn loom_publish_then_wake_never_exposes_wake_before_bit() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(128));
let woken = Arc::new(AtomicBool::new(false));
let producer = {
let ready = ready.clone();
let woken = woken.clone();
thread::spawn(move || {
ready.mark(70);
woken.store(true, Ordering::Release);
})
};
let saw_wake = woken.load(Ordering::Acquire);
let summary = ready.take_summary(0);
let bits = if summary & (1 << 1) != 0 { ready.take_word(1) } else { 0 };
if saw_wake {
assert_ne!(bits & (1 << 6), 0, "wake became visible before slot bit");
}
producer.join().unwrap();
let later = collect(&ready);
assert_ne!(bits | later.get(1).copied().unwrap_or(0), 0, "ready mark was lost");
});
}
#[test]
fn loom_mark_racing_summary_and_word_take_is_never_lost() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(128));
let producer = {
let ready = ready.clone();
thread::spawn(move || ready.mark(67))
};
let first = collect(&ready);
producer.join().unwrap();
let second = collect(&ready);
let observed = first.get(1).copied().unwrap_or(0) | second.get(1).copied().unwrap_or(0);
assert_ne!(observed & (1 << 3), 0);
});
}
#[test]
fn loom_concurrent_marks_in_different_words_survive_summary_coalescing() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(192));
let left = {
let ready = ready.clone();
thread::spawn(move || ready.mark(2))
};
let right = {
let ready = ready.clone();
thread::spawn(move || ready.mark(130))
};
left.join().unwrap();
right.join().unwrap();
let words = collect(&ready);
assert_ne!(words[0] & (1 << 2), 0);
assert_ne!(words[2] & (1 << 2), 0);
});
}
#[test]
fn loom_duplicate_marks_have_exactly_one_wake_transition() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(64));
let left = {
let ready = ready.clone();
thread::spawn(move || ready.mark(5))
};
let right = {
let ready = ready.clone();
thread::spawn(move || ready.mark(5))
};
let transitions =
usize::from(left.join().unwrap()) + usize::from(right.join().unwrap());
assert_eq!(transitions, 1, "duplicate marks emitted duplicate owner wakes");
assert_eq!(ready.take_single(), 1 << 5);
assert_eq!(ready.take_single(), 0);
});
}
#[test]
fn loom_restore_merges_with_concurrent_marks_in_same_word() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(128));
ready.mark(65);
ready.mark(66);
let summary = ready.take_summary(0);
assert_ne!(summary & (1 << 1), 0);
let taken = ready.take_word(1);
let producer = {
let ready = ready.clone();
thread::spawn(move || ready.mark(73))
};
ready.restore_word(1, taken & !(1 << 1));
producer.join().unwrap();
let words = collect(&ready);
assert_ne!(words[1] & (1 << 2), 0, "restored bit was lost");
assert_ne!(words[1] & (1 << 9), 0, "concurrent bit was overwritten");
});
}
#[test]
fn loom_unvisited_summary_restore_merges_with_new_word() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(192));
ready.mark(1);
ready.mark(70);
let taken_summary = ready.take_summary(0);
let producer = {
let ready = ready.clone();
thread::spawn(move || ready.mark(130))
};
ready.restore_summary(0, taken_summary & !1);
producer.join().unwrap();
let summary = ready.take_summary(0);
assert_ne!(summary & (1 << 1), 0);
assert_ne!(summary & (1 << 2), 0);
assert_eq!(summary & !low_mask(3), 0);
});
}
#[test]
fn loom_register_before_check_cannot_strand_a_later_mark() {
loom::model(|| {
let ready = Arc::new(ReadySet::new(64));
let registered = Arc::new(AtomicBool::new(false));
let woken = Arc::new(AtomicBool::new(false));
let producer = {
let ready = ready.clone();
let registered = registered.clone();
let woken = woken.clone();
thread::spawn(move || {
ready.mark(4);
if registered.load(Ordering::Acquire) {
woken.store(true, Ordering::Release);
}
})
};
registered.store(true, Ordering::Release);
let first = ready.take_single();
producer.join().unwrap();
let second = ready.take_single();
assert_ne!(first | second, 0);
if first == 0 && second == 0 {
assert!(woken.load(Ordering::Acquire));
}
});
}
}