diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index fd85871..f0f9ba9 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -43,6 +43,29 @@ jobs: - if: steps.version.outputs.changed == 'true' run: cargo check --no-default-features + # Checked before the upload rather than during it. This crate has no + # repository-level `CARGO_REGISTRY_TOKEN`, and every publish run so far + # short-circuited on "version unchanged", so the upload path has never + # actually run here. The first real release would be the first time the + # token is exercised, and an absent one surfaces as an authentication + # error from `cargo publish` after the package is built, which reads like + # a registry problem rather than a missing secret. + - if: steps.version.outputs.changed == 'true' + run: | + if [ -z "${CARGO_REGISTRY_TOKEN:-}" ]; then + echo "CARGO_REGISTRY_TOKEN is empty." + echo + echo "Set it on pathscale/nagoya, or at the organisation level with this" + echo "repository in scope:" + echo " gh secret set CARGO_REGISTRY_TOKEN --repo pathscale/nagoya" + echo + echo "Nothing downstream resolves until this crate publishes: worktable" + echo "requires nagoya ^0.1.1, and nagoya requires ps-st3 ^0.6.1." + exit 1 + fi + env: + CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} + - if: steps.version.outputs.changed == 'true' run: cargo publish --token "$CARGO_REGISTRY_TOKEN" env: diff --git a/Cargo.toml b/Cargo.toml index 4fdfdd3..ee102e6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "nagoya" -version = "0.1.0" +version = "0.1.1" edition = "2021" rust-version = "1.85" license = "MIT OR Apache-2.0" @@ -14,7 +14,7 @@ categories = ["asynchronous", "concurrency", "no-std"] default = ["std"] # Off, this crate does not link `std`. What it costs: `Runtime`, which owns # threads, and the parker `block_on` uses when there is one to park. -std = ["ps-st3/host", "dep:futures-io", "futures-io/std", "dep:futures-util", "futures-util/std", "futures-util/io"] +std = ["ps-st3/host", "async-task/std", "dep:futures-io", "futures-io/std", "dep:futures-util", "futures-util/std", "futures-util/io"] [dependencies] diff --git a/src/io.rs b/src/io.rs index 4e17cde..2d0e2ad 100644 --- a/src/io.rs +++ b/src/io.rs @@ -79,9 +79,11 @@ pub trait Read { async move { let mut filled = 0; while filled < buffer.len() { - match self.read(&mut buffer[filled..]).await? { - 0 => return Err(Error::new(ErrorKind::UnexpectedEof)), - n => filled += n, + match self.read(&mut buffer[filled..]).await { + Ok(0) => return Err(Error::new(ErrorKind::UnexpectedEof)), + Ok(n) => filled += n, + Err(error) if error.kind() == ErrorKind::Interrupted => continue, + Err(error) => return Err(error), } } Ok(()) @@ -114,6 +116,9 @@ pub trait Read { return Ok(filled - start); } Ok(n) => buffer.truncate(filled + n), + Err(error) if error.kind() == ErrorKind::Interrupted => { + buffer.truncate(filled); + } Err(error) => { buffer.truncate(filled); return Err(error); @@ -144,11 +149,13 @@ pub trait Write { async move { let mut written = 0; while written < buffer.len() { - match self.write(&buffer[written..]).await? { + match self.write(&buffer[written..]).await { // Nothing taken and no error: the destination will not // accept more, and looping would spin forever. - 0 => return Err(Error::new(ErrorKind::WriteZero)), - n => written += n, + Ok(0) => return Err(Error::new(ErrorKind::WriteZero)), + Ok(n) => written += n, + Err(error) if error.kind() == ErrorKind::Interrupted => continue, + Err(error) => return Err(error), } } Ok(()) diff --git a/src/lib.rs b/src/lib.rs index ae671e8..f724654 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -62,6 +62,13 @@ mod yield_now; pub use block_on::block_on; pub use par::{par_for, Cancel, ParFor}; +/// The idle policy a pool runs with, from the queues underneath. +/// +/// Re-exported because selecting one is a nagoya-level decision and every +/// caller that makes it would otherwise have to name `ps-st3` as a direct +/// dependency to spell the type, for one struct, in a graph that already +/// contains it transitively. +pub use st3::fanout::Tuning; pub use time::{ now_ns, poll as poll_timers, set_clock, sleep, sleep_until, timeout, Elapsed, Sleep, Timeout, }; @@ -69,8 +76,11 @@ pub use yield_now::{yield_now, YieldNow}; /// A handle to a spawned task's output. /// -/// Awaiting it yields `Some(output)`, or `None` if the task was cancelled by -/// dropping the handle before it finished. +/// Awaiting yields `Some(output)`, or `None` if its runnable was canceled +/// because the pool was destroyed before it could run. With `std` and +/// unwinding enabled, a task panic is rethrown when its handle is awaited, +/// rather than terminating the worker. Without `std`, panic handling belongs +/// to the host. Abort-on-panic builds cannot isolate a panic in either mode. /// /// **Dropping this detaches the task; it does not cancel it.** That is Tokio's /// behaviour and it was this crate's before, and it is worth stating because @@ -80,11 +90,11 @@ pub use yield_now::{yield_now, YieldNow}; /// [`cancel`](JoinHandle::cancel) is how to ask for the other thing. pub struct JoinHandle { /// `None` only between `Drop` taking it and the handle going away. - task: Option>, + task: Option>, } impl JoinHandle { - fn task(&mut self) -> &mut async_task::Task { + fn task(&mut self) -> &mut async_task::FallibleTask { self.task .as_mut() .expect("the task is taken only by `Drop`") @@ -122,7 +132,7 @@ impl core::fmt::Debug for JoinHandle { &self .task .as_ref() - .is_some_and(async_task::Task::is_finished), + .is_some_and(async_task::FallibleTask::is_finished), ) .finish() } @@ -140,16 +150,33 @@ impl Future for JoinHandle { type Output = Option; fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - Pin::new(self.get_mut().task()).poll(context).map(Some) + Pin::new(self.get_mut().task()).poll(context) } } +/// Host-supplied identity of the execution lane currently running a pool worker. +/// +/// Implementations must return an identity for the calling lane, not for the +/// lane that originally created the executor. Return `None` off a worker or +/// for a different pool. Invalid worker indices fall back to shared routing. +/// Keep this lookup bounded and nonblocking: it runs on every task wake. +/// Context implementations must not strongly own their pool; use a weak +/// reference or a host-owned identity to avoid a queue ownership cycle. +/// +/// This interface requires only `core`/`alloc`. A no_std host can use its +/// scheduler's execution-local storage without acquiring OS-thread ownership. +pub trait WorkerContext: Send + Sync { + /// The calling lane's worker index for `pool`, if it is currently running it. + fn current_worker(&self, pool: &Pool) -> Option; +} + /// A pool with a place to put the next task. /// /// Work goes to one injector that every worker drains, so there is nothing for /// a spawn to choose: the pool decides who runs it. pub struct Executor { pool: Arc, + context: Option>, } impl Executor { @@ -159,7 +186,32 @@ impl Executor { /// which is what keeps this free of an operating system. #[must_use] pub fn new(pool: Arc) -> Self { - Self { pool } + Self { + pool, + context: None, + } + } + + /// Use host-provided execution-local identity for locality scheduling. + /// + /// Unlike `Runtime`, this does not create threads or require `std`. + #[must_use] + pub fn with_worker_context(pool: Arc, context: Arc) -> Self { + Self { + pool, + context: Some(context), + } + } + + /// Run a caller-owned worker thread with nagoya's thread-local identity. + /// + /// The runner must belong to this executor's pool. Returns `false` if the + /// worker is already running, as `Pool::run` does. For a custom/no_std host, + /// use `with_worker_context` and drive the pool directly. + #[cfg(feature = "std")] + pub fn run_worker(&self, runner: st3::fanout::Runner) -> bool { + let _current = task::mark_current(&self.pool, runner.id()); + self.pool.run(runner) } /// Another handle to the same pool, for a task that spawns. @@ -168,7 +220,10 @@ impl Executor { /// is one: another handle to the same pool. #[must_use] pub fn clone_handle(&self) -> Self { - Self::new(self.pool.clone()) + Self { + pool: self.pool.clone(), + context: self.context.clone(), + } } /// The pool this executor spawns onto. @@ -184,7 +239,7 @@ impl Executor { F::Output: Send + 'static, { JoinHandle { - task: Some(task::spawn(future, self.pool.clone())), + task: Some(task::spawn(future, self.pool.clone(), self.context.clone()).fallible()), } } } diff --git a/src/par.rs b/src/par.rs index 90d11b2..5066170 100644 --- a/src/par.rs +++ b/src/par.rs @@ -27,7 +27,9 @@ //! that it will not be, and neither is worth the unsafe until something asks //! for it. Move what you need into the closure, or put it in an `Arc`. -use alloc::sync::Arc; +#[cfg(feature = "std")] +use alloc::boxed::Box; +use alloc::sync::{Arc, Weak}; use core::future::Future; use core::ops::Range; use core::pin::Pin; @@ -97,6 +99,10 @@ const PIECES_PER_WORKER: usize = 8; /// happens until it is polled: the first poll is what hands the work over, so /// a loop that is created and dropped costs one allocation and no work. /// +/// With `std` and unwinding enabled, a body panic cancels remaining work and +/// is rethrown by the awaiting future after all published pieces retire. +/// Without `std`, panic handling is the host's responsibility. +/// /// ``` /// # use core::sync::atomic::{AtomicUsize, Ordering}; /// # use alloc::sync::Arc; @@ -129,6 +135,8 @@ where outstanding: AtomicUsize::new(1), finished: AtomicBool::new(false), waiter: Mutex::new(None), + #[cfg(feature = "std")] + panic: Mutex::new(None), }) as Arc, pool, range, @@ -147,6 +155,10 @@ trait Split: Send + Sync { fn retire(&self); fn is_finished(&self) -> bool; fn park(&self, waker: &Waker); + #[cfg(feature = "std")] + fn record_panic(&self, payload: Box); + #[cfg(feature = "std")] + fn resume_panic(&self); } struct Shared { @@ -157,6 +169,8 @@ struct Shared { outstanding: AtomicUsize, finished: AtomicBool, waiter: Mutex>, + #[cfg(feature = "std")] + panic: Mutex>>, } impl Split for Shared @@ -169,7 +183,6 @@ where // between two items the body is running and nothing here can interrupt // it. if cancel.is_cancelled() { - self.retire(); return; } let mut range = range; @@ -177,7 +190,7 @@ where // the pool and keeping the left. Recursion here would put the whole // split depth on one worker's stack, which is a stack overflow waiting // for a large enough range. - while range.end - range.start > leaf { + while range.end.saturating_sub(range.start) > leaf { let middle = range.start + (range.end - range.start) / 2; let right = middle..range.end; range = range.start..middle; @@ -186,15 +199,18 @@ where // the piece finish and retire against a count that has not been // raised yet, which would complete the loop early. self.outstanding.fetch_add(1, Ordering::AcqRel); - let half = self.clone(); - let handle = pool.clone(); - let switch = cancel.clone(); - pool.submit_fn(move || half.work(&handle, right, leaf, &switch)); + let piece = Piece { + shared: self.clone(), + pool: Arc::downgrade(pool), + range: right, + leaf, + cancel: cancel.clone(), + }; + pool.submit_fn(move || piece.run()); } for index in range { (self.body)(index); } - self.retire(); } fn retire(&self) { @@ -213,7 +229,69 @@ where } fn park(&self, waker: &Waker) { - *self.waiter.lock() = Some(waker.clone()); + let old = self.waiter.lock().replace(waker.clone()); + drop(old); + } + + #[cfg(feature = "std")] + fn record_panic(&self, payload: Box) { + let mut pending = self.panic.lock(); + if pending.is_none() { + *pending = Some(payload); + } else { + drop(pending); + drop(payload); + } + } + + #[cfg(feature = "std")] + fn resume_panic(&self) { + let payload = self.panic.lock().take(); + if let Some(payload) = payload { + std::panic::resume_unwind(payload); + } + } +} + +// One counted piece, including while it is queued. Dropping an unrun closure +// must retire it too; keeping retirement only at the end of work() loses that +// accounting on panic, cancellation, and pool destruction. +struct Piece { + shared: Arc, + pool: Weak, + range: Range, + leaf: usize, + cancel: Cancel, +} + +impl Piece { + fn run(self) { + let Some(pool) = self.pool.upgrade() else { + self.cancel.cancel(); + return; + }; + #[cfg(feature = "std")] + { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.shared + .clone() + .work(&pool, self.range.clone(), self.leaf, &self.cancel); + })); + if let Err(payload) = result { + self.cancel.cancel(); + self.shared.record_panic(payload); + } + } + #[cfg(not(feature = "std"))] + self.shared + .clone() + .work(&pool, self.range.clone(), self.leaf, &self.cancel); + } +} + +impl Drop for Piece { + fn drop(&mut self) { + self.shared.retire(); } } @@ -294,12 +372,21 @@ impl Future for ParFor { if let Some((shared, pool, range, leaf)) = this.state.take() { this.started = Some(shared.clone()); + if range.is_empty() || this.cancel.is_cancelled() { + shared.retire(); + return Poll::Ready(()); + } // Hand the whole range to the pool rather than splitting it here. // Splitting is itself work, and doing it on the pool means it is // parallel too: each piece splits its own half. - let handle = pool.clone(); - let switch = this.cancel.clone(); - pool.submit_fn(move || shared.work(&handle, range, leaf, &switch)); + let piece = Piece { + shared, + pool: Arc::downgrade(&pool), + range, + leaf, + cancel: this.cancel.clone(), + }; + pool.submit_fn(move || piece.run()); } let Some(shared) = this.started.as_ref() else { @@ -309,12 +396,16 @@ impl Future for ParFor { }; if shared.is_finished() { + #[cfg(feature = "std")] + shared.resume_panic(); return Poll::Ready(()); } // Registered before the flag is read again, so a piece that finishes // between the two still finds a waker to call. shared.park(context.waker()); if shared.is_finished() { + #[cfg(feature = "std")] + shared.resume_panic(); return Poll::Ready(()); } Poll::Pending diff --git a/src/runtime.rs b/src/runtime.rs index 35d2194..72ad80b 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -25,7 +25,7 @@ use alloc::sync::Arc; use core::future::Future; -use st3::fanout::{Pool, StdHost}; +use st3::fanout::{Pool, StdHost, Tuning}; use crate::{Executor, JoinHandle}; @@ -49,14 +49,34 @@ impl Runtime { /// returning a runtime that silently runs nothing would be worse. #[must_use] pub fn new(workers: usize) -> Self { + Self::with_tuning(workers, Tuning::default(), "nagoya") + } + + /// A runtime with `workers` threads, at `tuning`, whose threads are named + /// `{label}-{id}`. + /// + /// Worker identity is installed automatically. Callers owning their own + /// threads can instead use [`Executor::run_worker`]; custom/no_std hosts + /// can supply [`crate::WorkerContext`] without creating an owned runtime. + /// + /// # Panics + /// + /// If a thread cannot be started, or `workers` exceeds `usize::BITS`, + /// the underlying pool's worker bitmap capacity. + #[must_use] + pub fn with_tuning(workers: usize, tuning: Tuning, label: &str) -> Self { let workers = workers.max(1); + assert!( + workers <= usize::BITS as usize, + "worker count exceeds the pool bitmap capacity" + ); let host = Arc::new(StdHost::new(workers)); - let pool = Pool::new(workers, 1024, host); + let pool = Pool::with_tuning(workers, 1024, host, tuning); for id in 0..workers { let pool = pool.clone(); let runner = pool.runner(id); std::thread::Builder::new() - .name(alloc::format!("nagoya-{id}")) + .name(alloc::format!("{label}-{id}")) .spawn(move || { // Tells the scheduler that a wake happening on this thread // belongs to worker `id`, so it can skip the injector. See @@ -85,16 +105,28 @@ impl Runtime { pub fn executor(&self) -> &Executor { &self.executor } + + /// The pool underneath, for the APIs that take one directly. + /// + /// [`crate::par_for`] takes an `Arc` because it schedules closures + /// without building async tasks. + #[must_use] + pub fn pool(&self) -> &Arc { + self.executor.pool() + } } /// Threads for the shared pool. /// -/// The machine's parallelism, which is also what `tokio::spawn` gave the -/// callers this replaces. Keeping the number the same is the point: changing -/// the thread count and the runtime in one step makes any measurement of the -/// swap unreadable. +/// Use the machine's parallelism up to the pool's supported bitmap capacity. fn shared_threads() -> usize { - std::thread::available_parallelism().map_or(2, core::num::NonZeroUsize::get) + supported_default_threads( + std::thread::available_parallelism().map_or(2, core::num::NonZeroUsize::get), + ) +} + +fn supported_default_threads(available: usize) -> usize { + available.clamp(1, usize::BITS as usize) } /// One pool per process, started on first use. @@ -128,6 +160,23 @@ mod tests { use crate::block_on; use core::sync::atomic::{AtomicUsize, Ordering}; + #[test] + fn automatic_worker_count_respects_the_bitmap_capacity() { + assert_eq!(supported_default_threads(0), 1); + assert_eq!(supported_default_threads(2), 2); + assert_eq!(supported_default_threads(usize::MAX), usize::BITS as usize); + } + + #[test] + fn an_owned_runtime_installs_worker_identity() { + let runtime = Runtime::with_tuning(1, Tuning::locality(), "identity-test"); + let pool = runtime.pool().clone(); + let task = runtime.spawn(async move { crate::task::current_worker(&pool) }); + let worker = block_on(task); + runtime.pool().shut_down(); + assert_eq!(worker, Some(Some(0))); + } + #[test] fn the_shared_pool_is_one_pool() { assert!(core::ptr::eq(background(), background())); diff --git a/src/sync.rs b/src/sync.rs index 92f5ece..d9b4447 100644 --- a/src/sync.rs +++ b/src/sync.rs @@ -26,11 +26,11 @@ //! depend on it, that is the moment to write the queue properly rather than to //! discover it in production. //! -//! **Not `no_std`-free of allocation.** The waiter lists are `VecDeque`, -//! so this needs `alloc`, which the rest of the crate needs anyway. +//! **Requires `alloc`.** Contended waits use reusable indexed slots. +//! Uncontended acquisition does not allocate a waiter. -use alloc::collections::VecDeque; use alloc::sync::Arc; +use alloc::vec::Vec; use core::cell::UnsafeCell; use core::future::Future; use core::ops::{Deref, DerefMut}; @@ -40,43 +40,203 @@ use core::task::{Context, Poll, Waker}; use spin::Mutex; -/// The waiter list every primitive here shares. -/// -/// A spin lock rather than a blocking one because it is held for the length of -/// a push or a pop and nothing else, and because a blocking mutex would need -/// the very thing this module exists to avoid. -#[derive(Default)] -struct Waiters(Mutex>); +#[derive(Clone, Copy, PartialEq, Eq)] +enum WakeKind { + Queued, + One, + All, +} + +struct Waiter { + previous: Option, + next: Option, + kind: WakeKind, + waker: Option, +} + +enum WaiterSlot { + Live(Waiter), + Free(Option), +} + +// Stable slot indices belong to futures, not wakers. A granted slot remains +// live until its future acknowledges or cancels it. This permits O(1) removal +// without an intrusive pointer, a per-waiter Arc, or a scan on cancellation. +struct WaitQueue { + slots: Vec, + free: Option, + head: Option, + tail: Option, +} + +impl WaitQueue { + const fn new() -> Self { + Self { + slots: Vec::new(), + free: None, + head: None, + tail: None, + } + } + + fn get(&self, id: usize) -> &Waiter { + match &self.slots[id] { + WaiterSlot::Live(waiter) => waiter, + WaiterSlot::Free(_) => unreachable!("registration already released"), + } + } + + fn get_mut(&mut self, id: usize) -> &mut Waiter { + match &mut self.slots[id] { + WaiterSlot::Live(waiter) => waiter, + WaiterSlot::Free(_) => unreachable!("registration already released"), + } + } + + fn enqueue(&mut self, id: usize) { + let tail = self.tail; + let waiter = self.get_mut(id); + waiter.previous = tail; + waiter.next = None; + waiter.kind = WakeKind::Queued; + if let Some(tail) = tail { + self.get_mut(tail).next = Some(id); + } else { + self.head = Some(id); + } + self.tail = Some(id); + } + + fn unlink(&mut self, id: usize) { + let waiter = self.get(id); + let (previous, next) = (waiter.previous, waiter.next); + if let Some(previous) = previous { + self.get_mut(previous).next = next; + } else { + self.head = next; + } + if let Some(next) = next { + self.get_mut(next).previous = previous; + } else { + self.tail = previous; + } + let waiter = self.get_mut(id); + waiter.previous = None; + waiter.next = None; + } + + // Return any replaced waker so its destructor runs outside the queue lock. + fn register(&mut self, token: &mut Option, waker: Option<&Waker>) -> Option { + let id = match *token { + Some(id) => { + if self.get(id).kind != WakeKind::Queued { + self.enqueue(id); + } + id + } + None => { + let waiter = Waiter { + previous: None, + next: None, + kind: WakeKind::Queued, + waker: None, + }; + let id = if let Some(id) = self.free { + self.free = match &self.slots[id] { + WaiterSlot::Free(next) => *next, + WaiterSlot::Live(_) => unreachable!("free-list entry is live"), + }; + self.slots[id] = WaiterSlot::Live(waiter); + id + } else { + self.slots.push(WaiterSlot::Live(waiter)); + self.slots.len() - 1 + }; + *token = Some(id); + self.enqueue(id); + id + } + }; + let waiter = self.get_mut(id); + if let Some(waker) = waker { + if !waiter + .waker + .as_ref() + .is_some_and(|old| old.will_wake(waker)) + { + return waiter.waker.replace(waker.clone()); + } + } + None + } + + fn remove(&mut self, token: &mut Option) -> Option { + let id = token.take()?; + if self.get(id).kind == WakeKind::Queued { + self.unlink(id); + } + let old = core::mem::replace(&mut self.slots[id], WaiterSlot::Free(self.free)); + self.free = Some(id); + match old { + WaiterSlot::Live(waiter) => Some(waiter), + WaiterSlot::Free(_) => unreachable!("registration already released"), + } + } + + // Outer Option means a waiter was granted, even if enable() has not + // supplied a waker yet. + fn grant_one(&mut self) -> Option> { + let id = self.head?; + self.unlink(id); + let waiter = self.get_mut(id); + waiter.kind = WakeKind::One; + Some(waiter.waker.take()) + } + + fn grant_all(&mut self) -> Vec { + let mut wakers = Vec::new(); + while let Some(id) = self.head { + self.unlink(id); + let waiter = self.get_mut(id); + waiter.kind = WakeKind::All; + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + wakers + } +} + +struct Waiters(Mutex); impl Waiters { const fn new() -> Self { - Self(Mutex::new(VecDeque::new())) + Self(Mutex::new(WaitQueue::new())) } - fn push(&self, waker: &Waker) { - let mut queue = self.0.lock(); - // Replacing an equivalent waker rather than pushing a second one: a - // future polled twice before it completes would otherwise leave a stale - // entry that wakes a task which has already moved on. - if queue.iter().any(|w| w.will_wake(waker)) { - return; + fn register(&self, token: &mut Option, waker: &Waker) { + let old = self.0.lock().register(token, Some(waker)); + drop(old); + } + + fn remove(&self, token: &mut Option) -> Option { + if token.is_none() { + return None; } - queue.push_back(waker.clone()); + let removed = self.0.lock().remove(token); + removed.map(|waiter| waiter.kind) } fn wake_one(&self) { - let waker = self.0.lock().pop_front(); + let waker = self.0.lock().grant_one().flatten(); if let Some(waker) = waker { waker.wake(); } } fn wake_all(&self) { - // Drained under the lock, woken outside it: `wake` runs arbitrary - // scheduler code, and holding a spin lock across that invites the - // woken task to spin on the lock its waker still holds. - let drained: VecDeque = core::mem::take(&mut *self.0.lock()); - for waker in drained { + let wakers = self.0.lock().grant_all(); + for waker in wakers { waker.wake(); } } @@ -91,10 +251,8 @@ pub struct Notify { permit: AtomicBool, /// Bumped by [`Notify::notify_waiters`], which leaves no permit. /// - /// A `Notified` snapshots this when it is created, so a broadcast landing - /// between creation and the first poll is still observed. That is the race - /// `tokio`'s `Notified::enable` exists for, and taking it at construction - /// costs one relaxed load instead of an intrusive list. + /// Snapshotted at construction so broadcasts before the first poll are + /// observed. Single notifications are owned by registered waiter slots. generation: AtomicUsize, waiters: Waiters, } @@ -126,14 +284,37 @@ impl Notify { /// Wake one waiter, or leave a permit for the next one to arrive. pub fn notify_one(&self) { - self.permit.store(true, Ordering::Release); - self.waiters.wake_one(); + let waker = { + let mut queue = self.waiters.0.lock(); + self.grant_one(&mut queue) + }; + if let Some(waker) = waker { + waker.wake(); + } + } + + fn grant_one(&self, queue: &mut WaitQueue) -> Option { + match queue.grant_one() { + Some(waker) => waker, + None => { + // Registration and this fallback share a lock: a waiter + // cannot appear between the empty check and permit store. + self.permit.store(true, Ordering::Release); + None + } + } } /// Wake every current waiter. Leaves no permit. pub fn notify_waiters(&self) { - self.generation.fetch_add(1, Ordering::AcqRel); - self.waiters.wake_all(); + let wakers = { + let mut queue = self.waiters.0.lock(); + self.generation.fetch_add(1, Ordering::AcqRel); + queue.grant_all() + }; + for waker in wakers { + waker.wake(); + } } /// Wait for a notification, taking a stored permit if one is there. @@ -141,6 +322,9 @@ impl Notify { Notified { notify: self, generation: self.generation.load(Ordering::Acquire), + registration: None, + done: false, + reserved_one: false, } } } @@ -150,22 +334,18 @@ pub struct Notified<'a> { notify: &'a Notify, /// The broadcast generation when this was created. generation: usize, + registration: Option, + done: bool, + // enable() may reserve a notification before poll delivers it. Dropping + // during that window must transfer the notification, not consume it. + reserved_one: bool, } impl Notified<'_> { - /// Whether a notification is already waiting for this future. - /// - /// **This is not `tokio`'s `enable` and does not need to be.** There, the - /// future must be linked into the waiter list before the caller reads any - /// state, or a notification in between is lost. Here the two ways of - /// notifying are both already durable across that window: `notify_one` - /// leaves a permit that outlives it, and `notify_waiters` bumps a - /// generation this future snapshotted when it was created. So the race the - /// call guards against cannot happen, and the method exists to keep the - /// registration point visible at the call site, which is worth having. + /// Register before checking external state, reserving a notification if + /// one is already available. Returns whether this future is now ready. pub fn enable(self: Pin<&mut Self>) -> bool { - self.notify.permit.load(Ordering::Acquire) - || self.notify.generation.load(Ordering::Acquire) != self.generation + self.get_mut().ready_or_register(None) } } @@ -173,13 +353,9 @@ impl Future for Notified<'_> { type Output = (); fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> { - if self.is_notified() { - return Poll::Ready(()); - } - // Registered before the second look, so a notification landing between - // the two still finds a waker to call. - self.notify.waiters.push(context.waker()); - if self.is_notified() { + let this = self.get_mut(); + if this.ready_or_register(Some(context.waker())) { + this.reserved_one = false; return Poll::Ready(()); } Poll::Pending @@ -187,17 +363,83 @@ impl Future for Notified<'_> { } impl Notified<'_> { - /// Take a permit, or observe that a broadcast has happened since creation. - fn is_notified(&self) -> bool { - if self - .notify - .permit - .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire) - .is_ok() + fn ready_or_register(&mut self, waker: Option<&Waker>) -> bool { + if self.done { + return true; + } + if self.registration.is_none() + && self.notify.generation.load(Ordering::Acquire) != self.generation { + self.done = true; + return true; + } + // The uncontended stored-permit path does not allocate or lock. + if self.registration.is_none() + && self + .notify + .permit + .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.done = true; + self.reserved_one = true; + return true; + } + let mut queue = self.notify.waiters.0.lock(); + let assigned = self + .registration + .map(|id| queue.get(id).kind) + .filter(|kind| *kind != WakeKind::Queued); + let notification = if assigned.is_some() { + assigned + } else if self.notify.generation.load(Ordering::Acquire) != self.generation { + Some(WakeKind::All) + } else if self.notify.permit.swap(false, Ordering::AcqRel) { + Some(WakeKind::One) + } else { + None + }; + if let Some(kind) = notification { + let removed = queue.remove(&mut self.registration); + self.done = true; + self.reserved_one = kind == WakeKind::One; + drop(queue); + drop(removed); return true; } - self.notify.generation.load(Ordering::Acquire) != self.generation + let old = queue.register(&mut self.registration, waker); + drop(queue); + drop(old); + false + } +} + +impl Drop for Notified<'_> { + fn drop(&mut self) { + if self.reserved_one { + self.notify.notify_one(); + return; + } + if self.registration.is_none() { + return; + } + let (removed, waker) = { + let mut queue = self.notify.waiters.0.lock(); + let removed = queue.remove(&mut self.registration); + let waker = if removed + .as_ref() + .is_some_and(|entry| entry.kind == WakeKind::One) + { + self.notify.grant_one(&mut queue) + } else { + None + }; + (removed, waker) + }; + drop(removed); + if let Some(waker) = waker { + waker.wake(); + } } } @@ -328,7 +570,10 @@ impl Semaphore { /// Wait for a permit. pub fn acquire(&self) -> Acquire<'_> { - Acquire { semaphore: self } + Acquire { + semaphore: self, + registration: None, + } } /// Hand the semaphore `n` more permits than it was created with. @@ -399,23 +644,39 @@ impl Drop for SemaphorePermit<'_> { /// The future returned by [`Semaphore::acquire`]. pub struct Acquire<'a> { semaphore: &'a Semaphore, + registration: Option, } impl<'a> Future for Acquire<'a> { type Output = SemaphorePermit<'a>; fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - if let Some(permit) = self.semaphore.try_acquire() { + let this = self.get_mut(); + if let Some(permit) = this.semaphore.try_acquire() { + this.semaphore.waiters.remove(&mut this.registration); return Poll::Ready(permit); } - self.semaphore.waiters.push(context.waker()); - if let Some(permit) = self.semaphore.try_acquire() { + this.semaphore + .waiters + .register(&mut this.registration, context.waker()); + if let Some(permit) = this.semaphore.try_acquire() { + this.semaphore.waiters.remove(&mut this.registration); return Poll::Ready(permit); } Poll::Pending } } +impl Drop for Acquire<'_> { + fn drop(&mut self) { + if self.semaphore.waiters.remove(&mut self.registration) == Some(WakeKind::One) + && self.semaphore.available_permits() != 0 + { + self.semaphore.waiters.wake_one(); + } + } +} + /// State value meaning "one writer holds this"; any other value is the number /// of readers. const WRITER: usize = usize::MAX; @@ -509,12 +770,18 @@ impl RwLock { /// Wait for a read lock. pub fn read(&self) -> Read<'_, T> { - Read { lock: self } + Read { + lock: self, + registration: None, + } } /// Wait for a write lock. pub fn write(&self) -> Write<'_, T> { - Write { lock: self } + Write { + lock: self, + registration: None, + } } } @@ -524,12 +791,18 @@ impl RwLock { /// The guard a caller can hold across an await without naming a lifetime, /// which is what `tokio::sync`'s owned guards were used for here. pub fn read_owned(self: Arc) -> ReadOwned { - ReadOwned { lock: Some(self) } + ReadOwned { + lock: Some(self), + registration: None, + } } /// The write half of [`Self::read_owned`]. pub fn write_owned(self: Arc) -> WriteOwned { - WriteOwned { lock: Some(self) } + WriteOwned { + lock: Some(self), + registration: None, + } } /// An owned read lock, or `None` if a writer holds it. @@ -663,18 +936,24 @@ impl Drop for OwnedRwLockWriteGuard { /// The future returned by [`RwLock::read`]. pub struct Read<'a, T: ?Sized> { lock: &'a RwLock, + registration: Option, } impl<'a, T: ?Sized> Future for Read<'a, T> { type Output = RwLockReadGuard<'a, T>; fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - if self.lock.try_read_raw() { - return Poll::Ready(RwLockReadGuard { lock: self.lock }); + let this = self.get_mut(); + if this.lock.try_read_raw() { + this.lock.waiters.remove(&mut this.registration); + return Poll::Ready(RwLockReadGuard { lock: this.lock }); } - self.lock.waiters.push(context.waker()); - if self.lock.try_read_raw() { - return Poll::Ready(RwLockReadGuard { lock: self.lock }); + this.lock + .waiters + .register(&mut this.registration, context.waker()); + if this.lock.try_read_raw() { + this.lock.waiters.remove(&mut this.registration); + return Poll::Ready(RwLockReadGuard { lock: this.lock }); } Poll::Pending } @@ -683,18 +962,24 @@ impl<'a, T: ?Sized> Future for Read<'a, T> { /// The future returned by [`RwLock::write`]. pub struct Write<'a, T: ?Sized> { lock: &'a RwLock, + registration: Option, } impl<'a, T: ?Sized> Future for Write<'a, T> { type Output = RwLockWriteGuard<'a, T>; fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { - if self.lock.try_write_raw() { - return Poll::Ready(RwLockWriteGuard { lock: self.lock }); + let this = self.get_mut(); + if this.lock.try_write_raw() { + this.lock.waiters.remove(&mut this.registration); + return Poll::Ready(RwLockWriteGuard { lock: this.lock }); } - self.lock.waiters.push(context.waker()); - if self.lock.try_write_raw() { - return Poll::Ready(RwLockWriteGuard { lock: self.lock }); + this.lock + .waiters + .register(&mut this.registration, context.waker()); + if this.lock.try_write_raw() { + this.lock.waiters.remove(&mut this.registration); + return Poll::Ready(RwLockWriteGuard { lock: this.lock }); } Poll::Pending } @@ -703,6 +988,7 @@ impl<'a, T: ?Sized> Future for Write<'a, T> { /// The future returned by [`RwLock::read_owned`]. pub struct ReadOwned { lock: Option>>, + registration: Option, } impl Future for ReadOwned { @@ -712,11 +998,14 @@ impl Future for ReadOwned { let this = self.get_mut(); let lock = this.lock.as_ref().expect("polled after completion"); if lock.try_read_raw() { + lock.waiters.remove(&mut this.registration); let lock = this.lock.take().expect("just checked"); return Poll::Ready(OwnedRwLockReadGuard { lock }); } - lock.waiters.push(context.waker()); + lock.waiters + .register(&mut this.registration, context.waker()); if lock.try_read_raw() { + lock.waiters.remove(&mut this.registration); let lock = this.lock.take().expect("just checked"); return Poll::Ready(OwnedRwLockReadGuard { lock }); } @@ -727,6 +1016,7 @@ impl Future for ReadOwned { /// The future returned by [`RwLock::write_owned`]. pub struct WriteOwned { lock: Option>>, + registration: Option, } impl Future for WriteOwned { @@ -736,11 +1026,14 @@ impl Future for WriteOwned { let this = self.get_mut(); let lock = this.lock.as_ref().expect("polled after completion"); if lock.try_write_raw() { + lock.waiters.remove(&mut this.registration); let lock = this.lock.take().expect("just checked"); return Poll::Ready(OwnedRwLockWriteGuard { lock }); } - lock.waiters.push(context.waker()); + lock.waiters + .register(&mut this.registration, context.waker()); if lock.try_write_raw() { + lock.waiters.remove(&mut this.registration); let lock = this.lock.take().expect("just checked"); return Poll::Ready(OwnedRwLockWriteGuard { lock }); } @@ -748,6 +1041,34 @@ impl Future for WriteOwned { } } +impl Drop for Read<'_, T> { + fn drop(&mut self) { + self.lock.waiters.remove(&mut self.registration); + } +} + +impl Drop for Write<'_, T> { + fn drop(&mut self) { + self.lock.waiters.remove(&mut self.registration); + } +} + +impl Drop for ReadOwned { + fn drop(&mut self) { + if let Some(lock) = &self.lock { + lock.waiters.remove(&mut self.registration); + } + } +} + +impl Drop for WriteOwned { + fn drop(&mut self) { + if let Some(lock) = &self.lock { + lock.waiters.remove(&mut self.registration); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -801,6 +1122,7 @@ mod tests { assert_eq!(semaphore.available_permits(), 2); } + #[cfg(feature = "std")] #[test] fn a_barrier_releases_only_once_everyone_has_arrived() { use crate::runtime::Runtime; diff --git a/src/task.rs b/src/task.rs index 9da7a75..56589c2 100644 --- a/src/task.rs +++ b/src/task.rs @@ -19,6 +19,8 @@ use core::ptr::NonNull; use async_task::{Runnable, Task}; use st3::fanout::{Act, Job, Pool}; +use crate::WorkerContext; + // Which pool worker this thread is, while it is running one. // // A raw pointer rather than an id, so a process with two pools cannot hand one @@ -31,29 +33,34 @@ std::thread_local! { /// Mark this thread as `worker` of `pool` until the guard drops. /// -/// Called by [`crate::runtime::Runtime`] around its `Pool::run`. A caller -/// driving a pool by hand can call it too; without it nothing breaks, wakes -/// simply take the injector as they always did. +/// Used by the owned runtime and Executor::run_worker. A custom/no_std host +/// can instead supply WorkerContext to its executor. #[cfg(feature = "std")] pub(crate) fn mark_current(pool: &Arc, worker: usize) -> CurrentGuard { - CURRENT.with(|current| current.set(Some((Arc::as_ptr(pool), worker)))); - CurrentGuard + let previous = CURRENT.with(|current| current.replace(Some((Arc::as_ptr(pool), worker)))); + CurrentGuard { + previous, + _thread_bound: core::marker::PhantomData, + } } -/// Clears the current-worker marker on drop, including on unwind. +/// Restores a nested marker on drop, including on unwind. Must not cross threads. #[cfg(feature = "std")] -pub(crate) struct CurrentGuard; +pub(crate) struct CurrentGuard { + previous: Option<(*const Pool, usize)>, + _thread_bound: core::marker::PhantomData>, +} #[cfg(feature = "std")] impl Drop for CurrentGuard { fn drop(&mut self) { - CURRENT.with(|current| current.set(None)); + CURRENT.with(|current| current.set(self.previous)); } } /// This thread's worker id, if it is running `pool`. #[cfg(feature = "std")] -fn current_worker(pool: &Arc) -> Option { +pub(crate) fn current_worker(pool: &Arc) -> Option { CURRENT.with(|current| match current.get() { Some((owner, worker)) if core::ptr::eq(owner, Arc::as_ptr(pool)) => Some(worker), _ => None, @@ -66,7 +73,11 @@ fn current_worker(_pool: &Arc) -> Option { } /// Put `future` on `pool` and hand back the handle to its result. -pub(crate) fn spawn(future: F, pool: Arc) -> Task +pub(crate) fn spawn( + future: F, + pool: Arc, + context: Option>, +) -> Task where F: Future + Send + 'static, F::Output: Send + 'static, @@ -74,15 +85,20 @@ where // Called for the first poll and for every wake after it. The `Runnable` is // a single heap allocation that already holds the future and its state, so // turning it into a job is `into_raw`: no box, no closure, nothing new. + let weak_pool = Arc::downgrade(&pool); let schedule = move |runnable: Runnable| { + let Some(pool) = weak_pool.upgrade() else { + // The executor/host owns pool lifetime, not a queued task. + drop(runnable); + return; + }; let pointer = runnable.into_raw(); // SAFETY: `Job::from_raw` asks for a pointer that stays valid until the // job runs, work that is `Send`, an `execute` correct for the pointer, // and one job per pointer. `Runnable::into_raw` gives up ownership of // an allocation that lives until `from_raw` takes it back, which // `poll_once` does exactly once because running consumes the job; - // `F` and `F::Output` are `Send`; and `Runnable::run` catches no panic, - // which is the same as every other job this pool takes. + // `F` and `F::Output` are `Send`. let job = unsafe { Job::from_raw(pointer, poll_once) }; // A wake that happens *on* a worker goes straight to that worker. // @@ -96,17 +112,34 @@ where // Off a worker thread this is the old path, which is what a wake from // a timer thread or an application thread wants: it has no locality to // preserve and the injector is where anyone can find the job. - match current_worker(&pool) { + match worker_for(&pool, context.as_deref()) { Some(worker) => pool.submit_local(worker, job), None => pool.submit_job(job), } }; + // Keep failures attached to the join result instead of unwinding a worker. + // This feature is enabled only by nagoya's std feature. In no_std the host + // retains responsibility for its panic policy. + #[cfg(feature = "std")] + let (runnable, task) = async_task::Builder::new() + .propagate_panic(true) + .spawn(move |()| future, schedule); + #[cfg(not(feature = "std"))] let (runnable, task) = async_task::spawn(future, schedule); runnable.schedule(); task } +fn worker_for(pool: &Arc, context: Option<&dyn WorkerContext>) -> Option { + let worker = match context { + Some(context) => context.current_worker(pool), + None => current_worker(pool), + }; + // A bad host hint must not index outside the pool or panic inside a waker. + worker.filter(|id| *id < pool.workers()) +} + /// Run one poll of the task this pointer owns, or release it unrun. /// /// The `Act::Drop` arm is not a formality. A pool dropped with work still @@ -129,3 +162,62 @@ unsafe fn poll_once(pointer: NonNull<()>, act: Act) { Act::Drop => drop(runnable), } } + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use st3::fanout::StdHost; + + fn pool() -> Arc { + Pool::new(2, 256, Arc::new(StdHost::new(2))) + } + + #[test] + fn routing_depends_on_identity_not_probabilistic_thread_migration() { + let a = pool(); + let b = pool(); + assert_eq!(worker_for(&a, None), None); + { + let _outer = mark_current(&a, 1); + assert_eq!(worker_for(&a, None), Some(1)); + assert_eq!(worker_for(&b, None), None); + { + let _inner = mark_current(&b, 0); + assert_eq!(worker_for(&a, None), None); + assert_eq!(worker_for(&b, None), Some(0)); + } + assert_eq!(worker_for(&a, None), Some(1)); + } + assert_eq!(worker_for(&a, None), None); + } + + struct HostContext { + owner: alloc::sync::Weak, + worker: usize, + } + + impl WorkerContext for HostContext { + fn current_worker(&self, pool: &Pool) -> Option { + self.owner + .upgrade() + .and_then(|owner| core::ptr::eq(owner.as_ref(), pool).then_some(self.worker)) + } + } + + #[test] + fn host_identity_routes_without_a_thread_local_marker() { + let a = pool(); + let b = pool(); + let context = HostContext { + owner: Arc::downgrade(&a), + worker: 1, + }; + assert_eq!(worker_for(&a, Some(&context)), Some(1)); + assert_eq!(worker_for(&b, Some(&context)), None); + let invalid = HostContext { + owner: Arc::downgrade(&a), + worker: usize::MAX, + }; + assert_eq!(worker_for(&a, Some(&invalid)), None); + } +} diff --git a/src/time.rs b/src/time.rs index d876fb1..87f7eeb 100644 --- a/src/time.rs +++ b/src/time.rs @@ -22,13 +22,14 @@ //! a runtime that is only ever polled should not own a thread it does not use. //! //! Without `std` there is no thread and no clock, so the consumer supplies -//! both: [`set_clock`] once, then [`poll`] from a tick interrupt or its own -//! loop. That is the whole no_std contract, and it is why the registry and the -//! driver are separate things in this file. +//! both: [`set_clock`] once, then [`poll`] from a normal execution context. +//! A tick interrupt must only signal deferred work; it must not call `poll`, +//! register/drop timers, or invoke scheduler wakers. Those paths use locks and +//! allocation and are not interrupt-safe. The host's deferred loop processes +//! the signal and calls `poll(now_ns())` after returning from the interrupt. //! //! [`Host`]: st3::fanout::Host -use alloc::collections::BinaryHeap; use alloc::sync::Arc; use alloc::vec::Vec; use core::cmp::Ordering as CmpOrdering; @@ -42,13 +43,13 @@ use spin::Mutex; /// One pending timer, shared between the future and the registry. /// -/// The future owns one end and the heap the other, so a cancelled timer can be -/// disarmed without finding and removing its heap entry: the entry survives to -/// its deadline and fires nothing. +/// The heap index is maintained under `PENDING`. Cancellation uses it to +/// remove an entry in O(log n), without scanning or retaining a tombstone. #[derive(Debug)] struct Slot { fired: AtomicBool, waker: Mutex>, + index: AtomicUsize, } /// A deadline and what to wake at it. @@ -72,7 +73,7 @@ impl PartialOrd for Entry { } } impl Ord for Entry { - /// Reversed, because `BinaryHeap` is a max-heap and the earliest deadline + /// Reversed, because the registry is a max-heap and the earliest deadline /// is the one wanted first. fn cmp(&self, other: &Self) -> CmpOrdering { other @@ -82,7 +83,92 @@ impl Ord for Entry { } } -static PENDING: Mutex>> = Mutex::new(None); +const NOT_QUEUED: usize = usize::MAX; + +// Indexed max-heap using Entry's reversed deadline ordering. Every move updates +// its slot index; index reads and writes are serialized by the registry lock. +#[derive(Default)] +struct TimerHeap { + entries: Vec, +} + +impl TimerHeap { + fn peek(&self) -> Option<&Entry> { + self.entries.first() + } + + fn swap(&mut self, a: usize, b: usize) { + self.entries.swap(a, b); + self.entries[a].slot.index.store(a, Ordering::Relaxed); + self.entries[b].slot.index.store(b, Ordering::Relaxed); + } + + fn sift_up(&mut self, mut index: usize) { + while index > 0 { + let parent = (index - 1) / 2; + if self.entries[index] <= self.entries[parent] { + break; + } + self.swap(index, parent); + index = parent; + } + } + + fn sift_down(&mut self, mut index: usize) { + // Checking for a child this way avoids overflowing 2 * index + 1. + while index < self.entries.len() / 2 { + let left = 2 * index + 1; + let right = left + 1; + let child = if right < self.entries.len() && self.entries[right] > self.entries[left] { + right + } else { + left + }; + if self.entries[index] >= self.entries[child] { + break; + } + self.swap(index, child); + index = child; + } + } + + fn push(&mut self, entry: Entry) { + let index = self.entries.len(); + self.entries.push(entry); + self.entries[index] + .slot + .index + .store(index, Ordering::Relaxed); + self.sift_up(index); + } + + fn remove(&mut self, index: usize) -> Entry { + let entry = self.entries.swap_remove(index); + entry.slot.index.store(NOT_QUEUED, Ordering::Relaxed); + if index < self.entries.len() { + self.entries[index] + .slot + .index + .store(index, Ordering::Relaxed); + if index > 0 && self.entries[index] > self.entries[(index - 1) / 2] { + self.sift_up(index); + } else { + self.sift_down(index); + } + } + entry + } + + fn pop(&mut self) -> Option { + if self.entries.is_empty() { + None + } else { + Some(self.remove(0)) + } + } +} + +static PENDING: Mutex> = Mutex::new(None); static SEQ: AtomicU64 = AtomicU64::new(0); static CLOCK: AtomicUsize = AtomicUsize::new(0); @@ -135,9 +221,10 @@ pub fn now_ns() -> u64 { /// Fire every timer due at `now`, and report when the next one is due. /// -/// The driver for this is a thread with `std`. Without one, call it from a tick -/// interrupt or an idle loop; nothing else pumps the registry, and a timer that -/// is never polled never fires. +/// The driver is a thread with `std`, or the host's deferred-work/idle loop +/// without it. This function takes ordinary spin locks, allocates, and invokes +/// arbitrary wakers. Do not call it from an interrupt, signal handler, or any +/// context that can preempt another timer operation on the same execution lane. /// /// Returns `None` when nothing is pending, so a caller with nothing else to do /// can sleep until something registers rather than spinning on an empty heap. @@ -173,7 +260,7 @@ pub fn poll(now: u64) -> Option { /// earliest pending timer. fn arm(deadline: u64, slot: &Arc) -> bool { let mut guard = PENDING.lock(); - let heap = guard.get_or_insert_with(BinaryHeap::new); + let heap = guard.get_or_insert_with(TimerHeap::default); let earliest = heap.peek().is_none_or(|entry| deadline < entry.deadline); heap.push(Entry { deadline, @@ -185,9 +272,8 @@ fn arm(deadline: u64, slot: &Arc) -> bool { /// A future that is ready once its deadline has passed. /// -/// Dropping one disarms it. The heap entry survives until its deadline, holding -/// nothing and waking nothing, so a cancelled `timeout` costs one dead entry -/// for as long as it had left to run rather than forever. +/// Dropping an armed sleep removes its heap entry in O(log n). Registration, +/// polling, and dropping must run outside interrupt/signal-handler contexts. #[derive(Debug)] pub struct Sleep { deadline: u64, @@ -263,6 +349,7 @@ impl Future for Sleep { let slot = Arc::new(Slot { fired: AtomicBool::new(false), waker: Mutex::new(Some(cx.waker().clone())), + index: AtomicUsize::new(NOT_QUEUED), }); let earliest = arm(this.deadline, &slot); this.slot = Some(slot); @@ -278,10 +365,25 @@ impl Future for Sleep { impl Drop for Sleep { fn drop(&mut self) { - // Disarm rather than unregister. Finding this entry in the heap costs a - // linear scan under the lock; leaving it costs one wake of nothing. if let Some(slot) = &self.slot { - *slot.waker.lock() = None; + let removed = { + let mut guard = PENDING.lock(); + let index = slot.index.load(Ordering::Relaxed); + if index == NOT_QUEUED { + None + } else { + Some( + guard + .as_mut() + .expect("queued timer has a registry") + .remove(index), + ) + } + }; + // Release heap ownership and the captured task outside PENDING. + drop(removed); + let waker = slot.waker.lock().take(); + drop(waker); } } } @@ -420,6 +522,83 @@ mod tests { use core::time::Duration; use std::time::Instant; + fn slot() -> alloc::sync::Arc { + alloc::sync::Arc::new(super::Slot { + fired: core::sync::atomic::AtomicBool::new(false), + waker: spin::Mutex::new(None), + index: core::sync::atomic::AtomicUsize::new(super::NOT_QUEUED), + }) + } + + #[test] + fn indexed_heap_removal_preserves_order_and_indices() { + use core::sync::atomic::Ordering; + let mut heap = super::TimerHeap::default(); + let mut slots = alloc::vec::Vec::new(); + for (seq, deadline) in [90, 10, 50, 10, 80, 20, 60, 30, 70, 40] + .into_iter() + .enumerate() + { + let slot = slot(); + heap.push(super::Entry { + deadline, + seq: seq as u64, + slot: slot.clone(), + }); + slots.push(slot); + } + for id in [2, 0, 7] { + let index = slots[id].index.load(Ordering::Relaxed); + let removed = heap.remove(index); + assert!(alloc::sync::Arc::ptr_eq(&removed.slot, &slots[id])); + assert_eq!(slots[id].index.load(Ordering::Relaxed), super::NOT_QUEUED); + for (index, entry) in heap.entries.iter().enumerate() { + assert_eq!(entry.slot.index.load(Ordering::Relaxed), index); + } + } + let mut order = alloc::vec::Vec::new(); + while let Some(entry) = heap.pop() { + assert_eq!(entry.slot.index.load(Ordering::Relaxed), super::NOT_QUEUED); + order.push((entry.deadline, entry.seq)); + } + assert_eq!( + order, + [ + (10, 1), + (10, 3), + (20, 5), + (40, 9), + (60, 6), + (70, 8), + (80, 4) + ] + ); + } + + #[test] + fn dropping_an_armed_sleep_releases_registry_ownership_immediately() { + use core::future::Future; + use core::pin::Pin; + use core::task::Context; + use std::sync::Arc; + use std::task::{Wake, Waker}; + struct Noop; + impl Wake for Noop { + fn wake(self: Arc) {} + } + let waker = Waker::from(Arc::new(Noop)); + let mut sleep = super::sleep_until(u64::MAX); + assert!(Pin::new(&mut sleep) + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + let weak = Arc::downgrade(sleep.slot.as_ref().unwrap()); + drop(sleep); + assert!( + weak.upgrade().is_none(), + "canceled timer is still retained by the registry" + ); + } + #[test] fn a_sleep_waits_at_least_its_duration() { let started = Instant::now(); diff --git a/tests/executor.rs b/tests/executor.rs index a994b68..8a39092 100644 --- a/tests/executor.rs +++ b/tests/executor.rs @@ -268,3 +268,102 @@ fn a_task_can_yield_its_worker() { assert_eq!(long, Some(1), "the yielding task did not let the other in"); assert_eq!(short, Some(0)); } + +#[test] +fn a_reversed_empty_range_finishes_without_starting_workers() { + let pool = Pool::new(1, 256, Arc::new(StdHost::new(1))); + let mut loop_ = Box::pin(nagoya::par_for(pool, 10..5, |_| panic!("empty range ran"))); + struct Noop; + impl std::task::Wake for Noop { + fn wake(self: Arc) {} + } + let waker = std::task::Waker::from(Arc::new(Noop)); + assert!(loop_ + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_ready()); +} + +#[test] +fn unrun_tasks_do_not_keep_their_containing_pool_alive() { + struct Dropped(Arc); + impl Drop for Dropped { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + let pool = Pool::new(1, 256, Arc::new(StdHost::new(1))); + let weak_pool = Arc::downgrade(&pool); + let executor = Executor::new(pool.clone()); + let drops = Arc::new(AtomicUsize::new(0)); + let captured = Dropped(drops.clone()); + let handle = executor.spawn(async move { + drop(captured); + }); + drop(executor); + drop(pool); + assert!(weak_pool.upgrade().is_none()); + assert_eq!(drops.load(Ordering::Relaxed), 1); + // Pool destruction canceled the runnable, not the caller's join handle. + assert_eq!(block_on(handle), None); +} + +#[test] +fn dropping_a_pool_retires_queued_parallel_pieces() { + let pool = Pool::new(1, 256, Arc::new(StdHost::new(1))); + let weak_pool = Arc::downgrade(&pool); + let mut loop_ = Box::pin(nagoya::par_for(pool.clone(), 0..10, |_| { + panic!("unrun body") + })); + struct Noop; + impl std::task::Wake for Noop { + fn wake(self: Arc) {} + } + let waker = std::task::Waker::from(Arc::new(Noop)); + let mut context = Context::from_waker(&waker); + assert!(loop_.as_mut().poll(&mut context).is_pending()); + drop(pool); + assert!(weak_pool.upgrade().is_none()); + assert!(loop_.as_mut().poll(&mut context).is_ready()); +} + +#[cfg(feature = "std")] +#[test] +fn a_panicking_task_does_not_take_the_only_worker_with_it() { + let running = Running::new(1); + let failed = running.executor.spawn(async { panic!("task failure") }); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + block_on(nagoya::timeout(std::time::Duration::from_secs(2), failed)) + })); + assert!( + outcome.is_err(), + "the task panic was not propagated to its awaiter" + ); + let next = running.executor.spawn(async { 42 }); + assert_eq!( + block_on(nagoya::timeout(std::time::Duration::from_secs(2), next)), + Ok(Some(42)), + ); +} + +#[cfg(feature = "std")] +#[test] +fn a_panicking_parallel_body_retires_its_piece_and_preserves_the_worker() { + let running = Running::new(1); + let loop_ = nagoya::par_for(running.executor.pool().clone(), 0..32, |_| { + panic!("body failure") + }) + .leaf(1); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + block_on(nagoya::timeout(std::time::Duration::from_secs(2), loop_)) + })); + assert!( + outcome.is_err(), + "the parallel-loop panic was not propagated" + ); + let next = running.executor.spawn(async { 42 }); + assert_eq!( + block_on(nagoya::timeout(std::time::Duration::from_secs(2), next)), + Ok(Some(42)), + ); +} diff --git a/tests/io_regressions.rs b/tests/io_regressions.rs new file mode 100644 index 0000000..b5e01df --- /dev/null +++ b/tests/io_regressions.rs @@ -0,0 +1,95 @@ +use nagoya::block_on; +use nagoya::io::{Error, ErrorKind, Read, Write}; + +struct InterruptedReader { + bytes: &'static [u8], + cursor: usize, + interrupt: bool, +} + +impl Read for InterruptedReader { + async fn read(&mut self, buffer: &mut [u8]) -> Result { + if self.interrupt { + self.interrupt = false; + return Err(Error::new(ErrorKind::Interrupted)); + } + let n = buffer.len().min(2).min(self.bytes.len() - self.cursor); + buffer[..n].copy_from_slice(&self.bytes[self.cursor..self.cursor + n]); + self.cursor += n; + self.interrupt = true; + Ok(n) + } +} + +#[derive(Default)] +struct InterruptedWriter { + bytes: Vec, + interrupt: bool, +} + +impl Write for InterruptedWriter { + async fn write(&mut self, buffer: &[u8]) -> Result { + if self.interrupt { + self.interrupt = false; + return Err(Error::new(ErrorKind::Interrupted)); + } + let n = buffer.len().min(2); + self.bytes.extend_from_slice(&buffer[..n]); + self.interrupt = true; + Ok(n) + } + + async fn flush(&mut self) -> Result<(), Error> { + Ok(()) + } +} + +#[test] +fn read_exact_keeps_progress_across_interruptions() { + let mut reader = InterruptedReader { + bytes: b"abcdef", + cursor: 0, + interrupt: true, + }; + let mut bytes = [0; 6]; + block_on(reader.read_exact(&mut bytes)).unwrap(); + assert_eq!(&bytes, b"abcdef"); + assert_eq!(reader.cursor, 6); +} + +#[test] +fn read_to_end_keeps_the_prefix_and_drops_scratch_space_on_interruption() { + let mut reader = InterruptedReader { + bytes: b"abcdef", + cursor: 0, + interrupt: true, + }; + let mut bytes = b"prefix:".to_vec(); + assert_eq!(block_on(reader.read_to_end(&mut bytes)).unwrap(), 6); + assert_eq!(&bytes, b"prefix:abcdef"); +} + +#[test] +fn write_all_neither_repeats_nor_skips_bytes_after_interruption() { + let mut writer = InterruptedWriter { + interrupt: true, + ..Default::default() + }; + block_on(writer.write_all(b"abcdef")).unwrap(); + assert_eq!(&writer.bytes, b"abcdef"); +} + +#[test] +fn real_eof_is_still_an_error_for_read_exact() { + let mut reader = InterruptedReader { + bytes: b"ab", + cursor: 0, + interrupt: false, + }; + let mut bytes = [0; 3]; + assert_eq!( + block_on(reader.read_exact(&mut bytes)).unwrap_err().kind(), + ErrorKind::UnexpectedEof, + ); + assert_eq!(&bytes[..2], b"ab"); +} diff --git a/tests/sync_regressions.rs b/tests/sync_regressions.rs new file mode 100644 index 0000000..a5a6d84 --- /dev/null +++ b/tests/sync_regressions.rs @@ -0,0 +1,186 @@ +//! Deterministic polling tests: no executor scheduling or sleeps are involved. +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll, Wake, Waker}; + +use nagoya::sync::{Notify, RwLock, Semaphore}; + +#[derive(Default)] +struct Count(AtomicUsize); + +impl Wake for Count { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +fn counter() -> (Arc, Waker) { + let count = Arc::new(Count::default()); + let waker = Waker::from(count.clone()); + (count, waker) +} + +fn poll(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +#[test] +fn canceling_the_first_acquirer_does_not_consume_the_only_wake() { + let semaphore = Semaphore::new(0); + let (a_count, a_waker) = counter(); + let (b_count, b_waker) = counter(); + let mut a = Box::pin(semaphore.acquire()); + let mut b = Box::pin(semaphore.acquire()); + assert!(poll(a.as_mut(), &a_waker).is_pending()); + assert!(poll(b.as_mut(), &b_waker).is_pending()); + drop(a); + semaphore.add_permits(1); + assert_eq!(a_count.0.load(Ordering::Relaxed), 0); + assert_eq!(b_count.0.load(Ordering::Relaxed), 1); + assert!(poll(b.as_mut(), &b_waker).is_ready()); +} + +#[test] +fn canceling_an_awakened_acquirer_hands_the_wake_on() { + let semaphore = Semaphore::new(0); + let (_, a_waker) = counter(); + let (b_count, b_waker) = counter(); + let mut a = Box::pin(semaphore.acquire()); + let mut b = Box::pin(semaphore.acquire()); + assert!(poll(a.as_mut(), &a_waker).is_pending()); + assert!(poll(b.as_mut(), &b_waker).is_pending()); + semaphore.add_permits(1); + drop(a); + assert_eq!(b_count.0.load(Ordering::Relaxed), 1); + assert!(poll(b.as_mut(), &b_waker).is_ready()); +} + +#[test] +fn two_notifications_belong_to_two_registered_waiters() { + let notify = Notify::new(); + // A task may await multiple futures with the very same waker. + let (_, waker) = counter(); + let mut a = Box::pin(notify.notified()); + let mut b = Box::pin(notify.notified()); + assert!(poll(a.as_mut(), &waker).is_pending()); + assert!(poll(b.as_mut(), &waker).is_pending()); + notify.notify_one(); + notify.notify_one(); + assert!(poll(a.as_mut(), &waker).is_ready()); + assert!(poll(b.as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_pending()); +} + +#[test] +fn enable_registers_each_waiter_without_needing_a_waker() { + let notify = Notify::new(); + let (_, waker) = counter(); + let mut a = Box::pin(notify.notified()); + let mut b = Box::pin(notify.notified()); + assert!(!a.as_mut().enable()); + assert!(!b.as_mut().enable()); + notify.notify_one(); + notify.notify_one(); + assert!(poll(a.as_mut(), &waker).is_ready()); + assert!(poll(b.as_mut(), &waker).is_ready()); +} + +#[test] +fn canceling_a_notified_waiter_transfers_its_notification() { + let notify = Notify::new(); + let (_, a_waker) = counter(); + let (b_count, b_waker) = counter(); + let mut a = Box::pin(notify.notified()); + let mut b = Box::pin(notify.notified()); + assert!(poll(a.as_mut(), &a_waker).is_pending()); + assert!(poll(b.as_mut(), &b_waker).is_pending()); + notify.notify_one(); + drop(a); + assert_eq!(b_count.0.load(Ordering::Relaxed), 1); + assert!(poll(b.as_mut(), &b_waker).is_ready()); +} + +#[test] +fn canceled_assigned_notification_becomes_a_spare_permit() { + let notify = Notify::new(); + let (_, waker) = counter(); + let mut a = Box::pin(notify.notified()); + assert!(!a.as_mut().enable()); + notify.notify_one(); + drop(a); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_pending()); +} + +#[test] +fn dropping_an_enabled_ready_future_returns_its_reserved_permit() { + let notify = Notify::new(); + let (_, waker) = counter(); + notify.notify_one(); + let mut a = Box::pin(notify.notified()); + assert!(a.as_mut().enable()); + // enable acknowledges readiness but has not delivered the result to poll. + drop(a); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_pending()); +} + +#[test] +fn a_repolled_waiter_updates_its_waker_without_duplicate_registration() { + let semaphore = Semaphore::new(0); + let (old_count, old) = counter(); + let (new_count, new) = counter(); + let mut acquire = Box::pin(semaphore.acquire()); + assert!(poll(acquire.as_mut(), &old).is_pending()); + assert!(poll(acquire.as_mut(), &new).is_pending()); + semaphore.add_permits(1); + assert_eq!(old_count.0.load(Ordering::Relaxed), 0); + assert_eq!(new_count.0.load(Ordering::Relaxed), 1); + assert!(poll(acquire.as_mut(), &new).is_ready()); +} + +#[test] +fn broadcasts_reach_old_futures_but_not_new_ones() { + let notify = Notify::new(); + let (_, waker) = counter(); + let mut unpolled = Box::pin(notify.notified()); + let mut registered = Box::pin(notify.notified()); + assert!(poll(registered.as_mut(), &waker).is_pending()); + notify.notify_waiters(); + assert!(poll(unpolled.as_mut(), &waker).is_ready()); + assert!(poll(registered.as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_pending()); +} + +#[test] +fn an_old_broadcast_does_not_consume_a_new_spare_notification() { + let notify = Notify::new(); + let (_, waker) = counter(); + let mut old = Box::pin(notify.notified()); + notify.notify_waiters(); + notify.notify_one(); + assert!(poll(old.as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_ready()); + assert!(poll(Box::pin(notify.notified()).as_mut(), &waker).is_pending()); +} + +#[test] +fn canceling_lock_waits_releases_their_wakers_and_reuses_slots() { + let lock = Arc::new(RwLock::new(0)); + let held = lock.try_write().unwrap(); + let (count, waker) = counter(); + for _ in 0..128 { + let mut read = Box::pin(lock.clone().read_owned()); + assert!(poll(read.as_mut(), &waker).is_pending()); + drop(read); + assert_eq!(Arc::strong_count(&count), 2); + } + drop(held); + assert_eq!(count.0.load(Ordering::Relaxed), 0); + assert!(lock.try_write().is_some()); +}