diff --git a/Cargo.toml b/Cargo.toml index dc3b1c6..215d6fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,9 +7,6 @@ description = "Low-overhead parallel and async work scheduler" repository = "https://github.com/NthTensor/Forte" rust-version = "1.96.0" -# Lie to cargo -links = "forte" - [workspace] resolver = "2" members = ["ci"] diff --git a/build.rs b/build.rs deleted file mode 100644 index 959d5a5..0000000 --- a/build.rs +++ /dev/null @@ -1,6 +0,0 @@ -//! This build-script is intentionally left blank. It exists so that cargo will -//! let us use the `link` field, which ensures that only one version of forte -//! may be used in a given build. This is important to prevent fighting -//! threadpools. - -fn main() {} diff --git a/src/job.rs b/src/job.rs index 51aed0b..6024d5e 100644 --- a/src/job.rs +++ b/src/job.rs @@ -12,11 +12,11 @@ use alloc::collections::VecDeque; use alloc::vec::Vec; use core::any::Any; use core::cell::UnsafeCell; +use core::marker::PhantomData; use core::mem::ManuallyDrop; use core::ptr::NonNull; use crate::latch::Latch; -use crate::platform::*; use crate::thread_pool::Worker; use crate::unwind; @@ -27,6 +27,13 @@ use crate::unwind; /// to be executed, and a function pointer that is capable of executing it. /// /// It is analogous to the chili type `JobShared` or the rayon type `JobRef`. +/// +/// # Panics +/// +/// Execute functions must never allow panics to unwind out of +/// `JobRef::execute`; this is a safety requirement of [`JobRef::new`]. Panics +/// must be caught and either handed to the pool's panic handler or re-emitted +/// when it is safe to do so. pub struct JobRef { /// A non-null pointer to some type-erased data which can be executed as a /// job by the `execute_fn`. This will usually point to either an instance @@ -44,9 +51,15 @@ impl JobRef { /// /// # Safety /// - /// The caller must ensure that `JobRef::execute` will only called on the - /// returned `JobRef` when it would be sound to call `execute_fn` on - /// `job_pointer`. + /// The caller must ensure that: + /// + /// * `JobRef::execute` will only be called on the returned `JobRef` when it + /// is sound to call `execute_fn` on `job_pointer`. Because a `JobRef` is + /// `Send`, `execute` may run on any thread, so this must hold for whatever + /// thread runs it (e.g. `!Send` job data must only be touched on its + /// owning thread). + /// + /// * `execute_fn` will not unwind when called on `job_pointer`. #[inline(always)] pub unsafe fn new( job_pointer: NonNull<()>, @@ -68,37 +81,32 @@ impl JobRef { /// Executes the `JobRef` by passing the execute function on the job pointer. #[inline(always)] pub fn execute(self, worker: &Worker) { - // SAFETY: The caller of `JobRef::new` defines the conditions under - // which this call is sound, and must ensure that this will not be - // called unless these conditions are met. + // SAFETY: `self` is consumed, so this is the one execution of this + // `JobRef`. Its `execute_fn`/`job_pointer` came from `JobRef::new` + // (the only constructor), whose contract requires the caller to ensure + // this call will be sound on whatever thread this is. unsafe { (self.execute_fn)(self.job_pointer, worker) } } } -// SAFETY: Every `JobRef` contains a function pointer and a data pointer. -// Function pointers are always `Send`, but the data pointer may or may not be -// valid for cross-thread access (the value it points to may or may not be -// `Sync`). -// -// However, even when this data is not thread-safe, `JobRef` still needs to be -// `Send`. This is because we need to be able to pass pointers to `!Send` job -// data between threads. For example, if we have a thread that owns a `Future` -// that is `!Send`, and we receive a wakeup notification on an IO polling -// thread, the IO thread must send the owning thread a `JobRef` containing a -// pointer to that `!Send` future. +// SAFETY: A `JobRef` is a function pointer plus a type-erased data pointer. +// Function pointers are `Send`; the data pointer may reference `!Send` data +// (e.g. a `!Send` future whose waker lives on another thread), so `Send` cannot +// be derived structurally. We nonetheless need it since, for instance, passing +// a pointer to `!Send` job data between threads is how an IO thread hands a +// wakeup back to the thread that owns the future. // -// This is only sound because the only method that can actually cause unsound -// cross-thread memory access is `JobRef::execute`. This function is safe, -// because the caller cannot know the soundness requirements of the underlying -// job being pointed to (due to type-erasure). However, `JobRef::new` is -// `unsafe`, and requires the caller to ensure that `execute` will only be -// called if it is correct for the execute function to be called on the -// job_pointer. +// Sending a `JobRef` between threads is sound. The only operation that +// dereferences the data pointer is `JobRef::execute`: `JobRef` has no `Drop` +// impl (so transfer-then-drop leaks the pointee rather than touching it), and +// `id` only reads the pointer's address as a `usize`. `JobRef::new`'s contract +// already requires the caller to ensure `execute` is called only when it is +// sound to do so *on the thread that runs it*, so moving the `JobRef` to +// another thread cannot introduce an unsound access the `new` caller was not +// already obliged to rule out. // -// Since every `JobRef` must be constructed with a call to `new`, it is not -// possible for _entirely safe_ code to violate the `!Send` condition. It is -// unfortunate that the soundness justification has to be squeezed into a single -// function, but thus are the constraints of type-erasure. +// The fields are private and `new` is the only constructor, so that obligation +// always rests on an `unsafe` call. Entirely safe code cannot fabricate a `JobRef`. unsafe impl Send for JobRef {} // ----------------------------------------------------------------------------- @@ -109,54 +117,54 @@ unsafe impl Send for JobRef {} /// conventions. /// /// Note: This is !Sync because of the unsafe cell. +#[derive(Default)] pub struct JobQueue { + /// The queued jobs. + /// + /// Every method briefly takes `&mut *job_refs.get()` for a single + /// `VecDeque` operation. That `&mut` is sound because it is always unique: + /// + /// * `JobQueue` is `!Sync` (it holds an `UnsafeCell`), so only one thread + /// ever accesses it. + /// + /// * No method returns a reference into `job_refs`, so no borrow escapes to + /// overlap a later call. + /// + /// * No user code or `Drop` glue runs while the borrow is live. `JobRef` + /// has no `Drop` impl, and `VecDeque`/`JobRef::id` never call back into + /// the queue. job_refs: UnsafeCell>, } impl JobQueue { - /// Creates a new job queue. - pub fn new() -> JobQueue { - JobQueue { - job_refs: UnsafeCell::new(VecDeque::new()), - } - } - /// Insert a job at the back of the queue (the side with the newest jobs). pub fn push_new(&self, job_ref: JobRef) { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; job_refs.push_back(job_ref); } /// Insert a job at the front of the queue (the side with the oldest jobs). pub fn push_old(&self, job_ref: JobRef) { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; job_refs.push_front(job_ref); } /// Removes the newest job in the queue. pub fn pop_newest(&self) -> Option { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; job_refs.pop_back() } /// Removes the oldest job in the queue. pub fn pop_oldest(&self) -> Option { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; job_refs.pop_front() } @@ -164,10 +172,8 @@ impl JobQueue { /// Attempt to remove the given job-ref from the back of the queue. #[inline(always)] pub fn recover_newest(&self, id: (usize, usize)) -> bool { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; if job_refs.back().map(JobRef::id) == Some(id) { let _ = job_refs.pop_back(); @@ -184,10 +190,8 @@ impl JobQueue { /// the newest jobs). Each chunk is of size `CHUNK_SIZE`. Afterwards, at most /// `CHUNK_SIZE` jobs will remain in the queue. pub fn split(&self) -> Vec> { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; let mut len = job_refs.len(); let num_chunks = len / Self::CHUNK_SIZE; @@ -203,10 +207,8 @@ impl JobQueue { /// Appends a chunk of jobs (expected to be provided by `split`) to the /// queue. Jobs are added to the end (the side with the newest jobs). pub fn append(&self, mut split_refs: VecDeque) { - // SAFETY: `JobQueue` is `!Sync`, so this can only be called from one - // thread. We ensure no other references to the inner value exist by not - // returning any references from this API, making this exclusive access - // sound. + // SAFETY: We have unique access to `job_refs` (see the field + // invariant). let job_refs = unsafe { &mut *self.job_refs.get() }; job_refs.append(&mut split_refs); } @@ -231,8 +233,14 @@ union StackJobData { /// /// This is analogous to the chili type `JobStack` and the rayon type `StackJob`. pub struct StackJob { + /// Job-completion signal. Set only by this job's `execute`. The field is + /// private and never handed out as `&Latch`, so an observed completion + /// means `execute` has run. completed: Latch, data: UnsafeCell>, + /// Makes this type `!Send`. Since `Worker` is also `!Send`, this ensures + /// `new` and `wait` are both given the same worker. + _not_send: PhantomData<*const ()>, } impl StackJob @@ -240,14 +248,15 @@ where F: FnOnce(&Worker) -> T + Send, T: Send, { - /// Creates a new `StackJob` owned by the current worker. + /// Creates a new `StackJob` owned by the given worker. #[inline(always)] - pub fn new(func: F, latch: Latch) -> StackJob { + pub fn new(func: F, worker: &Worker) -> StackJob { StackJob { data: UnsafeCell::new(StackJobData { func: ManuallyDrop::new(func), }), - completed: latch, + completed: worker.new_latch(), + _not_send: PhantomData, } } @@ -263,8 +272,7 @@ where /// * After this call, the `StackJob` will not be moved or dropped until one /// of these conditions is met: /// - /// * (A) A call to `check` on the `StackJob`'s latch returns something - /// other than `Pending`. + /// * (A) `wait` reports that the job has completed. /// /// * (B) The `JobRef` has been dropped without `execute` being called. #[inline(always)] @@ -278,30 +286,38 @@ where // defined by the safety comment for this function. Then: // // * `this` is an aligned pointer to an initialized `StackJob`, - // which will not be invalidated until a `check` on the latch it - // contains returns something other than `Pending`. + // which will not be invalidated until `wait` reports the job complete. // // We created `job_pointer` from a ref to `self`, which is a // `StackJob`, so it must have pointed to an aligned `StackJob` at // some point. // // If the caller allows the `JobRef` to be executed, they must also - // ensure that the pointer will not be invalidated unless a call to - // `check` on the job's `Latch` has returned something other than - // `Pending`. + // ensure the pointer stays valid until `wait` has reported the job + // complete. // // * `StackJob::execute` is called at most once on any `StackJob`. // // The caller ensures only one `JobRef` is ever created for this // `StackJob`. Since `JobRef::execute` consumes that `JobRef`, it // cannot be called multiple times. + // + // `JobRef::new` also requires that the execute function not unwind. + // `StackJob::execute` never unwinds: it catches panics from the + // closure (recording them in the latch, to be re-thrown by whoever + // waits on it), and holds an abort guard while it runs, which turns + // any other panic into an abort. unsafe { JobRef::new(job_pointer, Self::execute) } } - /// Returns a reference to the latch embedded in this stack job. + /// Waits for this job to complete, running other work on `worker` + /// meanwhile. Returns `true` if the job completed with an error. #[inline(always)] - pub fn completion_latch(&self) -> &Latch { - &self.completed + pub fn wait(&self, worker: &Worker) -> bool { + // Note: This really only works if `&Worker` is the same worker that was + // used to create this job. Since `Worker` and `StackJob` are both + // `!Send + !Sync`, we will just assume this is the case. + worker.wait_for(&self.completed) } /// Unwraps the stack job back into a closure. This allows the closure to be @@ -337,19 +353,24 @@ where /// /// # Safety /// - /// This may only be called if a `check` on the enclosed latch has returned - /// `Ok`. + /// May only be called after `wait` has returned `false`, meaning the job completed + /// with a value. #[inline(always)] pub unsafe fn unwrap_output(mut self) -> T { - // Synchronize with the fence in `StackJob::execute`, establishing a - // happens-after relationship with the following read. - fence(Ordering::Acquire); // SAFETY: We have exclusive access to the active union field `output`. // We take `self` by value, so no other `unwrap_*` method can also hold - // it, and `check` returned `Ok`, so `execute` ran. Since `execute` runs - // at most once, it is no longer running and cannot alias `data`. + // it, and `wait` returned `false`, so `execute` ran (by the invariant on + // `completed`, only `execute` sets the latch). Since `execute` runs at + // most once, it is no longer running and cannot alias `data`. // - // `output` is the active variant because `check` returning `Ok` means + // The caller observed completion via `wait`, which loops on + // `Latch::check`, an `Acquire` load that synchronizes with the + // `Release` store in the `Latch::set` call inside `execute`. Since + // `execute` writes the `output` field before that store, that write + // must happen-before this read. So there can be no data-race on this + // load. + // + // `output` is the active variant because `wait` returning `false` means // `execute` called `set` with `error_flag == false`, which always // follows a write of the `output` field, after which the union is not // written again. @@ -363,19 +384,23 @@ where /// /// # Safety /// - /// This may only be called if a `check` on the enclosed latch has returned - /// `Error`. + /// May only be called after `wait` has returned `true`, meaning the job completed + /// with an error. #[inline(always)] pub unsafe fn unwrap_error(mut self) -> Box { - // Synchronize with the fence in `StackJob::execute`, establishing a - // happens-after relationship with the following read. - fence(Ordering::Acquire); // SAFETY: We have exclusive access to the active union field `error`. // We take `self` by value, so no other `unwrap_*` method can also hold - // it, and `check` returned `Error`, so `execute` ran. Since `execute` - // runs at most once, it is no longer running and cannot alias `data`. + // it, and `wait` returned `true`, so `execute` ran (by the invariant on + // `completed`, only `execute` sets the latch). Since `execute` runs at + // most once, it is no longer running and cannot alias `data`. + // + // The caller observed completion via `wait`, which loops on + // `Latch::check`, an `Acquire` load that synchronizes with the + // `Release` store in the `Latch::set` call inside `execute`. Since + // `execute` writes the `error` field before that store, that write must + // happen-before this read. So there can be no data-race on this load. // - // `error` is the active variant because `check` returning `Error` means + // `error` is the active variant because `wait` returning `true` means // `execute` called `set` with `error_flag == true`, which always // follows a write of the `error` field, after which the union is not // written again. @@ -392,8 +417,7 @@ where /// The caller must ensure that: /// /// * `this` is an aligned pointer to an initialized `StackJob`, which - /// will not be invalidated until a `check` on the latch it contains - /// returns something other than `Pending`. + /// will not be invalidated until `wait` reports the job complete. /// /// * This function is called at most once on any `StackJob`. #[inline(always)] @@ -401,9 +425,12 @@ where // SAFETY: The pointer `this` is non-null, aligned, and the caller // ensures it points to an initialized `StackJob`. // - // `StackJobs` are always accessed immutably except for `unwrap_func`, + // `StackJob`s are always accessed immutably except for `unwrap_func`, // `unwrap_output`, and `unwrap_error`. The caller ensures these will not - // race this call, so the pointer is valid for immutable access. + // race this call, so the pointer is valid for immutable access. The + // `&Self` stays valid for the whole body: the only field mutated below + // is `data`, an `UnsafeCell`, so forming `&mut` from `data.get()` does + // not alias-violate this shared reference. let this = unsafe { this.cast::().as_ref() }; // Create an abort guard. If the closure panics, this will convert the // panic into an abort. Doing so prevents use-after-free for other @@ -412,10 +439,12 @@ where // Run the function and record the result. Produces a boolean flag that // is true in the event of a panic. let error_flag = { - // SAFETY: Only `unwrap_func`, `unwrap_output` and `unwrap_error` - // access `data`. Due to their individual safety contracts, they can - // only be called in a way that will not race with this function, so - // we must have unique access. + // SAFETY: `data` is an `UnsafeCell`, so the outer `&Self` does not + // alias it. `execute` runs at most once (per this function's + // contract), so no other `execute` races it, and + // `unwrap_func`/`unwrap_output`/`unwrap_error` cannot race it (per + // their contracts). This `&mut` is therefore unique for its + // lifetime. let data_ref = unsafe { &mut *this.data.get() }; // SAFETY: Each `StackJob` is constructed using field `func`, and // this is the only place we write to the union after construction. @@ -441,32 +470,28 @@ where } } }; - // This synchronizes with the `Acquire` fence within `unwrap_output` / - // `unwrap_error`, establishing a happens-before relationship that makes - // the preceding write to the `output`/`error` union field visible to - // the reader. + // This publishes the write to `data` with a `Release` store, so any + // thread that later observes completion through `wait` can load `data` + // without a race. // - // This is required because latches do not synchronize memory. - fence(Ordering::Release); - // SAFETY: This casts a reference to a raw pointer, which means the - // pointer must be aligned, non-null, and point to an initialized latch. + // SAFETY: The pointer is derived from `&this.completed`, a valid + // reference, so it is non-null, aligned, and points to an initialized + // `Latch`. // // We also meet Variant 2 of the `set` safety contract: // // * The latch has not been `set` since it was created or last `reset`, // and calls to `set` do not race. // - // This is the only place where this latch is set, and the caller - // ensures this is called at most once. Therefore `set` cannot have - // been called already, and there can be no other calls to `set` that - // would race with this one. + // By the invariant on `completed`, only this `execute` sets the latch, + // and it runs at most once, so no prior or racing `set` exists. // // * The latch will not be dropped or moved until after `check` returns // something other than `Pending`. // - // The caller ensures that this `StackJob` is not dropped until a - // `check` on the latch returns something other than `Pending`, and - // nothing removes the latch from the `StackJob`. + // The owner does not drop the `StackJob` until `wait` returns, which + // happens only once the latch is set, and nothing removes the latch + // from the `StackJob`. unsafe { Latch::set(&this.completed, error_flag) }; // Forget the abort guard, re-enabling panics. core::mem::forget(abort_guard); @@ -489,8 +514,12 @@ where F: FnOnce(&Worker), { /// Allocates a new `HeapJob` on the heap. + /// + /// # Safety + /// + /// The caller must ensure that the function `f` does not unwind. #[inline(always)] - pub fn new(f: F) -> Box { + pub unsafe fn new(f: F) -> Box { Box::new(HeapJob { f }) } @@ -537,9 +566,15 @@ where // function, so `JobRef::execute` and hence `HeapJob::execute` can // only be called during that interval. // - // * Accessing `f` will not violate a `!Send` requirement. + // * Using or dropping `f` will not violate a `!Send` requirement. + // + // This function's `!Send` safety clause confines `execute` (the only + // code that uses or drops `f`) to the construction thread. An + // unexecuted `JobRef` leaks the `HeapJob` (no `Drop` impl), so `f` is + // never dropped elsewhere. // - // This is ensured by the executability condition. + // `JobRef::new` also requires that the execute function not unwind. + // This is left to the caller of `HeapJob::new` to ensure. unsafe { JobRef::new(job_pointer, Self::execute) } } @@ -549,7 +584,8 @@ where /// /// The caller must ensure that: /// - /// * `this` is an aligned pointer to an initialized `HeapJob`. + /// * `this` was produced by `Box::into_raw` on a `Box>`, and this + /// call takes ownership of that allocation. /// /// * This function is called at most once on any `HeapJob`. /// @@ -559,13 +595,15 @@ where /// the `HeapJob` was constructed. #[inline(always)] unsafe fn execute(this: NonNull<()>, worker: &Worker) { - // SAFETY: The caller ensures that: - // - // * `this` was created by `Box::into_raw`. - // - // * This function is called at most once. + // SAFETY: The first clause satisfies the creation and + // allocation-ownership requirements of `Box::from_raw`, and the + // at-most-once clause rules out a prior reclaim. let this = unsafe { Box::from_raw(this.cast::().as_ptr()) }; - // Run the job. + // Run the job. Reading whatever `f` captures is sound because they are + // still live (per the third clause of the safety contract). `f` runs + // and is dropped on this thread, which the fourth clause requires to be + // the construction thread when `f` is `!Send`. `f` does not unwind, per + // `HeapJob::new`'s contract. (this.f)(worker); } } diff --git a/src/latch.rs b/src/latch.rs index 3c7cca9..668cd0f 100644 --- a/src/latch.rs +++ b/src/latch.rs @@ -13,7 +13,7 @@ use crate::platform::*; // States /// The default state of a latch is `LOCKED`. When in the locked state, `check` -/// returns `false` and `wait` blocks. +/// returns `Pending` and `wait` blocks. const LOCKED: u32 = 0b000; /// The latch enters the `SIGNAL` state when it is set (with error flag false). @@ -36,7 +36,7 @@ const ASLEEP: u32 = 0b100; /// The latch begins as *unset* (In the `LOCKED` state), and can later be *set* /// by any thread (entering the `SIGNAL`) state. Each latch is "owned" by a /// single thread at a time; other threads may set the latch, but only the -/// owning thread may wait on it. +/// owning thread may call `wait` or `check`. /// /// The general idea and spirit for latches (as well as some of the /// documentation) is due to rayon. However the implementation is specific to @@ -44,9 +44,15 @@ const ASLEEP: u32 = 0b100; /// /// ## Memory Ordering /// -/// Latches _do not synchronize memory_. They are only used for signaling. If -/// the thread that sets a latch wishes to transmit a value to the thread -/// waiting for that latch, explicit fences must be used. +/// `set` stores the latch state with `Release` and `check` loads it with +/// `Acquire`. A `check` that observes a set latch therefore *happens-after* +/// the `set`. +/// +/// Therefore, after the owning thread observes a signal through `check`: +/// +/// 1. It may reclaim the latch (drop or `reset` it). +/// +/// 2. It may access any writes ordered before the call to `set`. pub struct Latch { /// Holds the internal state of the latch. This tracks if the latch has been /// set or not. @@ -55,9 +61,14 @@ pub struct Latch { semaphore: &'static Semaphore, } +/// The result of a [`Latch::check`]: whether the latch has been set, and if so +/// whether it was set with an error. pub enum Status { + /// The latch has not been set. Pending, + /// The latch was set without an error. Ok, + /// The latch was set with an error. Error, } @@ -70,10 +81,16 @@ impl Latch { } } - /// Checks to see if the latch has been set. Returns true if it has been. + /// Checks to see if the latch has been set. + /// + /// # Memory Ordering + /// + /// A call to `check` that observes a signal establishes a happens-after + /// relationship with the `set` that produced the signal, which can be used + /// to safely reclaim the latch or transmit values between threads. #[inline(always)] pub fn check(&self) -> Status { - match self.state.load(Ordering::Relaxed) { + match self.state.load(Ordering::Acquire) { SIGNAL => Status::Ok, ERROR => Status::Error, _ => Status::Pending, @@ -87,10 +104,8 @@ impl Latch { /// /// # Memory Ordering /// - /// This does not synchronize memory. To synchronize memory with the thread - /// setting the latch, call `fence(Ordering::Acquire)` after this function. - /// The other thread must issue a corresponding `fence(Ordering::Release)` - /// call. + /// A call to `wait` does not synchronize memory, and so must be used in + /// conjunction with `check` (see the type-level "Memory Ordering" section). #[cold] pub fn wait(&self, seat_bitmask: u32, waiting_bitmask: &'static AtomicU32) { // First, check if the latch has been set. @@ -120,9 +135,10 @@ impl Latch { /// /// # Memory Ordering /// - /// This does not synchronize memory. To synchronize memory with the waiting - /// thread, call `fence(Ordering::Release)` before this function. The other - /// thread must issue a corresponding `fence(Ordering::Acquire)` call. + /// A call to `check` that observes the signal produced by this `set` + /// establishes a happens-after relationship. Stores on this thread made + /// before this call will be visible to the owning thread after `check` + /// returns a status other than `Pending`. /// /// # Safety /// @@ -136,7 +152,7 @@ impl Latch { /// /// 2. The latch has not been `set` since it was created or last `reset`, /// calls to `set` do not race, and the latch will not be dropped or - /// moved until after `check` returns something other than `Pending`. + /// moved until after `check` returns a status other than `Pending`. #[inline(always)] pub unsafe fn set(latch: *const Latch, error: bool) { // First we store a reference to the semaphore (which is 'static) so @@ -163,13 +179,16 @@ impl Latch { // In the event of a race with `wait`, this may cause `wait` to return. // Otherwise the other thread will sleep within `wait`. // - // SAFETY: The latch is still valid to access immutably, following the - // same logic as above. + // SAFETY: The latch must still be valid to dereference, up to the point + // where we begin this store, by the same reasoning as above. // - // NOTE: This store will mean `check` no longer returns `Pending`, - // invalidating the argument in Variant 2. The latch pointer therefore - // may become dangling after this line. - unsafe { (*latch).state.store(state, Ordering::Relaxed) }; + // The `Release` store synchronizes with the `Acquire` load in `check` + // to supply the second half of the Variant 2 argument: the latch does + // not merely live *up to* this store, the store is also ordered + // *before* any later free. The owner may free the latch as soon as + // `check` reports non-`Pending`, and that check is ordered after this + // store, so there can be no data race. + unsafe { (*latch).state.store(state, Ordering::Release) }; // Finally we try to signal the target thread on it's semaphore, just in // case it missed the notification and is currently waiting. This // guarantees that the other thread will make progress. @@ -201,7 +220,7 @@ pub struct Semaphore { impl Semaphore { /// Creates a new signal. - pub const fn new() -> Self { + pub(crate) const fn new() -> Self { Semaphore { state: AtomicU32::new(LOCKED), } diff --git a/src/lib.rs b/src/lib.rs index 027bf56..4b25d44 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -178,7 +178,6 @@ //! * *Block on.* Waits for a future to complete (outside of an async context). //! * *Broadcast.* Runs the same operation across all workers. //! -//! #![no_std] @@ -209,10 +208,27 @@ pub struct FnOnceMarker(); #[doc(hidden)] pub struct FutureMarker(); +// ----------------------------------------------------------------------------- +// Trait sealing + +/// Seals [`Spawn`] and [`SpawnScoped`] so they cannot be implemented outside +/// this crate. +mod sealed { + use core::future::Future; + + use crate::FnOnceMarker; + use crate::FutureMarker; + use crate::Worker; + + pub trait Sealed {} + + impl Sealed for F {} + impl Sealed for Fut {} +} + // ----------------------------------------------------------------------------- // Top-level exports -pub use latch::Latch; pub use scope::Scope; pub use scope::SpawnScoped; pub use thread_pool::Broadcast; @@ -233,6 +249,25 @@ pub use thread_pool::spawn; pub use thread_pool::spawn_broadcast; pub use thread_pool::spawn_on; +// ----------------------------------------------------------------------------- +// Internal primitives + +/// Low-level scheduling primitives for building custom executors on forte. +/// +/// These are fully public but `#[doc(hidden)]`. You can use this API, but it's +/// not advisable, and most of it is very unsafe. +#[doc(hidden)] +pub mod internals { + pub use crate::job::HeapJob; + pub use crate::job::JobQueue; + pub use crate::job::JobRef; + pub use crate::job::StackJob; + pub use crate::latch::Latch; + pub use crate::latch::Semaphore; + pub use crate::latch::Status; + pub use crate::thread_pool::Scheduler; +} + // ----------------------------------------------------------------------------- // Platform Support @@ -240,23 +275,19 @@ pub use thread_pool::spawn_on; // Currently there are no alternative implementations, but there may be in // future. mod platform { - + pub use alloc::sync::Arc; pub use core::sync::atomic::AtomicBool; pub use core::sync::atomic::AtomicPtr; pub use core::sync::atomic::AtomicU32; pub use core::sync::atomic::AtomicUsize; pub use core::sync::atomic::Ordering; pub use core::sync::atomic::fence; - - pub use alloc::sync::Arc; + use std::sync::LazyLock; pub use std::sync::Mutex; pub use std::thread::Builder as ThreadBuilder; pub use std::thread::JoinHandle; - pub use std::thread_local; - pub use std::thread::available_parallelism; - - use std::sync::LazyLock; + pub use std::thread_local; pub struct Lazy(LazyLock); impl Lazy { diff --git a/src/scope.rs b/src/scope.rs index 7435843..8e9f3bf 100644 --- a/src/scope.rs +++ b/src/scope.rs @@ -232,13 +232,17 @@ impl<'scope, 'env> Scope<'scope, 'env> { // operation. self.count.fetch_add(participants as u32, Ordering::Relaxed); // (*) + // The jobs can outlive this stack frame, so they cannot borrow `f`; + // they share ownership of it through an `Arc` instead. + let f = Arc::new(f); + // Create a new job for each member let member_data = self.thread_pool.get_member_data(); for (i, member_index) in members.iter_bits().enumerate() { - let func = &f; + let func = Arc::clone(&f); let scope = self; let op = move |worker: &Worker| { - // Run the job + // Run an instance of the job job. let result = unwind::halt_unwinding(|| { func(Broadcast { worker, @@ -251,6 +255,17 @@ impl<'scope, 'env> Scope<'scope, 'env> { if let Err(err) = result { scope.store_panic(err); }; + // Dropping the last handle drops `f`, whose drop glue may read + // `'scope` data. That data is only guaranteed to live until the + // scope's counter reaches zero, so the handle must be dropped + // before this job's count is removed. + // + // Because `f` is user code its drop glue may also panic. The + // panic is caught and stored on the scope so that it cannot + // unwind out of this job. + if let Err(err) = unwind::halt_unwinding(|| drop(func)) { + scope.store_panic(err); + } // SAFETY: This corresponds to one of the increments performed in // a batch with the `fetch_add` at the start of this function. // It was incremented `p` times, and this will be called `p` @@ -258,32 +273,40 @@ impl<'scope, 'env> Scope<'scope, 'env> { unsafe { scope.remove_reference() }; }; - let job = HeapJob::new(op); + // SAFETY: `HeapJob::new` may only be passed functions that do not + // unwind. The function `op` catches panics from the broadcast + // operation `func`, and makes no other calls that can panic. + let job = unsafe { HeapJob::new(op) }; // SAFETY: `HeapJob::into_job_ref` requires: // // * The `JobRef` will not outlive any of the items closed over by // the `op`. // - // The only non-copy captures are `scope: &'scope Scope` and - // `func: &&F + 'scope`, so we must show that the `JobRef` will - // not outlive `'scope`. + // The captures of `op` are `i` and `participants` (both `Copy`), + // `scope: &'scope Scope`, and `func: Arc`. // - // This is ensured via the scope's lifetime extension logic: we - // incremented the scope's counter on the line marked with (*), - // and we know the scope will not complete (extending the lifetime - // of `'scope`) until there is a corresponding call to - // `remove_reference()`. + // The job owns `func`, so the `JobRef` cannot outlive it. The + // `Arc` keeps the value `f` alive until the last handle is + // dropped; each handle is dropped inside `op` when a job runs, + // when this function's own handle goes out of scope, or leaks + // along with an unexecuted `JobRef`. // - // All `n` added references are not removed until the op has run - // `n` times on `n` threads, so the `op` cannot outlive the data - // it borrows. + // The capture `scope`, and any `'scope` data borrowed by `f`, + // are kept alive by the scope's lifetime extension logic: we + // incremented the scope's counter once per job on the line + // marked with (*), and the scope will not complete (extending + // the lifetime of `'scope`) until there is a corresponding call + // to `remove_reference()`. Each job holds its increment until + // after it has called `func` and dropped its `Arc` handle, so + // neither the `Scope` nor anything `f` borrows can be freed + // while a job might still use (or drop) them. // // * If `op` is `!Send` then `JobRef::execute` will only be called // on this thread. // - // `op` is unconditionally `Send`. Since `F` and `Scope` are - // `Sync,` the enclosed references `&Scope` and `&&F` are `Send`. + // `op` is unconditionally `Send`: `F: Send + Sync` makes + // `Arc: Send`, and `Scope: Sync` makes `&Scope: Send`. let job_ref = unsafe { job.into_job_ref() }; member_data.broadcasts[member_index].push(job_ref); member_data.semaphores[member_index].signal(); @@ -312,19 +335,39 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// `add_reference`, a direct `fetch_add` on the underlying counter, or the /// implicit initial increment the scope starts with. unsafe fn remove_reference(&self) { - let counter = self.count.fetch_sub(1, Ordering::Relaxed); + // This is `Release` so that this job's accesses to the scope are + // ordered before the decrement. + // + // The std-lib's `Arc` has the same basic pattern. + let counter = self.count.fetch_sub(1, Ordering::Release); if counter == 1 { + // This is the final decrement. + // + // The `Acquire` fence synchronizes with the release sequence formed + // by every other job's `Release` decrement of `count`, so all of + // their accesses to the scope happen-before this point. Combined + // with the `Release`/`Acquire` handshake in + // `Latch::set`/`Latch::check`, this orders every job's use of the + // scope before the owner returns from `complete` and frees it. + // Without it, remote jobs' atomic ops on `count` (and their reads + // and writes of `panic`) would race the owner's deallocation. + fence(Ordering::Acquire); // Alerts the owning thread that the scope has completed. // // This should never panic, because the counter can only go to zero // once, when the scope has been dropped and all work has been // completed. // - // SAFETY: The owning thread must call `Scope::complete` before - // dropping any `Scope`, and `Scope::complete` does not return until - // the latch is set, which happens only here, after the count - // reaches zero. Therefore, the `completed` field of this `Scope` - // must still be a live latch. + // SAFETY: We meet Variant 2 of the `Latch::set` safety contract: + // + // * The count reaches zero exactly once, so this branch runs + // exactly once, and this is the only place the latch is set. So + // no other call to `set` can race with this one, and `set` cannot + // have already been called. + // + // * This latch will not be dropped until `complete` returns, and + // `complete` will not return until it observes the signal we are + // about to send. unsafe { Latch::set(&self.completed, false) }; } } @@ -332,14 +375,15 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// Stores a panic so that it can be propagated when the scope is complete. /// If called multiple times, only the first panic is stored, and the /// remainder are dropped. + /// + /// This function never unwinds, but may abort in circumstances. #[cold] fn store_panic(&self, err: Box) { + // Create an abort guard that will prevent this function from unwinding. + let abort_guard = AbortOnDrop; // Check if the panic pointer has already been set. This lets us avoid // allocating a second time, and means we can immediately drop the panic // we have just been passed. - // - // Dropping this panic may itself trigger a panic, but this will simply - // trigger the scope's abort guard, causing an abort rather than UB. if self.panic.load(Ordering::Relaxed).is_null() { let nil = ptr::null_mut(); let err_ptr = Box::into_raw(Box::new(err)); @@ -360,19 +404,24 @@ impl<'scope, 'env> Scope<'scope, 'env> { Ordering::Release, Ordering::Relaxed, ) - .is_ok() + .is_err() { - // Ownership is now transferred into the panic field. - } else { - // Another panic raced in ahead of us, so we need to drop this one. + // Another panic raced in ahead of us, so we need to drop this + // error. Dropping the payload may itself panic, but this will + // trigger the abort guard instead of unwinding. // - // SAFETY: This was created by `Box::into_raw` just above. It is - // possible that this will panic, because it's a `Box`, - // however in the worst case this will simply trigger the - // scope's abort guard, causing an abort rather than UB. - let _: Box<_> = unsafe { Box::from_raw(err_ptr) }; + // SAFETY: This was created by `Box::into_raw` just above. + let err = unsafe { Box::from_raw(err_ptr) }; + drop(err); } + } else { + // A panic is already stored, so this payload is dropped. Dropping + // it may itself panic, but this will trigger the abort guard + // instead of unwinding. + drop(err); } + // At this point we can allow panics to unwind again. + core::mem::forget(abort_guard); } /// Propagates any panic captured while the scope was executing. @@ -423,6 +472,8 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// A trait for types that can be spawned onto a [`Scope`]. /// +/// This trait is sealed and cannot be implemented outside of forte. +/// /// It is implemented for: /// /// * Closures that satisfy `for<'worker> FnOnce(&'worker Worker) + Send + 'scope`. @@ -451,7 +502,7 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// }); /// ``` /// Hopefully rustc will fix this type inference failure eventually. -pub trait SpawnScoped<'scope, M>: 'scope { +pub trait SpawnScoped<'scope, M>: crate::sealed::Sealed + 'scope { /// Similar to [`spawn`][crate::Worker::spawn] but adds the work to a /// [`Scope`]. This work will be polled to completion some-time before the /// scome completes, and may borrow data that outlives the scope. @@ -481,15 +532,21 @@ where ) { // Create a job to execute the spawned function in the scope. let scope_ptr = ScopePtr::new(scope); - let job = HeapJob::new(move |worker| { + + // Create the spawned operation + let op = move |worker: &Worker| { // Catch any panics and store them on the scope. let result = unwind::halt_unwinding(|| self(worker)); if let Err(err) = result { scope_ptr.store_panic(err); }; drop(scope_ptr); - }); + }; + // SAFETY: `HeapJob::new` may only be called with functions that do not + // unwind. The function `op` uses `halt_unwinding` to ensure this is the + // case. + let job = unsafe { HeapJob::new(op) }; // SAFETY: `HeapJob::into_job_ref` requires: // // * The `JobRef` will not outlive any of the items closed over by @@ -500,14 +557,17 @@ where // will not outlive `'scope`. // // This is ensured via the scope's lifetime extension logic: the scope - // will not complete so long as the `scope_ptr` is held, extending the - // lifetime of `'scope` until after `self` the job executes and is + // will not complete so long as the `scope_ptr` is held, extending + // `'scope` until after the job has executed and `self` has been // dropped. // // * If `op` is `!Send` then `JobRef::execute` will only be called // on this thread. // - // `op` is unconditionally `Send`. + // Since its only non-`Copy` captures are an instance of `F` and the + // `Send` `scope_ptr`, `op` is `Send` iff `F` is. If `F: Send` the + // clause is vacuous. If `F: !Send`, `op` is `!Send`, and the caller's + // `spawn_scoped` safety contract confines execution to this thread. let job_ref = unsafe { job.into_job_ref() }; // Send the job to a queue to be executed. @@ -594,7 +654,12 @@ struct ScopeFutureJob<'scope, 'env, S: for<'w> Scheduler<'w>, Fut> { /// to allow mutable access within an `Arc`. The `state` field acts as a /// kind of mutex that ensures exclusive access by preventing the job from /// being queued or executed multiple times simultaneously. - future: UnsafeCell, + /// + /// The future is wrapped in `ManuallyDrop` so the `Arc`'s drop glue never + /// runs its destructor. That lets a waker free the allocation on any thread + /// while `poll` remains the sole place the future is dropped, on its + /// confined thread. + future: UnsafeCell>, /// A scope pointer. This allows the job to interact with the scope, and /// also keeps the scope alive until the job is dropped. scope_ptr: ScopePtr<'scope, 'env>, @@ -627,7 +692,7 @@ where ) -> Arc { let scope_ptr = ScopePtr::new(scope); Arc::new(Self { - future: UnsafeCell::new(future), + future: UnsafeCell::new(ManuallyDrop::new(future)), scope_ptr, // The job starts in the WOKEN state because we always queue it // after creating it. @@ -666,6 +731,9 @@ where // `job_pointer` without decrementing it. Therefore, when `execute` // is called, there will still be a strong reference for it to // consume. + // + // `JobRef::new` also requires that `poll` not unwind. See the + // doc-comment for `poll` for an argument as to why this is the case. unsafe { JobRef::new(job_pointer, Self::poll) } } @@ -673,8 +741,8 @@ where /// handing that to the job's own `scheduler`. The `worker` param should be /// the same as calling `Worker::with_current`. fn schedule(self: Arc, worker: Option<&Worker>) { - // Clone a strong reference so the allocation — and therefore the - // `scheduler` field — stays alive for the entire `schedule` call. + // Clone a strong reference so the allocation, and therefore the + // `scheduler` field, stays alive for the entire `schedule` call. // // `into_job_ref` consumes `self` via `Arc::into_raw`, moving our strong // reference into the `JobRef` *without* decrementing the count. The act @@ -690,6 +758,12 @@ where /// Polls the future. /// + /// # Panics + /// + /// This function does not unwind. Panics that occur while running the job + /// are caught with `catch_unwind` and stored on the scope. Other panics, + /// such as those emitted by drop-glue, may cause aborts. + /// /// # Safety /// /// `this` must be a pointer produced by `Arc::into_raw` on an `Arc`. @@ -726,7 +800,10 @@ where // resources from being freed. // // * `drop` decrements the reference count, potentially allowing the job - // to be freed. + // to be freed. A waker is `Send`, so this may run on any thread, but + // freeing the job does not run the future's destructor (it is + // `ManuallyDrop`, dropped only by `poll` on the confined thread); the + // remaining fields are `Send`, so freeing anywhere is sound. // let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker)) }; @@ -742,8 +819,8 @@ where // is the case, and abort the program if it is not. // // We use Acquire ordering here to ensure we have the current state of - // the future. This synchronizes with the fence in the Poll::Pending - // branch. + // the future. This synchronizes with the previous poll's `Release` + // `compare_exchange` (see the SAFETY comment below). if this.state.swap(LOCKED, Ordering::Acquire) != WOKEN { // We abort because this function needs to be panic safe. abort(); @@ -763,61 +840,70 @@ where // swap, and cause an abort. Exclusive access is therefore // guaranteed. // - // In the event that `poll` has been called previously, the `Acquire` - // ordering synchronizes with the call to + // Exclusive access is not enough on its own: a previous poll's writes + // to the future must also be visible here, with no data race. There + // are two cases, by how this poll was scheduled: // - // fence(Ordering::Release) + // - Woken after returning `Pending`: the previous poll published its + // writes with the `Release` `compare_exchange(LOCKED, READY)` below. + // A `wake` reads `READY` and reschedules; its `swap(WOKEN, Relaxed)` + // is an RMW, so it stays within that `Release` store's release + // sequence. This `swap(LOCKED, Acquire)` reads `WOKEN` from it and + // therefore synchronizes-with the previous poll's `Release` store, + // ordering that poll's writes before ours. // - // later in this function. This ensures that we are not racing with - // another mutable reference to the same value. + // - Woken while running: the previous poll's `compare_exchange` failed + // (the state was already `WOKEN`), so that poll rescheduled the job + // itself. Its writes are sequenced before its `schedule` push, which + // the job queue orders (`Release`) before this poll's dequeue + // (`Acquire`). // // * The future will not move. // // The future does not move, because it is stored in a field within an // `Arc`, which has a stable heap-allocated address. - let future = unsafe { Pin::new_unchecked(&mut *this.future.get()) }; + let future = unsafe { Pin::new_unchecked(&mut **this.future.get()) }; // Create a new context from the waker, and poll the future. let mut cx = Context::from_waker(&waker); let result = unwind::halt_unwinding(|| future.poll(&mut cx)); // Update the job state depending on the outcome of polling the future. + // + // Out of an abundance of caution, we add an abort guard here. + let abort_guard = AbortOnDrop; match result { // The job completed without panicking. Ok(Poll::Ready(())) => { - // Drop the job without rescheduling it, leaving it in the - // LOCKED or WOKEN state so that it cannot be rescheduled. + // Drop the future in place. State stays LOCKED, so it is never + // re-polled and this runs exactly once. + // + // SAFETY: The LOCKED state gives us exclusive access. Dropping + // a `!Send` future here is sound because `poll` only runs where + // the scheduler enqueues it, and `spawn_local` (the only entry + // for `!Send` work) confines every (re)schedule to the origin + // worker's thread, so this drop happens on that thread. + unsafe { ManuallyDrop::drop(&mut *this.future.get()) }; + drop(this); } // The job is still pending, and has not yet panicked. Ok(Poll::Pending) => { - // Try to set the state back back idle so other threads can + // Try to set the state back to `READY` so other threads can // schedule it again. This will only fail if the job was woken // while running, and is already in the WOKEN state. // - // If successful, this effectively releases our exclusive - // ownership of the future. + // On success this publishes the poll's writes with `Release`; a + // later poll's opening `Acquire` swap synchronizes with it (see + // the SAFETY comment above), ordering this poll before the next. let rescheduled = this .state .compare_exchange( LOCKED, READY, - Ordering::Relaxed, + Ordering::Release, Ordering::Relaxed, ) .is_err(); - // Emit a fence here, which synchronizes with the `Acquire` swap - // at the start of this function to ensure that the next thread - // to poll this future will observe the most recent version of - // it. - // - // A fence is required here because the write to `state` that - // establishes the happens-before relationship may be caused by - // either (a) the `compare_exchange` call above, or (b) the - // `swap` call in `wake`. - // - // This fence lets `wake` use `Relaxed` ordering, and upgrades - // it to `Release` only when necessary. - fence(Ordering::Release); // If the job was woken while running, it should be queued // immediately. Conveniently, we know the state will already be // WOKEN, so we can leave it as it is. @@ -829,31 +915,47 @@ where // NOTE: Eventually we may want to chose not to re-evaluate // this sometimes, to break out of infinite async loops. this.schedule(Some(worker)); + } else { + drop(this); } } // The job panicked. Store the panic in the scope so it can be // resumed later. - Err(err) => this.scope_ptr.store_panic(err), + Err(err) => { + this.scope_ptr.store_panic(err); + // Drop the future here, on its confined thread; see the Ready + // branch. + // + // SAFETY: The LOCKED state gives us exclusive access. + unsafe { ManuallyDrop::drop(&mut *this.future.get()) }; + drop(this); + } } + core::mem::forget(abort_guard); - // At this point, if we have not converted `this` into a job-ref and - // queued it, it will be dropped. If no wakers for this task are being - // held, then this will cause the reference counter to decrement to zero - // and free the task. + // In the branches above where the task is dropped (rather than + // converted into a job-ref and queued), the drop decrements the + // reference count, freeing the task if no wakers for it are being + // held. // - // The alternative, if there are no wakers being held for the task, is - // that the task will never wake and the scope will deadlock. + // I view this as preferable to the alternative: if no wakers are held + // for the task, the task will never wake and the scope will wait for + // its latch indefinitely, and deadlock. } /// Creates a new `RawWaker` from the provided pointer. /// /// # Safety /// - /// Must be called with a pointer created by calling `Arc::into_raw` on an - /// instance of `Arc` that is still alive. + /// `this` must be a pointer produced by `Arc::into_raw` on an `Arc`. + /// + /// This call borrows the reference: the caller must keep one strong + /// reference count alive for the duration of the call. It does not consume + /// a count. unsafe fn clone_as_waker(this: *const ()) -> RawWaker { - // SAFETY: This is called on a pointer created by `Arc::into_raw` on an - // instance of `Arc`. + // SAFETY: `this` came from `Arc::into_raw` and the caller keeps a strong + // reference alive (per the contract), so incrementing the count for the + // cloned waker is sound. unsafe { Arc::increment_strong_count(this.cast::()) }; RawWaker::new(this, &Self::VTABLE) } @@ -862,13 +964,22 @@ where /// /// # Safety /// - /// Must be called with a pointer created by calling `Arc::into_raw` on an - /// instance of `Arc` that is still alive. + /// `this` must be a pointer produced by `Arc::into_raw` on an `Arc`. + /// + /// This call takes ownership of exactly one strong reference count for that + /// allocation, consuming it via `Arc::from_raw` internally. The caller must + /// hold "ownership" of one such strong reference. unsafe fn wake(this: *const ()) { - // SAFETY: This is called on a pointer created by `Arc::into_raw` on an - // instance of `Arc`. + // SAFETY: `this` came from `Arc::into_raw` and the caller transfers + // ownership of one strong count (per the contract), so `from_raw` may + // reclaim it. let this = unsafe { Arc::from_raw(this.cast::()) }; + // This wake publishes no data of its own; it only needs the READY -> + // WOKEN transition to be atomic so exactly one waker enqueues the job. + // Hence the `Relaxed` ordering. The future's state is ordered by the + // poll's `Release` store and the enqueue/dequeue handoff, not by this + // swap. if this.state.swap(WOKEN, Ordering::Relaxed) == READY { Worker::with_current(|worker| this.schedule(worker)); } @@ -878,17 +989,25 @@ where /// /// # Safety /// - /// Must be called with a pointer created by calling `Arc::into_raw` on an - /// instance of `Arc` that is still alive. + /// `this` must be a pointer produced by `Arc::into_raw` on an `Arc`. + /// + /// This call borrows the reference: the caller must keep one strong + /// reference count alive for the duration of the call. It does not consume + /// a count. unsafe fn wake_by_ref(this: *const ()) { // We use manually drop here to prevent us from consuming the arc on // drop. This functions like an `&Arc` rather than an `Arc`. // - // SAFETY: This is called on a pointer created by `Arc::into_raw` on an - // instance of `Arc`. + // SAFETY: `this` came from `Arc::into_raw` and the caller keeps a strong + // reference alive (per the contract). Wrapping in `ManuallyDrop` borrows + // that count rather than consuming it. let this = unsafe { ManuallyDrop::new(Arc::from_raw(this.cast::())) }; + // As in `wake`, this publishes no data of its own; it only needs the + // READY -> WOKEN transition to be atomic. Hence, again, the `Relaxed` + // ordering. The future's state is ordered by the poll's `Release` store + // and the enqueue/dequeue handoff, not by this swap. if this.state.swap(WOKEN, Ordering::Relaxed) == READY { // Clone the waker, convert it into a job-ref and queue it. let this = ManuallyDrop::into_inner(this.clone()); @@ -900,14 +1019,18 @@ where /// /// # Safety /// - /// Must be called with a pointer created by calling `Arc::into_raw` on an - /// instance of `Arc` that is still alive. + /// `this` must be a pointer produced by `Arc::into_raw` on an `Arc`. + /// + /// This call takes ownership of exactly one strong reference count for that + /// allocation, releasing it via `Arc::decrement_strong_count` internally. + /// The caller must hold "ownership" of one such strong reference. unsafe fn drop_as_waker(this: *const ()) { // Rather than converting back into an arc, we can just decrement the // counter here. // - // SAFETY: This is called on a pointer created by `Arc::into_raw` on an - // instance of `Arc`. + // SAFETY: `this` came from `Arc::into_raw` and the caller transfers + // ownership of one strong count (per the contract), so decrementing it + // is sound. unsafe { Arc::decrement_strong_count(this.cast::()) }; } } @@ -928,8 +1051,10 @@ mod scope_ptr { /// reference scope from being deallocated. pub struct ScopePtr<'scope, 'env>(*const Scope<'scope, 'env>); - // SAFETY: Transferring ownership of the `*const Scope` is sound because the - // pointee is reached only bia atomic ops and is kept alive by the refcount. + // SAFETY: Sending the `*const Scope` is sound because the pointee outlives + // every `ScopePtr` (the refcount is held until drop), all cross-thread + // mutation of the `Scope` goes through atomics, and its remaining fields are + // written only during construction, so shared reads cannot race. unsafe impl Send for ScopePtr<'_, '_> {} impl<'scope, 'env> ScopePtr<'scope, 'env> { diff --git a/src/thread_pool.rs b/src/thread_pool.rs index 2250813..a2553ac 100644 --- a/src/thread_pool.rs +++ b/src/thread_pool.rs @@ -1,7 +1,9 @@ //! This module contains the api and worker logic for the Forte thread pool. +use alloc::boxed::Box; use alloc::format; use alloc::vec::Vec; +use core::any::Any; use core::array; use core::borrow::Borrow; use core::cell::Cell; @@ -15,6 +17,8 @@ use core::ptr; use core::ptr::NonNull; use core::task::Context; use core::task::Poll; +use std::eprintln; +use std::process::abort; use async_task::Runnable; use crossbeam_queue::SegQueue; @@ -90,18 +94,25 @@ pub struct ThreadPool { /// Holds controls for threads spawned and managed by the pool. Initialized /// on first call to `activate`, to allow for some non-static constructors. managed_workers: Mutex>, + /// Called with the payload of panics that have nowhere to propagate (see + /// [`ThreadPool::handle_panic`]). + panic_handler: Mutex>, } +/// The type of the callback registered by [`ThreadPool::set_panic_handler`]. +pub type PanicHandler = + Box) + Send + Sync + 'static>; + /// A public interface that can be temporarily claimed and used by a thread. /// Claiming a seat allows a thread to participate in the thread pool as a /// worker. -pub struct MemberData { +pub(crate) struct MemberData { /// The sharing side of each seat's work-stealing queue. These should only /// ever be accessed by the thread that currently owns the lease for this /// seat (to ensure the `!Sync` bound is respected). - sharers: [Sharer; 32], + pub sharers: [Sharer; 32], /// The stealing side of each seat's work-stealing queue. - stealers: [Stealer; 32], + pub stealers: [Stealer; 32], /// A set of queues used for transmitting work that must be executed on a /// particular worker. Used for broadcasts and cross-thread nonsend worker /// wakeups. @@ -130,14 +141,24 @@ impl MemberData { // SAFETY: The only `!Sync` field of `MemberData` is `sharers`. This is a // `Sharer` (aka `st3::Worker`). // -// Each `sharers[i]` is only ever touched by the single thread that currently -// holds seat `i`, and seat ownership is exclusive. Ownership thus moves between -// threads one at a time, exactly as if the `Sharer` (which is `Send`) were -// sent. +// Each `sharers[i]` is only ever reached through `sharing_queue()`, which +// indexes by the owner's `member_index`, so it is touched only by the single +// thread that currently holds seat `i`, and seat ownership is exclusive. +// Ownership thus moves between threads one at a time, exactly as if the +// `Sharer` (which is `Send`) were sent. +// +// The handoff is synchronized using the `claimed_bitmask` atomic. A resigning +// worker clears its seat bit with a `Release` RMW on `claimed_bitmask`, and the +// next worker to claim that seat does so with an `Acquire` RMW. When a worker +// resigns their seat, and that seat is then claimed by a different thread, a +// happens-before relationship is established, allowing ownership of the +// `sharers` to be released by the resigning thread and acquired by the claiming +// thread. // -// The handoff is synchronized: a resigning worker makes a `Release` store to -// the `claimed_bitmask`, and a joining worker makes an `Acquire` load of the -// same variable. So the next owner of the seat sees a consistent `Sharer`. +// This holds even when other seats are concurrently claimed and released in +// between: every write to `claimed_bitmask` is an RMW, so a resignation heads +// an unbroken release sequence, and the acquiring claim synchronizes with all +// previous resignations. unsafe impl Sync for MemberData {} /// Represents a worker thread that is managed by the pool, as opposed to @@ -173,7 +194,56 @@ impl ThreadPool { waiting_bitmask: CachePadded::new(AtomicU32::new(0)), wants_to_resign: CachePadded::new(AtomicU32::new(0)), managed_workers: Mutex::new(Vec::new()), + panic_handler: Mutex::new(None), + } + } + + /// Registers a callback to be invoked (via [`handle_panic`]) with the + /// payload of panics that have nowhere to propagate, such as panics from + /// closures or futures passed to [`spawn`]. This replaces any previously + /// registered handler. + /// + /// If no handler is registered, such panics abort the process. + /// + /// The handler must not panic; if it does, the process aborts. + /// + /// [`handle_panic`]: ThreadPool::handle_panic + /// [`spawn`]: ThreadPool::spawn + pub fn set_panic_handler(&self, handler: H) + where + H: Fn(Box) + Send + Sync + 'static, + { + let handler = Box::new(handler); + *self.panic_handler.lock().unwrap() = Some(handler); + } + + /// Disposes of a panic payload that has nowhere to propagate, by passing it + /// to the thread pool handler. If no handler is registered, this aborts the + /// process. + /// + /// # Panics + /// + /// This function never unwinds; if the handler (or anything else within + /// this function) panics, the process aborts. + pub fn handle_panic(&self, payload: Box) { + // Jobs call this function to discharge the `JobRef::new` requirement + // that they never unwind, so nothing in this function may panic: any + // panic (from the handler, a poisoned lock, or the write to stderr) + // is converted into an abort. + let abort_guard = unwind::AbortOnDrop; + { + let handler = self.panic_handler.lock().unwrap(); + match handler.as_ref() { + Some(handler) => handler(payload), + None => { + eprintln!( + "Forte: job panicked with no panic handler set; aborting" + ); + abort(); + } + } } + core::mem::forget(abort_guard); } /// Returns an opaque identifier for this thread pool. @@ -384,6 +454,11 @@ impl ThreadPool { /// /// Note: If the thread pool is full (it already has 32 active members) this /// waits for a vacancy before returning. + /// + /// # Panics + /// + /// Panics that propagate out of the closure may cause aborts, if this call + /// registered this thread as a temporary member of the pool. #[inline(always)] pub fn with_worker(&'static self, func: F) -> R where @@ -413,8 +488,13 @@ impl ThreadPool { // ----------------------------------------------------------------------------- // Schedulers +/// Something that can queue a `JobRef` for execution. pub trait Scheduler<'w>: Send + Sync { - // Is passed the result of `Worker::with_current` + /// Queues a job ref. + /// + /// This is passed the job ref that should be queued, and the outcome of + /// `Worker::with_current` in the current context. The worker parameter does + /// not have to be used, it is simply provided to avoid a lookup. fn schedule(&self, job_ref: JobRef, worker: Option<&'w Worker>); } @@ -435,6 +515,8 @@ pub use async_task::Task; /// A trait for types that can be spawned onto a [`ThreadPool`]. /// +/// This trait is sealed and cannot be implemented outside of forte. +/// /// It is implemented for: /// /// * Closures that satisfy `for<'worker> FnOnce(&'worker Worker) + Send + 'static`. @@ -463,7 +545,7 @@ pub use async_task::Task; /// THREAD_POOL.spawn(|_: &Worker| { }); /// ``` /// Hopefully rustc will fix this type inference failure eventually. -pub trait Spawn: 'static { +pub trait Spawn: crate::sealed::Sealed + 'static { /// The handle returned when spawning this type. type Output: 'static; @@ -496,7 +578,17 @@ where S: for<'w> Scheduler<'w> + 'static, { // Allocate a new job on the heap to store the closure. - let job = HeapJob::new(self); + let op = move |worker: &Worker| { + // Panics may not unwind out of the job, so are sent to the thread + // pool's handler. + if let Err(payload) = unwind::halt_unwinding(|| self(worker)) { + worker.thread_pool.handle_panic(payload); + } + }; + + // SAFETY: `op` wraps the user closure in `halt_unwinding` and routes any + // panic to `handle_panic`, which never unwinds, so `op` cannot unwind. + let job = unsafe { HeapJob::new(op) }; // Turn the job into an "owning" `JobRef` so it can be queued. // @@ -530,7 +622,7 @@ where /// * If the `Runnable` was created for a `!Send` future, this must only be /// called on the thread where the `Runnable` was created. #[inline(always)] -unsafe fn execute_runnable(this: NonNull<()>, _worker: &Worker) { +unsafe fn execute_runnable(this: NonNull<()>, worker: &Worker) { // SAFETY: `Runnable::from_raw` must be given a pointer produced by // `Runnable::into_raw` that has not already been consumed by a `from_raw`. // @@ -539,7 +631,17 @@ unsafe fn execute_runnable(this: NonNull<()>, _worker: &Worker) { let runnable = unsafe { Runnable::<()>::from_raw(this) }; // Poll the task. This will drop the future if the task is // canceled or the future completes. - runnable.run(); + // + // Polling (and dropping) a `!Send` future here is sound because the second + // clause of this function's safety contract requires the caller to call + // `execute_runnable` only on the thread the `Runnable` was created on. + // + // A panic from the future propagates out of `run`, has nowhere else to + // propagate, and must not unwind into the executing worker, so it is + // caught and passed to the pool's panic handler. + if let Err(payload) = unwind::halt_unwinding(|| runnable.run()) { + worker.thread_pool.handle_panic(payload); + } } impl Spawn for Fut @@ -590,11 +692,17 @@ where // SAFETY: `spawn_unchecked` has four obligations: // // * If `Self` is `!Send`, its `Runnable` must be used and dropped on - // the original thread. The `Runnable` is wrapped in a `JobRef` and - // handed to `scheduler`. + // the original thread. // - // The second clause of the `Spawn::spawn` safety contract - // requires the caller to ensure this. + // The second clause of the `Spawn::spawn` safety contract requires + // the caller to confine `scheduler` to the original thread, so + // `execute_runnable`, and by extension the `Runnable`, only + // ever run there. + // + // A `JobRef` has no `Drop` impl, so a `JobRef` that is never executed + // leaks its `Runnable` instead of dropping it. The `Runnable` is thus + // only ever dropped inside `execute_runnable` which, again, only runs + // on the correct original. // // * If `Self` is `!'static`, borrowed variables must outlive its // `Runnable`. @@ -776,7 +884,7 @@ pub struct Membership { impl ThreadPool { /// Returns member data, initializing it on the first call. - pub fn get_member_data(&'static self) -> &'static MemberData { + pub(crate) fn get_member_data(&'static self) -> &'static MemberData { self.member_data.get() } @@ -801,16 +909,28 @@ impl ThreadPool { #[cold] pub fn try_enroll(&'static self) -> Option { loop { + // This is only used to select a candidate seat; the seat may in + // fact be claimed by a concurrent call to `try_enroll` or + // `try_enroll_many`. let available_bitmask = - !self.claimed_bitmask.load(Ordering::Acquire); + !self.claimed_bitmask.load(Ordering::Relaxed); if available_bitmask == 0 { return None; } + + // Select the first free seat. let enrolled_index = available_bitmask.trailing_zeros() as usize; // TZCNT let enrolled_bitmask = 1 << enrolled_index; + + // This RMW tries to atomically take ownership of the seat. + // + // We use `Acquire` ordering here to synchronize with the `Release` + // RMW any previous owner would make during resignation. If the + // returned bits show the seat was already taken, we lost a race and + // retry. if self .claimed_bitmask - .fetch_or(enrolled_bitmask, Ordering::Relaxed) + .fetch_or(enrolled_bitmask, Ordering::Acquire) & enrolled_bitmask == 0 { @@ -830,7 +950,10 @@ impl ThreadPool { } let member_data = self.get_member_data(); loop { - let claimed_bitmask = self.claimed_bitmask.load(Ordering::Acquire); + // This is only used to select candidate seats; any seat may in fact + // be claimed by a concurrent call to `try_enroll` or + // `try_enroll_many`. + let claimed_bitmask = self.claimed_bitmask.load(Ordering::Relaxed); if claimed_bitmask == u32::MAX { return Vec::new(); } @@ -850,13 +973,21 @@ impl ThreadPool { available_bitmask &= available_bitmask - 1; } - // Attempt to claim all selected seats in one atomic step. + // This RMW tries to atomically claim the selected seats all at + // once. Doing this in a single step keeps the batch claim atomic + // with respect to other concurrent `try_enroll_many` calls, and can + // prevent livelocks. + // + // We use `Acquire` ordering on success to synchronize with the + // `Release` RMW any previous owner would make during resignation. + // If the CAS fails we lost a race and retry; since we claimed + // nothing, the failure ordering can be `Relaxed`. if self .claimed_bitmask .compare_exchange( claimed_bitmask, claimed_bitmask | enrolled_bitmask, - Ordering::Relaxed, + Ordering::Acquire, Ordering::Relaxed, ) .is_ok() @@ -949,6 +1080,13 @@ impl Membership { /// /// Rust's thread locals are fairly costly, so this function is expensive. /// If you can avoid calling it, do so. + /// + /// # Panics + /// + /// If a panic occurs within the closure, it is captured and the worker is + /// safely torn down before the panic is re-emitted. Panics within jobs run + /// by this worker are passed to the pool's panic handler, as described in + /// [`Worker::spawn`]. #[inline(always)] pub fn activate(self, f: F) -> R where @@ -957,20 +1095,24 @@ impl Membership { let worker = Worker { migrated: Cell::new(false), membership: self, - fifo_queue: JobQueue::new(), - lifo_queue: JobQueue::new(), + fifo_queue: JobQueue::default(), + lifo_queue: JobQueue::default(), nonsend_fifo_queue: Arc::new(SegQueue::new()), rng: XorShift64Star::new(), last_promote_tick: Cell::new(0), _phantom: PhantomData, }; + // Guard against any unexpected panics, so that the pointer is not left + // dangling. + let abort_guard = unwind::AbortOnDrop; + // Swap the local pointer to point to the newly allocated worker. let outer_ptr = WORKER_PTR.with(|ptr| ptr.replace(&worker)); // Run the function within the context created by the worker pointer, // and pass in a worker reference directly. - let result = f(&worker); + let result = unwind::halt_unwinding(|| f(&worker)); // Indicate that we want to resign. worker @@ -1039,9 +1181,15 @@ impl Membership { // Swap back to pointing to the previous value (possibly null). WORKER_PTR.with(|ptr| ptr.set(outer_ptr)); - // Return the intermediate values created while running the closure, - // namely the result and any jobs still remaining on the local queue. - result + // The worker is fully torn down, so panics may unwind freely again. + core::mem::forget(abort_guard); + + // If the closure panicked, re-raise it now that teardown is complete so + // that it propagates to the caller. Otherwise return its result. + match result { + Ok(result) => result, + Err(panic) => unwind::resume_unwinding(panic), + } } /// Returns a reference to the push-side `Sharer` queue for this @@ -1180,12 +1328,13 @@ impl Worker { // SAFETY: `WORKER_PTR` is a thread-local `Cell` holding a raw // pointer to a `Worker`. It is only written to by // `Membership::activate`, which stores the address of a `Worker` - // allocated within it's own stack frame. Before it returns, - // `activate` restores the previous value of `WORKER_PTR`, so that - // it is always either null or points to a live, immovable `Worker` - // on the current thread's call stack (but is never left dangling). + // allocated within its own stack frame. Before it returns, + // `activate` restores the previous value of `WORKER_PTR` (and it + // aborts on panic, so this reset is guaranteed to occur). So the + // `WORKER_PTR` is either null or a pointer to a live, immovable + // `Worker`. // - // If the pointer is non-null, it is therefore valid to dereference + // If the pointer is non-null, it is therefore sound to dereference // as a shared reference. Forming a `'static` reference is avoided // by passing the value into a closure, which bounds the reference's // lifetime to the closure body and prevents callers from retaining @@ -1211,9 +1360,10 @@ impl Worker { // pointer to a `Worker`. It is only written to by // `Membership::activate`, which stores the address of a `Worker` // allocated within its own stack frame. Before it returns, - // `activate` restores the previous value of `WORKER_PTR`, so that - // it is always either null or points to a live, immovable `Worker` - // on the current thread's call stack (but is never left dangling). + // `activate` restores the previous value of `WORKER_PTR` (and it + // aborts on panic, so this reset is guaranteed to occur). So the + // `WORKER_PTR` is either null or a pointer to a live, immovable + // `Worker` // // If the pointer is non-null, it is therefore sound to dereference // as a shared reference. Forming a `'static` reference is avoided @@ -1294,9 +1444,12 @@ impl Worker { // runner's fifo queue when executed. // // This reduces the cost of sharing a large number of small jobs. - let batch_job = HeapJob::new(move |worker| { + let op = move |worker: &Worker| { worker.fifo_queue.append(job_refs); - }); + }; + // SAFETY: `op` only calls `VecDeque::append` on `JobRef`s, which + // cannot panic. So this cannot unwind. + let batch_job = unsafe { HeapJob::new(op) }; // SAFETY: `HeapJob::into_job_ref` requires: // // * The data closed over by the `HeapJob` outlives the `JobRef`. @@ -1354,6 +1507,9 @@ impl Worker { /// /// The thread may go to sleep if it runs out of work to do, but will wake /// when the latch is set or more work becomes available. + /// + /// This function does not panic: jobs never allow panics from user code + /// to escape their execution (see the panic convention on `JobRef`). #[inline(always)] pub fn wait_for(&self, latch: &Latch) -> bool { loop { @@ -1509,9 +1665,16 @@ impl Worker { /// before the job runs, then swaps it back to what it was before. #[inline(always)] fn execute(&self, job_ref: JobRef, migrated: bool) { + // We restore the flag in a drop guard in case `job_ref.execute` panics. + struct RestoreFlag<'a>(&'a Cell, bool); + impl Drop for RestoreFlag<'_> { + fn drop(&mut self) { + self.0.set(self.1); + } + } let outer_migrated = self.migrated.replace(migrated); + let _restore = RestoreFlag(&self.migrated, outer_migrated); job_ref.execute(self); - self.migrated.set(outer_migrated); } } @@ -1540,10 +1703,14 @@ impl Worker { /// /// The panic behavior depends on the type of work being spawned: /// - /// * If a closure panics, it will be caught and ignored. - /// - /// * If a future panics, the [`Task`] will panic when awaited. + /// * If a closure panics, the panic has nowhere to propagate, so it is + /// caught and passed to the pool's panic handler (see + /// [`ThreadPool::set_panic_handler`]). If no handler is set, the + /// process aborts. /// + /// * If a future panics, the panic is likewise passed to the panic + /// handler. Awaiting the [`Task`] will then also panic, though not + /// with the original payload. #[inline] pub fn spawn(&self, work: S) -> S::Output where @@ -1591,12 +1758,8 @@ impl Worker { /// /// # Panics /// - /// The panic behavior depends on the type of work being spawned: - /// - /// * If a closure panics, it will be caught and ignored. - /// - /// * If a future panics, the [`Task`] will panic when awaited. - /// + /// Panics are passed to the pool's panic handler, as described in + /// [`spawn`](Worker::spawn). #[inline] pub fn spawn_local(&self, work: S) -> S::Output where @@ -1807,6 +1970,10 @@ impl Worker { /// only the panic from the first argument is propagated and the panic from /// the other argument is dropped (this may cause program aborts in some /// situations). + /// + /// Unrelated jobs executed while `join` waits may also panic; those + /// panics are passed to the pool's panic handler (see + /// [`ThreadPool::set_panic_handler`]) and are never re-thrown by `join`. #[inline(always)] pub fn join(&self, a: A, b: B) -> (RA, RB) where @@ -1818,7 +1985,7 @@ impl Worker { // Allocate a job to run the closure `a` on the stack. It is vital to // the correctness of this function that this stack-job never move until // it is freed. - let stack_job = StackJob::new(a, self.new_latch()); + let stack_job = StackJob::new(a, self); // SAFETY: We are only allowed to create a `JobRef` to this `StackJob` // if we can show that... @@ -1840,18 +2007,31 @@ impl Worker { // If `recover_newest` returns `true`, then the `JobRef` must have // been dropped without `execute` being called (satisfying B). // - // If `recover_newest` returns `false`, then we call `wait_for`, which - // will not allow the function to progress until `check` returns - // something other than `Pending` (satisfying A). + // If `recover_newest` returns `false`, then we call `wait_for`, + // which will not allow the function to progress until `check` + // returns something other than `Pending` (satisfying A). // - // In either case, we cannot move or drop the `StackJob` until we pass - // the branch marked with "(*)". We clearly do not. + // In either case, we cannot move or drop the `StackJob` -- by + // returning or by unwinding -- until we pass the branch marked with + // "(*)". We clearly do not return early, and unwinding is ruled out + // by the abort guard held from before the `JobRef` is pushed until + // after that branch: a panic that would unwind this frame in that + // region aborts the process instead. (Jobs executed while waiting + // never allow panics to escape their execution -- see the panic + // convention on `JobRef` -- so only unexpected panics can trip the + // guard.) let job_ref = unsafe { stack_job.as_job_ref() }; // Store the id of the `JobRef` for later, when we will need it to // safely recover the closure `a` for inline execution. let job_ref_id = job_ref.id(); + // Once the job is pushed, other threads may hold references to + // `stack_job` until we pass the branch marked with "(*)". A panic + // must not unwind this frame in that region (see the safety comment + // above), so it aborts the process instead. + let abort_guard = unwind::AbortOnDrop; + // Push the job onto the queue. self.lifo_queue.push_new(job_ref); @@ -1879,19 +2059,23 @@ impl Worker { result_a = unwind::halt_unwinding(|| a(self)); } else { // Wait for the job to complete. - if self.wait_for(stack_job.completion_latch()) { - // SAFETY: Since `wait_for` returned `true`, a `check` must have - // returned `Error`. + if stack_job.wait(self) { + // SAFETY: `wait` returned `true`, so the job completed with an + // error. let error = unsafe { stack_job.unwrap_error() }; result_a = Err(error); } else { - // SAFETY: Since `wait_for` returned `false`, a `check` must have - // returned `Ok`. + // SAFETY: `wait` returned `false`, so the job completed with a + // value. let output = unsafe { stack_job.unwrap_output() }; result_a = Ok(output); } } + // The stack job was either recovered or completed, so panics may + // unwind freely again. + core::mem::forget(abort_guard); + // Resume unwinding if either job panicked. match (result_a, result_b) { (Err(error), _) | (_, Err(error)) => { @@ -2135,6 +2319,11 @@ impl Worker { /// If the operation panics on one or more threads, exactly one panic will /// be propagated, only after all threads have completed (or themselves /// panicked). + /// + /// Unrelated jobs executed while `broadcast` waits may also panic; those + /// panics are passed to the pool's panic handler (see + /// [`ThreadPool::set_panic_handler`]) and are never re-thrown by + /// `broadcast`. #[inline(always)] pub fn broadcast(&self, f: F) -> Vec where @@ -2159,10 +2348,16 @@ impl Worker { participants, }) }; - (member_index, StackJob::new(op, self.new_latch())) + (member_index, StackJob::new(op, self)) }) .collect(); + // Once the jobs are sent, other threads hold references into `jobs` + // until every latch is set. A panic must not unwind this frame in + // that region (see the safety comment below), so it aborts the + // process instead. + let abort_guard = unwind::AbortOnDrop; + // Send the broadcast to each member, and wake them up. for (member_index, job) in &jobs { // SAFETY: We are only allowed to create a `JobRef` for this @@ -2181,8 +2376,15 @@ impl Worker { // // We call `wait_for` on each job's latch (marked with a *). This // does not allow the function to progress while `check` returns - // `Pending`. No `StackJob` is moved or dropped until after this - // function has been called on every `StackJob` and returned. + // `Pending`, and control cannot leave this function -- by + // returning or by unwinding -- before it has been called on + // every `StackJob` and returned: unwinding is ruled out by the + // abort guard held until every latch is set, which turns a + // panic that would unwind this frame into an abort. (Jobs + // executed while waiting never allow panics to escape their + // execution -- see the panic convention on `JobRef` -- so only + // unexpected panics can trip the guard.) No `StackJob` is moved + // or dropped before then. let job_ref = unsafe { job.as_job_ref() }; self.member_data.broadcasts[*member_index].push(job_ref); self.member_data.semaphores[*member_index].signal(); @@ -2191,9 +2393,12 @@ impl Worker { // Wait for each job to finish. let error_flags: Vec<_> = jobs .iter() - .map(|(_, job)| self.wait_for(job.completion_latch())) // (*) + .map(|(_, job)| job.wait(self)) // (*) .collect(); + // Every stack job is now complete, so panics may unwind freely again. + core::mem::forget(abort_guard); + // Allow workers to leave the pool again. self.thread_pool.unfreeze_membership(); @@ -2202,13 +2407,13 @@ impl Worker { .zip(error_flags) .map(|((_, job), error_flag)| { if error_flag { - // SAFETY: If `error_flag` is `true` then `check` has - // returned `Error`. + // SAFETY: `error_flag` (from `wait`) is `true`, so the job + // completed with an error. let error = unsafe { job.unwrap_error() }; unwind::resume_unwinding(error); } else { - // SAFETY: If `error_flag` is `false` then `check` has - // returned `Ok`. + // SAFETY: `error_flag` (from `wait`) is `false`, so the job + // completed with a value. unsafe { job.unwrap_output() } } }) @@ -2224,7 +2429,8 @@ impl Worker { /// /// # Panics /// - /// Panics are not propagated. + /// Panics in the operation are passed to the pool's panic handler, as + /// described in [`Worker::spawn`]. #[inline(always)] pub fn spawn_broadcast(&self, f: F) where @@ -2253,16 +2459,25 @@ impl Worker { // Send the broadcast to each member, and wake them up. for (i, member_index) in members.iter_bits().enumerate() { let func = Arc::clone(&f); + // Run the operation. A panic from it has nowhere to propagate, + // and must not unwind into the executing worker, so it is caught + // and passed to the pool's panic handler. let op = move |worker: &Worker| { - // Run the job - (*func)(Broadcast { - worker, - index: i, - participants, + let result = unwind::halt_unwinding(|| { + (*func)(Broadcast { + worker, + index: i, + participants, + }); }); + if let Err(payload) = result { + worker.thread_pool.handle_panic(payload); + } }; - let job = HeapJob::new(op); + // SAFETY: `op` cannot unwind. All user code is wrapped in + // `halt_unwinding` and `handle_panic` does not unwind. + let job = unsafe { HeapJob::new(op) }; // SAFETY: `HeapJob::into_job_ref` has two preconditions: // @@ -2270,14 +2485,18 @@ impl Worker { // `op`. // // `op` owns an `Arc` clone (`func`). Since `F: 'static`, that - // `Arc` — and everything it keeps alive — is itself `'static`, + // `Arc` and everything it keeps alive is itself `'static`, // so it outlives the `JobRef` regardless of when the job runs. // // * If `F: !Send` then the `JobRef` must only be executed on this // thread. // - // `op` is `Send` (`Arc: Send` because `F: Send + Sync`), so - // this does not apply. + // The `op` is `Send`, so this does not apply. + // + // * The function `op` must not unwind. + // + // `op` catches panics from the broadcast operation and passes + // them to `ThreadPool::handle_panic`, which never unwinds. let job_ref = unsafe { job.into_job_ref() }; self.member_data.broadcasts[member_index].push(job_ref); self.member_data.semaphores[member_index].signal(); diff --git a/src/time.rs b/src/time.rs index b6879c1..7469567 100644 --- a/src/time.rs +++ b/src/time.rs @@ -8,7 +8,8 @@ pub fn ticks() -> u64 { use core::arch::asm; let cnt: u64; // SAFETY: `mrs cntvct_el0` only reads the architectural virtual counter - // register and does not touch memory or the stack. + // register: it touches neither memory (`nomem`) nor the stack (`nostack`), + // and writes no condition flags (`preserves_flags`). unsafe { asm!( "mrs {}, cntvct_el0", @@ -25,8 +26,10 @@ pub fn ticks() -> u64 { pub fn ticks() -> u64 { use core::arch::asm; let cnt: u64; - // SAFETY: `rdtime` reads a timer CSR into a general-purpose register and - // does not access Rust memory. + // SAFETY: `rdtime` reads a timer CSR into its destination register: it + // touches neither memory (`nomem`) nor the stack (`nostack`), and writes no + // other register, leaving `fflags` and vector state untouched + // (`preserves_flags`). unsafe { asm!( "rdtime {}", @@ -45,7 +48,8 @@ pub fn ticks() -> u64 { use core::arch::x86::_rdtsc; #[cfg(target_arch = "x86_64")] use core::arch::x86_64::_rdtsc; - // SAFETY: `_rdtsc` reads the counter into a register and touches neither - // memory nor the stack. + // SAFETY: the `rdtsc` instruction is available on every `x86`/`x86_64` + // target Rust supports (baseline since i586), which is `_rdtsc`'s only + // requirement. unsafe { _rdtsc() } }