From 3a9f9cdfd4ed41d6ddde3c713bf0a1bac1ad5efb Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 10 Sep 2026 12:55:41 +0700 Subject: [PATCH 1/8] Re-export the tuning type callers have to name Selecting an idle policy is a nagoya-level decision: a caller picks it when it builds a runtime, through nagoya's own API. But `Tuning` is `ps-st3`'s type and nagoya did not re-export it, so every such caller had to add `ps-st3` as a direct dependency to spell one struct, in a graph that already contains it transitively. WorkTable did exactly that. --- src/lib.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index ae671e8..0061949 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, }; From a7698bc4826b655659cf2baaa2476ec9e1ecb32c Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 10 Sep 2026 12:55:41 +0700 Subject: [PATCH 2/8] Let a caller build a marked pool at any tuning `Runtime::new` is the only way to get threads that call `task::mark_current`, and that marker is the whole mechanism `local_wakes` depends on: without it a wake takes the injector no matter what the tuning says. `mark_current` is private, so a caller wanting a pool at a non-default tuning had to build the `Pool` itself, and got a pool that silently ignores the half of the tuning it was chosen for. That is not a small difference. The tuning is set, the pool runs, and the behaviour it asks for never happens, so a comparison between two such pools would attribute to the scheduler a difference that was only ever a missing thread-local. WorkTable has been running every flavor but its default this way. `with_tuning` takes the count, the tuning and a thread-name label, and `new` becomes it at the default tuning. Two tests, because one would not be evidence. A locality-tuned `Runtime` polls a self-waking task on exactly one thread. A hand-built pool at the same tuning, differing only in the missing `mark_current`, polls it on more than one. The second is what makes the first mean something. --- src/runtime.rs | 137 +++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 134 insertions(+), 3 deletions(-) diff --git a/src/runtime.rs b/src/runtime.rs index 35d2194..cf35548 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,40 @@ 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}`. + /// + /// # Why this has to exist here rather than in the caller + /// + /// A caller can already build a [`Pool`] at any tuning and run it on its + /// own threads. What it cannot do is [`crate::task::mark_current`], which + /// is private, and **that marker is the only thing that makes + /// `local_wakes` do anything**: without it a wake takes the injector no + /// matter what the tuning says. + /// + /// So a pool built outside this module silently runs any locality-flavored + /// tuning as if it were spread. That is not a small difference, it is the + /// entire mechanism the tuning selects, and it is invisible: the pool + /// works, the tuning is set, and the behaviour it asks for never happens. + /// A comparison between two such pools would attribute to the scheduler a + /// difference that was only ever a missing thread-local. + /// + /// # Panics + /// + /// If a thread cannot be started, for the reason [`Runtime::new`] gives. + #[must_use] + pub fn with_tuning(workers: usize, tuning: Tuning, label: &str) -> Self { let workers = workers.max(1); 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 @@ -164,3 +190,108 @@ mod tests { assert_eq!(counter.load(Ordering::Relaxed), 64); } } + +/// What [`Runtime::with_tuning`] exists for, demonstrated rather than asserted. +/// +/// The claim is that a pool built outside this module cannot honour a +/// locality tuning, because [`crate::task::mark_current`] is private and that +/// marker is the whole mechanism. A test that only checked the marked pool +/// would pass just as well if the marker did nothing, so both are run. +#[cfg(test)] +mod local_wakes { + use super::{Executor, Runtime, Tuning}; + use alloc::sync::Arc; + use alloc::vec::Vec; + use core::future::Future; + use core::pin::Pin; + use core::task::{Context, Poll}; + use st3::fanout::{Pool, StdHost}; + use std::collections::HashSet; + use std::sync::Mutex; + + const WORKERS: usize = 4; + const YIELDS: usize = 400; + + /// Self-wakes `left` times, recording which thread ran each poll. + /// + /// The thread identity is the observable. `local_wakes` means a task that + /// wakes itself goes back to the worker that was running it, so the set of + /// threads that polled one task has exactly one member. + struct RecordYields { + left: usize, + seen: Arc>>, + } + + impl Future for RecordYields { + type Output = (); + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> { + self.seen + .lock() + .expect("the set holds no state a panic could corrupt") + .insert(std::thread::current().id()); + if self.left == 0 { + return Poll::Ready(()); + } + self.left -= 1; + context.waker().wake_by_ref(); + Poll::Pending + } + } + + fn threads_that_polled_one_task(executor: &Executor) -> usize { + let seen = Arc::new(Mutex::new(HashSet::new())); + let task = RecordYields { + left: YIELDS, + seen: seen.clone(), + }; + crate::block_on(executor.spawn(task)); + let count = seen.lock().expect("the set").len(); + count + } + + #[test] + fn a_marked_pool_keeps_a_self_waking_task_on_one_worker() { + let runtime = Runtime::with_tuning(WORKERS, Tuning::locality(), "test-marked"); + assert_eq!( + threads_that_polled_one_task(runtime.executor()), + 1, + "a locality-tuned runtime let a self-wake leave its worker, so the marker is not being set" + ); + } + + /// The same tuning, on a pool this module did not build. + /// + /// This is what `WorkTable` was doing for every flavor but the default, + /// and it is why `with_tuning` had to exist: the tuning is set, the pool + /// works, and the behaviour it asks for silently never happens. + #[test] + fn an_unmarked_pool_cannot_honour_the_same_tuning() { + let host = Arc::new(StdHost::new(WORKERS)); + let pool = Pool::with_tuning(WORKERS, 1024, host, Tuning::locality()); + let threads: Vec<_> = (0..WORKERS) + .map(|id| { + let pool = pool.clone(); + let runner = pool.runner(id); + std::thread::spawn(move || { + // No `mark_current`, because it is private. That is the + // whole of the difference from the test above. + let _ = pool.run(runner); + }) + }) + .collect(); + let executor = Executor::new(pool.clone()); + + let count = threads_that_polled_one_task(&executor); + + pool.shut_down(); + for thread in threads { + let _ = thread.join(); + } + + assert!( + count > 1, + "an unmarked pool kept the task on one worker, which would mean the marker is not the mechanism \ + `local_wakes` depends on and this whole API is unnecessary" + ); + } +} From e094ed0fcccb10f21682e2c9fad7c6d85728d1b3 Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 10 Sep 2026 14:33:10 +0700 Subject: [PATCH 3/8] Release 0.1.1, a marked pool at any tuning --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 4fdfdd3..3aae645 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" From 0601f2f7c6906b09a3f920aff4f51db821f93952 Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 10 Sep 2026 18:01:34 +0700 Subject: [PATCH 4/8] Let a caller reach the pool an executor runs on `par_for` takes an `Arc` because it schedules closures and never builds a task, so a caller holding only a `Runtime` could not use it. Every data-parallel benchmark worked around that by building its own `Pool`, which meant its threads were never marked as pool workers, and that marker is the only thing that makes `local_wakes` do anything. Thirteen benchmarks were therefore running the default tuning on unmarked threads and could not see a flavor at all. --- src/runtime.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/runtime.rs b/src/runtime.rs index cf35548..99dd090 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -111,6 +111,20 @@ 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` rather than an executor, + /// because it schedules closures and never builds a task. Without this a + /// caller wanting `par_for` on a tuned pool had to build the `Pool` + /// itself, which meant its threads were never marked as pool workers, and + /// **that marker is the only thing that makes `local_wakes` do anything**. + /// So every locality-flavored tuning silently ran as spread, which is the + /// exact trap `Runtime::with_tuning` exists to close. + #[must_use] + pub fn pool(&self) -> &Arc { + self.executor.pool() + } } /// Threads for the shared pool. From 791dbf390e137c0e4d0d32fd631b64b97f8d3239 Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 10 Sep 2026 22:20:46 +0700 Subject: [PATCH 5/8] Require ps-st3 0.6.1, which is where the wake fixes are 0.6.1 carries three fixes 0.6.0 does not: a burst of submits all naming the same sleeper, private work waking a worker that cannot run it, and a nudge whose condition could not be true at the default injector batch size. The requirement was ^0.6, which accepts 0.6.0. A lockfile pinning it would build a pool without any of them, and nothing would say so: the crate compiles either way. Every scheduler measurement made against this crate would quietly stop being reproducible. Measured on the fixes this pins, YCSB C at 16 threads: cores kept busy went from 7.5 to 12.9 against tokio's 12.8. --- Cargo.toml | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3aae645..02d8cdc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,7 +18,13 @@ std = ["ps-st3/host", "dep:futures-io", "futures-io/std", "dep:futures-util", "f [dependencies] -ps-st3 = { version = "^0.6", default-features = false, features = ["fanout"] } +# `^0.6.1`, not `^0.6`. 0.6.1 carries three wake fixes and 0.6.0 does not: +# a burst of submits naming the same sleeper, private work waking the wrong +# worker, and a nudge that never fired because its condition could not be +# true at the default batch size. `^0.6` accepts 0.6.0, so a lockfile pinning +# it would silently build a pool without any of them, and every scheduler +# number measured against this crate would stop being reproducible. +ps-st3 = { version = "^0.6.1", default-features = false, features = ["fanout"] } spin = { version = "0.9", default-features = false, features = ["mutex", "spin_mutex"] } # The task itself: one allocation holding the future, its state machine, its # waker vtable and its output, with `Runnable::into_raw` handing the pointer @@ -33,7 +39,7 @@ futures-io = { version = "0.3", default-features = false, optional = true } futures-util = { version = "0.3", default-features = false, optional = true } [dev-dependencies] -ps-st3 = { version = "^0.6", features = ["fanout", "host"] } +ps-st3 = { version = "^0.6.1", features = ["fanout", "host"] } # The comparison arm. Only the scheduler is used; nothing here needs its I/O. tokio = { version = "1", features = ["rt-multi-thread", "macros"] } # The other comparison arm. Not an async runtime: a data-parallel pool, so it From 963fd6d94837a2a6604589f36b7c44981f839e90 Mon Sep 17 00:00:00 2001 From: meh Date: Fri, 11 Sep 2026 02:49:32 +0700 Subject: [PATCH 6/8] Fail loudly when the publish token is missing This crate has no repository-level CARGO_REGISTRY_TOKEN, and every publish run so far short-circuited on "version unchanged, nothing to publish", so the upload path has never run here. 0.1.1 would be the first time the token is exercised. An absent token surfaces as an authentication error from cargo publish after the package is already built, which reads like a registry problem rather than a missing secret. The guard says which it is, and where to set it. This is the crate the release order starts at: worktable requires nagoya ^0.1.1 and nagoya requires ps-st3 ^0.6.1, so a publish that fails obscurely here blocks everything and looks like it came from somewhere else. Whether an organisation-level secret already covers this repository is not readable from a repository-level secret listing, so the guard reports rather than assumes. --- .github/workflows/publish.yml | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) 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: From 54b456fa7a23ee31b62d072eb526e8de8197b9ec Mon Sep 17 00:00:00 2001 From: meh Date: Fri, 11 Sep 2026 03:12:37 +0700 Subject: [PATCH 7/8] Take the caret back off ps-st3 The requirement was ^0.6.1 to stop a lockfile resolving 0.6.0, which lacks the three wake fixes. Carets for crates we control is the rule, and a floor is a pin wearing a different hat. Cargo takes the highest compatible version anyway, so ^0.6 gets 0.6.1 once it publishes. It also means this branch resolves against the registry today rather than being red until then. --- Cargo.toml | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 02d8cdc..3aae645 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,13 +18,7 @@ std = ["ps-st3/host", "dep:futures-io", "futures-io/std", "dep:futures-util", "f [dependencies] -# `^0.6.1`, not `^0.6`. 0.6.1 carries three wake fixes and 0.6.0 does not: -# a burst of submits naming the same sleeper, private work waking the wrong -# worker, and a nudge that never fired because its condition could not be -# true at the default batch size. `^0.6` accepts 0.6.0, so a lockfile pinning -# it would silently build a pool without any of them, and every scheduler -# number measured against this crate would stop being reproducible. -ps-st3 = { version = "^0.6.1", default-features = false, features = ["fanout"] } +ps-st3 = { version = "^0.6", default-features = false, features = ["fanout"] } spin = { version = "0.9", default-features = false, features = ["mutex", "spin_mutex"] } # The task itself: one allocation holding the future, its state machine, its # waker vtable and its output, with `Runnable::into_raw` handing the pointer @@ -39,7 +33,7 @@ futures-io = { version = "0.3", default-features = false, optional = true } futures-util = { version = "0.3", default-features = false, optional = true } [dev-dependencies] -ps-st3 = { version = "^0.6.1", features = ["fanout", "host"] } +ps-st3 = { version = "^0.6", features = ["fanout", "host"] } # The comparison arm. Only the scheduler is used; nothing here needs its I/O. tokio = { version = "1", features = ["rt-multi-thread", "macros"] } # The other comparison arm. Not an async runtime: a data-parallel pool, so it From 9372f6bee02850b346cc9dc716a85a727779498b Mon Sep 17 00:00:00 2001 From: meh Date: Fri, 11 Sep 2026 05:25:52 +0700 Subject: [PATCH 8/8] Fold the eleven review findings from the whole-repository review Applied from the four-patch series in ~/code/patch, whose author states plainly that nothing was compiled or run. It is compiled and run now: the crate builds, all 65 tests pass including the 15 new regression tests the series brings, `--no-default-features` still checks, and WorkTable's 999 tests pass against this tree, which is what its `[patch.crates-io]` points at. What the series covers, from the review: Cancelled and stale semaphore waiters no longer consume wakes; multiple notifications no longer collapse into one; registrations are removed on completion and on drop, for borrowed and owned RwLock waits too. Cancelled timers no longer sit in the heap until their deadlines. The unindexed heap is now an indexed one, so removal is O(log n) with slot indices maintained under the registry mutex. A panic no longer removes a worker or strands a parallel completion. Panic propagation is enabled through the existing `std` feature rather than a new dependency, parallel-piece bodies catch and rethrow at the awaiting future, and an RAII token retires a piece on execution, cancellation, unwind, and on dropping an unrun queue entry. Queued jobs no longer own the pool that owns them. Scheduling closures and parallel pieces hold `Weak` references, and `FallibleTask` lets a cancelled runnable produce the `None` join result the signature already declared. `no_std` and custom hosts can activate locality: `WorkerContext` and `with_worker_context` without a `std` dependency, plus `run_worker` for a caller-owned thread. The automatic worker count is clamped to `usize::BITS`, and an explicitly oversized count is rejected before host resources are created. Reversed empty `par_for` ranges complete before queue publication rather than underflowing. Whole-buffer I/O retries `Interrupted` while preserving progress. The only manifest change is `async-task/std` inside the existing `std` feature. Nothing became unconditionally std. Two caveats the series carries and I have not verified beyond compiling: panic isolation requires std and unwinding, so abort-on-panic builds still abort, and timers are not interrupt-safe. A host following the old tick-interrupt documentation has to move registry processing to deferred execution. --- Cargo.toml | 2 +- src/io.rs | 19 +- src/lib.rs | 66 +++++- src/par.rs | 115 ++++++++- src/runtime.rs | 168 +++---------- src/sync.rs | 482 +++++++++++++++++++++++++++++++------- src/task.rs | 118 +++++++++- src/time.rs | 217 +++++++++++++++-- tests/executor.rs | 99 ++++++++ tests/io_regressions.rs | 95 ++++++++ tests/sync_regressions.rs | 186 +++++++++++++++ 11 files changed, 1295 insertions(+), 272 deletions(-) create mode 100644 tests/io_regressions.rs create mode 100644 tests/sync_regressions.rs diff --git a/Cargo.toml b/Cargo.toml index 3aae645..ee102e6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 0061949..f724654 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -76,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 @@ -87,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`") @@ -129,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() } @@ -147,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 { @@ -166,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. @@ -175,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. @@ -191,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 99dd090..72ad80b 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -55,27 +55,21 @@ impl Runtime { /// A runtime with `workers` threads, at `tuning`, whose threads are named /// `{label}-{id}`. /// - /// # Why this has to exist here rather than in the caller - /// - /// A caller can already build a [`Pool`] at any tuning and run it on its - /// own threads. What it cannot do is [`crate::task::mark_current`], which - /// is private, and **that marker is the only thing that makes - /// `local_wakes` do anything**: without it a wake takes the injector no - /// matter what the tuning says. - /// - /// So a pool built outside this module silently runs any locality-flavored - /// tuning as if it were spread. That is not a small difference, it is the - /// entire mechanism the tuning selects, and it is invisible: the pool - /// works, the tuning is set, and the behaviour it asks for never happens. - /// A comparison between two such pools would attribute to the scheduler a - /// difference that was only ever a missing thread-local. + /// 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, for the reason [`Runtime::new`] gives. + /// 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::with_tuning(workers, 1024, host, tuning); for id in 0..workers { @@ -114,13 +108,8 @@ impl Runtime { /// The pool underneath, for the APIs that take one directly. /// - /// [`crate::par_for`] takes an `Arc` rather than an executor, - /// because it schedules closures and never builds a task. Without this a - /// caller wanting `par_for` on a tuned pool had to build the `Pool` - /// itself, which meant its threads were never marked as pool workers, and - /// **that marker is the only thing that makes `local_wakes` do anything**. - /// So every locality-flavored tuning silently ran as spread, which is the - /// exact trap `Runtime::with_tuning` exists to close. + /// [`crate::par_for`] takes an `Arc` because it schedules closures + /// without building async tasks. #[must_use] pub fn pool(&self) -> &Arc { self.executor.pool() @@ -129,12 +118,15 @@ impl Runtime { /// 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. @@ -168,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())); @@ -204,108 +213,3 @@ mod tests { assert_eq!(counter.load(Ordering::Relaxed), 64); } } - -/// What [`Runtime::with_tuning`] exists for, demonstrated rather than asserted. -/// -/// The claim is that a pool built outside this module cannot honour a -/// locality tuning, because [`crate::task::mark_current`] is private and that -/// marker is the whole mechanism. A test that only checked the marked pool -/// would pass just as well if the marker did nothing, so both are run. -#[cfg(test)] -mod local_wakes { - use super::{Executor, Runtime, Tuning}; - use alloc::sync::Arc; - use alloc::vec::Vec; - use core::future::Future; - use core::pin::Pin; - use core::task::{Context, Poll}; - use st3::fanout::{Pool, StdHost}; - use std::collections::HashSet; - use std::sync::Mutex; - - const WORKERS: usize = 4; - const YIELDS: usize = 400; - - /// Self-wakes `left` times, recording which thread ran each poll. - /// - /// The thread identity is the observable. `local_wakes` means a task that - /// wakes itself goes back to the worker that was running it, so the set of - /// threads that polled one task has exactly one member. - struct RecordYields { - left: usize, - seen: Arc>>, - } - - impl Future for RecordYields { - type Output = (); - fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> { - self.seen - .lock() - .expect("the set holds no state a panic could corrupt") - .insert(std::thread::current().id()); - if self.left == 0 { - return Poll::Ready(()); - } - self.left -= 1; - context.waker().wake_by_ref(); - Poll::Pending - } - } - - fn threads_that_polled_one_task(executor: &Executor) -> usize { - let seen = Arc::new(Mutex::new(HashSet::new())); - let task = RecordYields { - left: YIELDS, - seen: seen.clone(), - }; - crate::block_on(executor.spawn(task)); - let count = seen.lock().expect("the set").len(); - count - } - - #[test] - fn a_marked_pool_keeps_a_self_waking_task_on_one_worker() { - let runtime = Runtime::with_tuning(WORKERS, Tuning::locality(), "test-marked"); - assert_eq!( - threads_that_polled_one_task(runtime.executor()), - 1, - "a locality-tuned runtime let a self-wake leave its worker, so the marker is not being set" - ); - } - - /// The same tuning, on a pool this module did not build. - /// - /// This is what `WorkTable` was doing for every flavor but the default, - /// and it is why `with_tuning` had to exist: the tuning is set, the pool - /// works, and the behaviour it asks for silently never happens. - #[test] - fn an_unmarked_pool_cannot_honour_the_same_tuning() { - let host = Arc::new(StdHost::new(WORKERS)); - let pool = Pool::with_tuning(WORKERS, 1024, host, Tuning::locality()); - let threads: Vec<_> = (0..WORKERS) - .map(|id| { - let pool = pool.clone(); - let runner = pool.runner(id); - std::thread::spawn(move || { - // No `mark_current`, because it is private. That is the - // whole of the difference from the test above. - let _ = pool.run(runner); - }) - }) - .collect(); - let executor = Executor::new(pool.clone()); - - let count = threads_that_polled_one_task(&executor); - - pool.shut_down(); - for thread in threads { - let _ = thread.join(); - } - - assert!( - count > 1, - "an unmarked pool kept the task on one worker, which would mean the marker is not the mechanism \ - `local_wakes` depends on and this whole API is unnecessary" - ); - } -} 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()); +}