Something went wrong. Try again.
Next Generation WASM Microkernel Operating System
wasm os rust microkernel
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172// Copyright 2025 Jonas Kruckenberg//// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or// http://opensource.org/licenses/MIT>, at your option. This file may not be// copied, modified, or distributed except according to those terms.
use crate::error::Closed;use crate::sync::wake_batch::WakeBatch;use alloc::sync::Arc;use cordyceps::{Linked, List, list};use core::cell::UnsafeCell;use core::marker::PhantomPinned;use core::pin::Pin;use core::ptr::NonNull;use core::sync::atomic::{AtomicUsize, Ordering};use core::task::{Context, Poll, Waker};use core::{fmt, mem, ptr};use mycelium_bitfield::{FromBits, bitfield, enum_from_bits};use pin_project::{pin_project, pinned_drop};use spin::{Mutex, MutexGuard};use util::{CachePadded, loom_const_fn};
/// A queue of waiting tasks which can be [woken in first-in, first-out/// order][wake], or [all at once][wake_all].////// This type is taken from [maitake-sync](https://github.com/hawkw/mycelium/blob/dd0020892564c77ee4c20ffbc2f7f5b046ad54c8/maitake-sync/src/wait_queue.rs#L1577).////// A `WaitQueue` allows any number of tasks to [wait] asynchronously and be/// woken when some event occurs, either [individually][wake] in first-in,/// first-out order, or [all at once][wake_all]. This makes it a vital building/// block of runtime services (such as timers or I/O resources), where it may be/// used to wake a set of tasks when a timer completes or when a resource/// becomes available. It can be equally useful for implementing higher-level/// synchronization primitives: for example, a `WaitQueue` plus an/// [`UnsafeCell`] is essentially an entire implementation of a fair/// asynchronous mutex. Finally, a `WaitQueue` can be a useful/// synchronization primitive on its own: sometimes, you just need to have a/// bunch of tasks wait for something and then wake them all up.////// # Implementation Notes////// This type is currently implemented using [intrusive doubly-linked/// list][ilist].////// The *[intrusive]* aspect of this map is important, as it means that it does/// not allocate memory. Instead, nodes in the linked list are stored in the/// futures of tasks trying to wait for capacity. This means that it is not/// necessary to allocate any heap memory for each task waiting to be woken.////// However, the intrusive linked list introduces one new danger: because/// futures can be *cancelled*, and the linked list nodes live within the/// futures trying to wait on the queue, we *must* ensure that the node/// is unlinked from the list before dropping a cancelled future. Failure to do/// so would result in the list containing dangling pointers. Therefore, we must/// use a *doubly-linked* list, so that nodes can edit both the previous and/// next node when they have to remove themselves. This is kind of a bummer, as/// it means we can't use something nice like this [intrusive queue by Dmitry/// Vyukov][2], and there are not really practical designs for lock-free/// doubly-linked lists that don't rely on some kind of deferred reclamation/// scheme such as hazard pointers or QSBR.////// Instead, we just stick a [`Mutex`] around the linked list, which must be/// acquired to pop nodes from it, or for nodes to remove themselves when/// futures are cancelled. This is a bit sad, but the critical sections for this/// mutex are short enough that we still get pretty good performance despite it.////// [intrusive]: https://fuchsia.dev/fuchsia-src/development/languages/c-cpp/fbl_containers_guide/introduction/// [2]: https://www.1024cores.net/home/lock-free-algorithms/queues/intrusive-mpsc-node-based-queue/// [`Waker`]: Waker/// [wait]: WaitQueue::wait/// [wake]: WaitQueue::wake/// [wake_all]: WaitQueue::wake_all/// [`UnsafeCell`]: UnsafeCell/// [ilist]: linked_list::List#[derive(Debug)]pub struct WaitQueue { /// The wait maps's state variable. state: CachePadded<AtomicUsize>, /// The linked list of waiters. /// /// # Safety /// /// This is protected by a mutex; the mutex *must* be acquired when /// manipulating the linked list, OR when manipulating waiter nodes that may /// be linked into the list. If a node is known to not be linked, it is safe /// to modify that node (such as by waking the stored [`Waker`]) without /// holding the lock; otherwise, it may be modified through the list, so the /// lock must be held when modifying the /// node. queue: Mutex<List<Waiter>>,}
bitfield! { #[derive(Eq, PartialEq)] struct State<usize> { /// The queue's state. const INNER: StateInner;
/// The number of times [`WaitQueue::wake_all`] has been called. const WAKE_ALLS = ..; }}
/// The queue's current state.#[derive(Debug, Copy, Clone, Eq, PartialEq)]#[repr(u8)]enum StateInner { /// No waiters are queued, and there is no pending notification. /// Waiting while the queue is in this state will enqueue the waiter; /// notifying while in this state will store a pending notification in the /// queue, transitioning to [`StateInner::Woken`]. Empty = 0b00, /// There are one or more waiters in the queue. Waiting while /// the queue is in this state will not transition the state. Waking while /// in this state will wake the first waiter in the queue; if this empties /// the queue, then the queue will transition to [`StateInner::Empty`]. Waiting = 0b01, /// The queue has a stored notification. Waiting while the queue /// is in this state will consume the pending notification *without* /// enqueueing the waiter and transition the queue to [`StateInner::Empty`]. /// Waking while in this state will leave the queue in this state. Woken = 0b10, /// The queue is closed. Waiting while in this state will return /// [`Closed`] without transitioning the queue's state. /// /// *Note*: This *must* correspond to all state bits being set, as it's set /// via a [`fetch_or`]. /// /// [`fetch_or`]: AtomicUsize::fetch_or Closed = 0b11,}
#[pin_project(PinnedDrop)]#[derive(Debug)]#[must_use = "futures do nothing unless `.await`ed or `poll`ed"]pub struct Wait<'a> { // The [`WaitQueue`] being waited on. queue: &'a WaitQueue, // Entry in the wait queue linked list. #[pin] waiter: Waiter,}
/// Future returned from [`WaitQueue::wait_owned()`].////// This is identical to the [`Wait`] future, except that it takes an/// [`Arc`] reference to the [`WaitQueue`], allowing the returned future to/// live for the `'static` lifetime.////// This future is fused, so once it has completed, any future calls to poll/// will immediately return [`Poll::Ready`].////// # Notes////// This future is `!Unpin`, as it is unsafe to [`core::mem::forget`] a/// `WaitOwned` future once it has been polled. For instance, the following/// code must not compile://////```compile_fail/// use maitake_sync::wait_queue::WaitOwned;////// // Calls to this function should only compile if `T` is `Unpin`./// fn assert_unpin<T: Unpin>() {}////// assert_unpin::<WaitOwned<'_>>();/// ```#[pin_project(PinnedDrop)]#[derive(Debug)]pub struct WaitOwned { // The `WaitQueue` being waited on. queue: Arc<WaitQueue>, // Entry in the wait queue. #[pin] waiter: Waiter,}
/// A waiter node which may be linked into a wait queue.#[pin_project]#[repr(C)]struct Waiter { // The intrusive linked list node. // // This *must* be the first field in the struct in order for the `Linked` // implementation to be sound. #[pin] node: UnsafeCell<WaiterInner>, // The future's state. state: WaitState,}
struct WaiterInner { /// Intrusive linked list pointers. links: list::Links<Waiter>, /// The node's waker waker: Wakeup, // This type is !Unpin due to the heuristic from: // <https://github.com/rust-lang/rust/pull/82834> _pin: PhantomPinned,}
bitfield! { #[derive(Eq, PartialEq)] struct WaitState<usize> { /// The waiter's state. const INNER: WaitStateInner; /// The number of times [`WaitQueue::wake_all`] has been called. const WAKE_ALLS = ..; }}
enum_from_bits! { /// The state of a [`Waiter`] node in a [`WaitQueue`]. #[derive(Debug, Eq, PartialEq)] enum WaitStateInner<u8> { /// The waiter has not yet been enqueued. /// /// The number of times [`WaitQueue::wake_all`] has been called is stored /// when the node is created, in order to determine whether it was woken by /// a stored wakeup when enqueueing. /// /// When in this state, the node is **not** part of the linked list, and /// can be dropped without removing it from the list. Start = 0b00,
/// The waiter is waiting. /// /// When in this state, the node **is** part of the linked list. If the /// node is dropped in this state, it **must** be removed from the list /// before dropping it. Failure to ensure this will result in dangling /// pointers in the linked list! Waiting = 0b01,
/// The waiter has been woken. /// /// When in this state, the node is **not** part of the linked list, and /// can be dropped without removing it from the list. Woken = 0b10, }}
#[derive(Clone, Debug)]enum Wakeup { Empty, Waiting(Waker), One, All, // Closed,}
// === impl WaitQueue ===
impl Default for WaitQueue { fn default() -> Self { Self::new() }}
impl WaitQueue { loom_const_fn! { pub const fn new() -> Self { Self { state: CachePadded(AtomicUsize::new(StateInner::Empty.into_usize())), queue: Mutex::new(List::new()), } } }
/// Wake the next task in the queue. /// /// If the queue is empty, a wakeup is stored in the `WaitQueue`, and the /// **next** call to [`wait().await`] will complete immediately. If one or more /// tasks are currently in the queue, the first task in the queue is woken. /// /// At most one wakeup will be stored in the queue at any time. If `wake()` /// is called many times while there are no tasks in the queue, only a /// single wakeup is stored. /// /// [`wait().await`]: Self::wait() #[inline] pub fn wake(&self) { // snapshot the queue's current state. let mut state = self.load();
// check if any tasks are currently waiting on this queue. if there are // no waiting tasks, store the wakeup to be consumed by the next call to // `wait`. loop { match state.get(State::INNER) { // if the queue is closed, bail. StateInner::Closed => return, // if there are waiting tasks, break out of the loop and wake one. StateInner::Waiting => break, _ => {} }
let next = state.with_inner(StateInner::Woken); // advance the state to `Woken`, and return (if we did so // successfully) match self.compare_exchange(state, next) { Ok(_) => return, Err(actual) => state = actual, } }
// okay, there are tasks waiting on the queue; we must acquire the lock // on the linked list and wake the next task from the queue. let mut queue = self.queue.lock();
// the queue's state may have changed while we were waiting to acquire // the lock, so we need to acquire a new snapshot before we take the // waker. state = self.load(); let waker = self.wake_locked(&mut queue, state); drop(queue);
//now that we've released the lock, wake the waiting task (if we //actually deuqueued one). if let Some(waker) = waker { waker.wake(); } }
/// Wake *all* tasks currently in the queue. /// /// All tasks currently waiting on the queue are woken. Unlike [`wake()`], a /// wakeup is *not* stored in the queue to wake the next call to [`wait()`] /// if the queue is empty. Instead, this method only wakes all currently /// registered waiters. Registering a task to be woken is done by `await`ing /// the [`Future`] returned by the [`wait()`] method on this queue. /// /// [`wake()`]: Self::wake /// [`wait()`]: Self::wait #[expect(clippy::missing_panics_doc, reason = "internal assertion")] pub fn wake_all(&self) -> usize { let mut batch = WakeBatch::new(); let mut waiters_remaining = true;
let mut queue = self.queue.lock(); let state = self.load();
match state.get(State::INNER) { StateInner::Woken | StateInner::Empty => { self.state .0 .fetch_add(State::ONE_WAKE_ALL, Ordering::SeqCst); return 0; } StateInner::Closed => return 0, StateInner::Waiting => { let next_state = State::new() .with_inner(StateInner::Empty) .with(State::WAKE_ALLS, state.get(State::WAKE_ALLS) + 1); self.compare_exchange(state, next_state) .expect("state should not have transitioned while locked"); } }
// As long as there are waiters remaining to wake, lock the queue, drain // another batch, release the lock, and wake them. let mut woken = 0; while waiters_remaining { waiters_remaining = Self::drain_to_wake_batch(&mut batch, &mut queue, Wakeup::All); woken += batch.len(); MutexGuard::unlocked(&mut queue, || batch.wake_all()); }
woken }
/// Wait to be woken up by this queue. /// /// Equivalent to: /// /// ```ignore /// async fn wait(&self); /// ``` /// /// This returns a [`Wait`] [`Future`] that will complete when the task is /// woken by a call to [`wake()`] or [`wake_all()`], or when the `WaitQueue` /// is dropped. /// /// Each `WaitQueue` holds a single wakeup. If [`wake()`] was previously /// called while no tasks were waiting on the queue, then `wait().await` /// will complete immediately, consuming the stored wakeup. Otherwise, /// `wait().await` waits to be woken by the next call to [`wake()`] or /// [`wake_all()`]. /// /// The [`Wait`] future is not guaranteed to receive wakeups from calls to /// [`wake()`] if it has not yet been polled. See the documentation for the /// [`Wait::subscribe()`] method for details on receiving wakeups from the /// queue prior to polling the `Wait` future for the first time. /// /// A `Wait` future **is** is guaranteed to receive wakeups from calls to /// [`wake_all()`] as soon as it is created, even if it has not yet been /// polled. /// /// # Returns /// /// The [`Future`] returned by this method completes with one of the /// following [outputs](Future::Output): /// /// - [`Ok`]`(())` if the task was woken by a call to [`wake()`] or /// [`wake_all()`]. /// - [`Err`]`(`[`Closed`]`)` if the task was woken by the `WaitQueue` being /// [`close`d](WaitQueue::close). /// /// # Cancellation /// /// A `WaitQueue` fairly distributes wakeups to waiting tasks in the order /// that they started to wait. If a [`Wait`] future is dropped, the task /// will forfeit its position in the queue. /// /// [`wake()`]: Self::wake /// [`wake_all()`]: Self::wake_all pub fn wait(&self) -> Wait<'_> { Wait { queue: self, waiter: self.waiter(), } }
/// Wait to be woken up by this queue, returning a future that's valid /// for the `'static` lifetime. /// /// This returns a [`WaitOwned`] future that will complete when the task /// is woken by a call to [`wake()`] or [`wake_all()`], or when the /// `WaitQueue` is [closed]. /// /// This is identical to the [`wait()`] method, except that it takes a /// [`Arc`] reference to the [`WaitQueue`], allowing the returned future /// to live for the `'static` lifetime. See the documentation for /// [`wait()`] for details on how to use the future returned by this /// method. /// /// # Returns /// /// The [`Future`] returned by this method completes with one of the /// following [outputs](Future::Output): /// /// - [`Ok`]`(())` if the task was woken by a call to [`wake()`] or /// [`wake_all()`]. /// - [`Err`]`(`[`Closed`]`)` if the task was woken by the `WaitQueue` /// being [closed]. /// /// # Cancellation /// /// A `WaitQueue` fairly distributes wakeups to waiting tasks in the /// order that they started to wait. If a [`WaitOwned`] future is /// dropped, the task will forfeit its position in the queue. /// /// [`wake()`]: Self::wake /// [`wake_all()`]: Self::wake_all /// [`wait()`]: Self::wait /// [closed]: Self::close pub fn wait_owned(self: &Arc<Self>) -> WaitOwned { let waiter = self.waiter(); let queue = self.clone(); WaitOwned { queue, waiter } }
/// Asynchronously poll the given function `f` until a condition occurs, /// using the [`WaitQueue`] to only re-poll when notified. /// /// This can be used to implement a "wait loop", turning a "try" function /// (e.g. "try_recv" or "try_send") into an asynchronous function (e.g. /// "recv" or "send"). /// /// In particular, this function correctly *registers* interest in the [`WaitQueue`] /// prior to polling the function, ensuring that there is not a chance of a race /// where the condition occurs AFTER checking but BEFORE registering interest /// in the [`WaitQueue`], which could lead to deadlock. /// /// This is intended to have similar behavior to `Condvar` in the standard library, /// but asynchronous, and not requiring operating system intervention (or existence). /// /// In particular, this can be used in cases where interrupts or events are used /// to signify readiness or completion of some task, such as the completion of a /// DMA transfer, or reception of an ethernet frame. In cases like this, the interrupt /// can wake the queue, allowing the polling function to check status fields for /// partial progress or completion. /// /// Consider using [`Self::wait_for_value()`] if your function does return a value. /// /// Consider using [`WaitCell::wait_for()`](super::wait_cell::WaitCell::wait_for) /// if you do not need multiple waiters. /// /// # Errors /// /// Returns [`Err`]`(`[`Closed`]`)` if the [`WaitQueue`] is closed. /// pub async fn wait_for<F: FnMut() -> bool>(&self, mut f: F) -> Result<(), Closed> { loop { let wait = self.wait(); let mut wait = core::pin::pin!(wait); let _ = wait.as_mut().subscribe()?; if f() { return Ok(()); } wait.await?; } }
/// Asynchronously poll the given function `f` until a condition occurs, /// using the [`WaitQueue`] to only re-poll when notified. /// /// This can be used to implement a "wait loop", turning a "try" function /// (e.g. "try_recv" or "try_send") into an asynchronous function (e.g. /// "recv" or "send"). /// /// In particular, this function correctly *registers* interest in the [`WaitQueue`] /// prior to polling the function, ensuring that there is not a chance of a race /// where the condition occurs AFTER checking but BEFORE registering interest /// in the [`WaitQueue`], which could lead to deadlock. /// /// This is intended to have similar behavior to `Condvar` in the standard library, /// but asynchronous, and not requiring operating system intervention (or existence). /// /// In particular, this can be used in cases where interrupts or events are used /// to signify readiness or completion of some task, such as the completion of a /// DMA transfer, or reception of an ethernet frame. In cases like this, the interrupt /// can wake the queue, allowing the polling function to check status fields for /// partial progress or completion, and also return the status flags at the same time. /// /// Consider using [`Self::wait_for()`] if your function does not return a value. /// /// Consider using [`WaitCell::wait_for_value()`](super::wait_cell::WaitCell::wait_for_value) /// if you do not need multiple waiters. /// /// # Errors /// /// Returns [`Err`]`(`[`Closed`]`)` if the [`WaitQueue`] is closed. /// pub async fn wait_for_value<T, F: FnMut() -> Option<T>>(&self, mut f: F) -> Result<T, Closed> { loop { let wait = self.wait(); let mut wait = core::pin::pin!(wait); match wait.as_mut().subscribe() { Poll::Ready(wr) => wr?, Poll::Pending => {} } if let Some(t) = f() { return Ok(t); } wait.await?; } }
/// Close the queue, indicating that it may no longer be used. /// /// Once a queue is closed, all [`wait()`] calls (current or future) will /// return an error. /// /// This method is generally used when implementing higher-level /// synchronization primitives or resources: when an event makes a resource /// permanently unavailable, the queue can be closed. /// /// [`wait()`]: Self::wait pub fn close(&self) { let state = self .state .0 .fetch_or(StateInner::Closed.into_bits(), Ordering::SeqCst); let state = State::from_bits(state); if state.get(State::INNER) != StateInner::Waiting { return; }
let mut queue = self.queue.lock(); let mut batch = WakeBatch::new(); let mut waiters_remaining = true;
while waiters_remaining { waiters_remaining = Self::drain_to_wake_batch(&mut batch, &mut queue, Wakeup::All); MutexGuard::unlocked(&mut queue, || batch.wake_all()); } }
/// Returns `true` if this `WaitQueue` is [closed](Self::close). #[must_use] pub fn is_closed(&self) -> bool { self.load().get(State::INNER) == StateInner::Closed }
/// Returns a [`Waiter`] entry in this queue. /// /// This is factored out into a separate function because it's used by both /// [`WaitQueue::wait`] and [`WaitQueue::wait_owned`]. fn waiter(&self) -> Waiter { // how many times has `wake_all` been called when this waiter is created? let current_wake_alls = self.load().get(State::WAKE_ALLS); let state = WaitState::new() .with(WaitState::WAKE_ALLS, current_wake_alls) .with(WaitState::INNER, WaitStateInner::Start); Waiter { state, node: UnsafeCell::new(WaiterInner { links: list::Links::new(), waker: Wakeup::Empty, _pin: PhantomPinned, }), } }
// fn try_wait(&self) -> Poll<Result<(), Closed>> { // let mut state = self.load(); // let initial_wake_alls = state.get(State::WAKE_ALLS); // while state.get(State::INNER) == StateInner::Woken { // match self.compare_exchange(state, state.with_inner(StateInner::Empty)) { // Ok(_) => return Poll::Ready(Ok(())), // Err(actual) => state = actual, // } // } // // match state.get(State::INNER) { // StateInner::Closed => Poll::Ready(Err(Closed(()))), // _ if state.get(State::WAKE_ALLS) > initial_wake_alls => Poll::Ready(Ok(())), // StateInner::Empty | StateInner::Waiting => Poll::Pending, // StateInner::Woken => Poll::Ready(Ok(())), // } // }
#[cold] #[inline(never)] fn wake_locked(&self, queue: &mut List<Waiter>, curr: State) -> Option<Waker> { let inner = curr.get(State::INNER);
// is the queue still in the `Waiting` state? it is possible that we // transitioned to a different state while locking the queue. if inner != StateInner::Waiting { // if there are no longer any queued tasks, try to store the // wakeup in the queue and bail. if let Err(actual) = self.compare_exchange(curr, curr.with_inner(StateInner::Woken)) { debug_assert!(actual.get(State::INNER) != StateInner::Waiting); self.store(actual.with_inner(StateInner::Woken)); }
return None; }
// otherwise, we have to dequeue a task and wake it. let node = queue .pop_back() .expect("if we are in the Waiting state, there must be waiters in the queue"); let waker = Waiter::wake(node, queue, Wakeup::One);
// if we took the final waiter currently in the queue, transition to the // `Empty` state. if queue.is_empty() { self.store(curr.with_inner(StateInner::Empty)); }
waker }
/// Drain waiters from `queue` and add them to `batch`. Returns `true` if /// the batch was filled while more waiters remain in the queue, indicating /// that this function must be called again to wake all waiters. fn drain_to_wake_batch( batch: &mut WakeBatch, queue: &mut List<Waiter>, wakeup: Wakeup, ) -> bool { while let Some(node) = queue.pop_back() { let Some(waker) = Waiter::wake(node, queue, wakeup.clone()) else { // this waiter was enqueued by `Wait::register` and doesn't have // a waker, just keep going. continue; };
if batch.add_waker(waker) { // there's still room in the wake set, just keep adding to it. continue; }
// wake set is full, drop the lock and wake everyone! break; }
!queue.is_empty() }
fn load(&self) -> State { State::from_bits(self.state.0.load(Ordering::SeqCst)) }
fn store(&self, state: State) { self.state.0.store(state.0, Ordering::SeqCst); }
fn compare_exchange(&self, current: State, new: State) -> Result<State, State> { self.state .0 .compare_exchange(current.0, new.0, Ordering::SeqCst, Ordering::SeqCst) .map(State::from_bits) .map_err(State::from_bits) }}
// === impl State & StateInner ===
impl State { const ONE_WAKE_ALL: usize = Self::WAKE_ALLS.first_bit();
fn with_inner(self, inner: StateInner) -> Self { self.with(Self::INNER, inner) }}impl StateInner { const fn into_usize(self) -> usize { self as u8 as usize }}
impl FromBits<usize> for StateInner { const BITS: u32 = 2; type Error = core::convert::Infallible;
fn try_from_bits(bits: usize) -> Result<Self, Self::Error> { #[expect( clippy::cast_possible_truncation, reason = "`bits` will only have 3 bits set" )] Ok(match bits as u8 { bits if bits == Self::Empty as u8 => Self::Empty, bits if bits == Self::Waiting as u8 => Self::Waiting, bits if bits == Self::Woken as u8 => Self::Woken, bits if bits == Self::Closed as u8 => Self::Closed, _ => { unreachable!(); } }) }
fn into_bits(self) -> usize { self as usize }}
// === impl Waiter ===
impl fmt::Debug for Waiter { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Waiter") .field("state", &self.state) .finish_non_exhaustive() }}
impl Waiter { /// Returns the [`Waker`] for the task that owns this `Waiter`. /// /// # Safety /// /// This is only safe to call while the list is locked. The `list` /// parameter ensures this method is only called while holding the lock, so /// this can be safe. /// /// Of course, that must be the *same* list that this waiter is a member of, /// and currently, there is no way to ensure that... #[inline(always)] fn wake(this: NonNull<Self>, list: &mut List<Self>, wakeup: Wakeup) -> Option<Waker> { Waiter::with_inner(this, list, |node| { let waker = mem::replace(&mut node.waker, wakeup); match waker { // the node has a registered waker, so wake the task. Wakeup::Waiting(waker) => Some(waker), // do nothing: the node was registered by `Wait::register` // without a waker, so the future will already be woken when it is // actually polled. Wakeup::Empty => None, // the node was already woken? this should not happen and // probably indicates a race! _ => unreachable!("tried to wake a waiter in the {:?} state!", waker), } }) }
/// # Safety /// /// This is only safe to call while the list is locked. The dummy `_list` /// parameter ensures this method is only called while holding the lock, so /// this can be safe. /// /// Of course, that must be the *same* list that this waiter is a member of, /// and currently, there is no way to ensure that... #[inline(always)] fn with_inner<T>( mut this: NonNull<Self>, _list: &mut List<Self>, f: impl FnOnce(&mut WaiterInner) -> T, ) -> T { // safety: this is only called while holding the lock on the queue, // so it's safe to mutate the waiter. unsafe { f(&mut *this.as_mut().node.get()) } }
fn poll_wait( mut self: Pin<&mut Self>, queue: &WaitQueue, waker: Option<&Waker>, ) -> Poll<Result<(), Closed>> { // Safety: we never move out of `ptr` below, only mutate its fields let ptr = unsafe { NonNull::from(Pin::into_inner_unchecked(self.as_mut())) }; let mut this = self.as_mut().project();
match this.state.get(WaitState::INNER) { WaitStateInner::Start => { let queue_state = queue.load();
// can we consume a pending wakeup? if queue .compare_exchange( queue_state.with_inner(StateInner::Woken), queue_state.with_inner(StateInner::Empty), ) .is_ok() { this.state.set(WaitState::INNER, WaitStateInner::Woken); return Poll::Ready(Ok(())); }
// okay, no pending wakeups. try to wait...
let mut waiters = queue.queue.lock(); let mut queue_state = queue.load();
// the whole queue was woken while we were trying to acquire // the lock! if queue_state.get(State::WAKE_ALLS) != this.state.get(WaitState::WAKE_ALLS) { this.state.set(WaitState::INNER, WaitStateInner::Woken); return Poll::Ready(Ok(())); }
// transition the queue to the waiting state 'to_waiting: loop { match queue_state.get(State::INNER) { // the queue is `Empty`, transition to `Waiting` StateInner::Empty => { match queue.compare_exchange( queue_state, queue_state.with_inner(StateInner::Waiting), ) { Ok(_) => break 'to_waiting, Err(actual) => queue_state = actual, } } // the queue is already `Waiting` StateInner::Waiting => break 'to_waiting, // the queue was woken, consume the wakeup. StateInner::Woken => { match queue.compare_exchange( queue_state, queue_state.with_inner(StateInner::Empty), ) { Ok(_) => { this.state.set(WaitState::INNER, WaitStateInner::Woken); return Poll::Ready(Ok(())); } Err(actual) => queue_state = actual, } } StateInner::Closed => return Poll::Ready(Err(Closed(()))), } }
// enqueue the node this.state.set(WaitState::INNER, WaitStateInner::Waiting); if let Some(waker) = waker { // safety: we may mutate the inner state because we are // holding the lock. unsafe { let node = this.node.as_mut().get(); debug_assert!(matches!((*node).waker, Wakeup::Empty)); (*node).waker = Wakeup::Waiting(waker.clone()); } }
waiters.push_front(ptr);
Poll::Pending } WaitStateInner::Waiting => { let _waiters = queue.queue.lock(); // safety: we may mutate the inner state because we are // holding the lock. unsafe { let node = &mut *this.node.get(); match node.waker { Wakeup::Waiting(ref mut curr_waker) => { match waker { Some(waker) if !curr_waker.will_wake(waker) => { *curr_waker = waker.clone(); } _ => {} } Poll::Pending } Wakeup::All | Wakeup::One => { this.state.set(WaitState::INNER, WaitStateInner::Woken); Poll::Ready(Ok(())) } // Wakeup::Closed => { // this.state.set(WaitState::INNER, WaitStateInner::Woken); // Poll::Ready(Err(Closed(()))) // } Wakeup::Empty => { if let Some(waker) = waker { node.waker = Wakeup::Waiting(waker.clone()); }
Poll::Pending } } } } WaitStateInner::Woken => Poll::Ready(Ok(())), } }
/// Release this `Waiter` from the queue. /// /// This is called from the `drop` implementation for the [`Wait`] and /// [`WaitOwned`] futures. fn release(mut self: Pin<&mut Self>, queue: &WaitQueue) { let state = *(self.as_mut().project().state); // Safety: we never move out of `ptr` below, only mutate its fields let ptr = NonNull::from(unsafe { Pin::into_inner_unchecked(self) });
// if we're not enqueued, we don't have to do anything else. if state.get(WaitState::INNER) != WaitStateInner::Waiting { return; }
let mut waiters = queue.queue.lock(); let state = queue.load(); // remove the node // safety: we have the lock on the queue, so this is safe. unsafe { waiters.remove(ptr); };
// if we removed the last waiter from the queue, transition the state to // `Empty`. if waiters.is_empty() && state.get(State::INNER) == StateInner::Waiting { queue.store(state.with_inner(StateInner::Empty)); }
// if the node has an unconsumed wakeup, it must be assigned to the next // node in the queue. let next_waiter = if Waiter::with_inner(ptr, &mut waiters, |node| matches!(&node.waker, Wakeup::One)) { queue.wake_locked(&mut waiters, state) } else { None };
drop(waiters);
if let Some(next) = next_waiter { next.wake(); } }}
// Safety: TODOunsafe impl Linked<list::Links<Self>> for Waiter { type Handle = NonNull<Waiter>;
fn into_ptr(r: Self::Handle) -> NonNull<Self> { r }
unsafe fn from_ptr(ptr: NonNull<Self>) -> Self::Handle { ptr }
unsafe fn links(target: NonNull<Self>) -> NonNull<list::Links<Waiter>> { // Safety: ensured by caller unsafe { // Safety: using `ptr::addr_of!` avoids creating a temporary // reference, which stacked borrows dislikes. let node = &*ptr::addr_of!((*target.as_ptr()).node); let links = ptr::addr_of_mut!((*node.get()).links); // Safety: since the `target` pointer is `NonNull`, we can assume // that pointers to its members are also not null, making this use // of `new_unchecked` fine. NonNull::new_unchecked(links) } }}
// === impl Wait ===
impl Wait<'_> { /// Returns `true` if this `WaitOwned` future is waiting for a /// notification from the provided [`WaitQueue`]. #[inline] #[must_use] pub fn waits_on(&self, queue: &WaitQueue) -> bool { ptr::eq(self.queue, queue) }
/// Returns `true` if `self` and `other` are waiting on a notification /// from the same [`WaitQueue`]. #[inline] #[must_use] pub fn same_queue(&self, other: &Wait<'_>) -> bool { ptr::eq(self.queue, other.queue) }
/// Eagerly subscribe this future to wakeups from [`WaitQueue::wake()`]. /// /// Polling a `Wait` future adds that future to the list of waiters that may /// receive a wakeup from a `WaitQueue`. However, in some cases, it is /// desirable to subscribe to wakeups *prior* to actually waiting for one. /// This method should be used when it is necessary to ensure a `Wait` /// future is in the list of waiters before the future is `poll`ed for the /// first time. /// /// In general, this method is used in cases where a [`WaitQueue`] must /// synchronize with some additional state, such as an `AtomicBool` or /// counter. If a task first checks that state, and then chooses whether or /// not to wait on the `WaitQueue` based on that state, then a race /// condition may occur where the `WaitQueue` wakes waiters *between* when /// the task checked the external state and when it first polled its `Wait` /// future to wait on the queue. This method allows registering the `Wait` /// future with the queue *prior* to checking the external state, without /// actually sleeping, so that when the task does wait for the `Wait` future /// to complete, it will have received any wakeup that was sent between when /// the external state was checked and the `Wait` future was first polled. /// /// # Returns /// /// This method returns a [`Poll`]`<`[`Result<(), Closed>`]`>` which is `Ready` a wakeup was /// already received. This method returns [`Poll::Ready`] in the following /// cases: /// /// 1. The [`WaitQueue::wake()`] method was called between the creation of the /// `Wait` and the call to this method. /// 2. This is the first call to `subscribe` or `poll` on this future, and the /// `WaitQueue` was holding a stored wakeup from a previous call to /// [`wake()`]. This method consumes the wakeup in that case. /// 3. The future has previously been `subscribe`d or polled, and it has since /// then been marked ready by either consuming a wakeup from the /// `WaitQueue`, or by a call to [`wake()`] or [`wake_all()`] that /// removed it from the list of futures ready to receive wakeups. /// 4. The `WaitQueue` has been [`close`d](WaitQueue::close), in which case /// this method returns `Poll::Ready(Err(Closed))`. /// /// If this method returns [`Poll::Ready`], any subsequent `poll`s of this /// `Wait` future will also immediately return [`Poll::Ready`]. /// /// If the [`Wait`] future subscribed to wakeups from the queue, and /// has not been woken, this method returns [`Poll::Pending`]. /// /// [`wake()`]: WaitQueue::wake /// [`wake_all()`]: WaitQueue::wake_all pub fn subscribe(self: Pin<&mut Self>) -> Poll<Result<(), Closed>> { let this = self.project(); this.waiter.poll_wait(this.queue, None) }}
impl Future for Wait<'_> { type Output = Result<(), Closed>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { let this = self.project(); this.waiter.poll_wait(this.queue, Some(cx.waker())) }}
#[pinned_drop]impl PinnedDrop for Wait<'_> { fn drop(mut self: Pin<&mut Self>) { let this = self.project(); this.waiter.release(this.queue); }}
// === impl WaitOwned ===
impl WaitOwned { /// Returns `true` if this `WaitOwned` future is waiting for a /// notification from the provided [`WaitQueue`]. #[inline] #[must_use] pub fn waits_on(&self, queue: &WaitQueue) -> bool { ptr::eq(&*self.queue, queue) }
/// Returns `true` if `self` and `other` are waiting on a notification /// from the same [`WaitQueue`]. #[inline] #[must_use] pub fn same_queue(&self, other: &Wait<'_>) -> bool { ptr::eq(&*self.queue, other.queue) }
/// Eagerly subscribe this future to wakeups from [`WaitQueue::wake()`]. /// /// Polling a `WaitOwned` future adds that future to the list of waiters /// that may receive a wakeup from a `WaitQueue`. However, in some /// cases, it is desirable to subscribe to wakeups *prior* to actually /// waiting for one. This method should be used when it is necessary to /// ensure a `WaitOwned` future is in the list of waiters before the /// future is `poll`ed for the rst time. /// /// In general, this method is used in cases where a [`WaitQueue`] must /// synchronize with some additional state, such as an `AtomicBool` or /// counter. If a task first checks that state, and then chooses whether or /// not to wait on the `WaitQueue` based on that state, then a race /// condition may occur where the `WaitQueue` wakes waiters *between* when /// the task checked the external state and when it first polled its /// `WaitOwned` future to wait on the queue. This method allows /// registering the `WaitOwned` future with the queue *prior* to /// checking the external state, without actually sleeping, so that when /// the task does wait for the `WaitOwned` future to complete, it will /// have received any wakeup that was sent between when the external /// state was checked and the `WaitOwned` future was first polled. /// /// # Returns /// /// This method returns a [`Poll`]`<`[`Result<(), Closed>`]`>` which is `Ready` /// a wakeup was already received. This method returns [`Poll::Ready`] /// in the following cases: /// /// 1. The [`WaitQueue::wake()`] method was called between the creation /// of the `WaitOwned` future and the call to this method. /// 2. This is the first call to `subscribe` or `poll` on this future, /// and the `WaitQueue` was holding a stored wakeup from a previous /// call to [`wake()`]. This method consumes the wakeup in that case. /// 3. The future has previously been `subscribe`d or polled, and it /// has since then been marked ready by either consuming a wakeup /// from the `WaitQueue`, or by a call to [`wake()`] or /// [`wake_all()`] that removed it from the list of futures ready to /// receive wakeups. /// 4. The `WaitQueue` has been [`close`d](WaitQueue::close), in which /// case this method returns `Poll::Ready(Err(Closed))`. /// /// If this method returns [`Poll::Ready`], any subsequent `poll`s of /// this `Wait` future will also immediately return [`Poll::Ready`]. /// /// If the [`WaitOwned`] future subscribed to wakeups from the queue, /// and has not been woken, this method returns [`Poll::Pending`]. /// /// [`wake()`]: WaitQueue::wake /// [`wake_all()`]: WaitQueue::wake_all pub fn subscribe(self: Pin<&mut Self>) -> Poll<Result<(), Closed>> { let this = self.project(); this.waiter.poll_wait(this.queue, None) }}
impl Future for WaitOwned { type Output = Result<(), Closed>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { let this = self.project(); this.waiter.poll_wait(&*this.queue, Some(cx.waker())) }}
#[pinned_drop]impl PinnedDrop for WaitOwned { fn drop(mut self: Pin<&mut Self>) { let this = self.project(); this.waiter.release(&*this.queue); }}