diff --git a/CHANGELOG.md b/CHANGELOG.md index 3648e2e..493f170 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,17 @@ This project is currently in early [pre-release], and there may be arbitrary bre ## [Unreleased] +### Added + +- `Scope::spawn_on` as a scoped version of `Worker::spawn_on` +- `Scope::spawn_local` as a scoped version of `Worker::spawn_local` + +### Changed + +- `Worker::spawn_on` now accepts a member number, and sends work to that thread. + +## [1.0.0-alpha.5] + ### Added - `Worker::spawn_local` for spawning `!Send` work. diff --git a/src/job.rs b/src/job.rs index 9ff40e6..51aed0b 100644 --- a/src/job.rs +++ b/src/job.rs @@ -23,8 +23,8 @@ use crate::unwind; // ----------------------------------------------------------------------------- // JobRef -/// A `JobRef` is a specialized v-table, containing a pointer to work that needs to -/// be executed, and a function pointer that is capable of executing it. +/// A `JobRef` is a specialized v-table, containing a pointer to work that needs +/// 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`. pub struct JobRef { @@ -75,11 +75,10 @@ impl JobRef { } } -// SAFETY: This is sound, but just barely. -// -// 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`). +// 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 @@ -108,6 +107,8 @@ unsafe impl Send for JobRef {} /// A queue of jobs. This is a simple wrapper around a vec dequeue that uses /// inner mutation, and has some more intuitively named methods to enforce /// conventions. +/// +/// Note: This is !Sync because of the unsafe cell. pub struct JobQueue { job_refs: UnsafeCell>, } @@ -122,7 +123,7 @@ impl JobQueue { /// 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 not `Sync`, so this can only be called from one + // 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. @@ -132,7 +133,7 @@ impl JobQueue { /// 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 not `Sync`, so this can only be called from one + // 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. @@ -142,7 +143,7 @@ impl JobQueue { /// Removes the newest job in the queue. pub fn pop_newest(&self) -> Option { - // SAFETY: `JobQueue` is not `Sync`, so this can only be called from one + // 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. @@ -152,7 +153,7 @@ impl JobQueue { /// Removes the oldest job in the queue. pub fn pop_oldest(&self) -> Option { - // SAFETY: `JobQueue` is not `Sync`, so this can only be called from one + // 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. @@ -163,7 +164,7 @@ 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 not `Sync`, so this can only be called from one + // 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. @@ -183,7 +184,7 @@ 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 not `Sync`, so this can only be called from one + // 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. @@ -202,7 +203,7 @@ 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 not `Sync`, so this can only be called from one + // 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. @@ -262,8 +263,8 @@ 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) A call to `check` on the `StackJob`'s latch returns something + /// other than `Pending`. /// /// * (B) The `JobRef` has been dropped without `execute` being called. #[inline(always)] @@ -294,8 +295,6 @@ where // The caller ensures only one `JobRef` is ever created for this // `StackJob`. Since `JobRef::execute` consumes that `JobRef`, it // cannot be called multiple times. - // - // unsafe { JobRef::new(job_pointer, Self::execute) } } @@ -320,26 +319,14 @@ where /// * If a `JobRef` did exist, it was never executed. #[inline(always)] pub unsafe fn unwrap_func(mut self) -> F { - // SAFETY: For this access to be valid, we must first establish that - // we have exclusive access to `data`. Only three other functions access - // `data`, and none of them can race with this function: - // - // * Since `JobRef::execute` is not called, `StackJob::execute` is not - // called, and cannot be running. - // - // * Since this function and `unwrap_output` both consume the - // `StackJob`, and a `StackJob` cannot be duplicated, `unwrap_output` - // cannot be running now. - // - // * Since this function and `unwrap_error` both consume the `StackJob`, - // and a `StackJob` cannot be duplicated, `unwrap_error` cannot be - // running now. - // - // Next, we must establish that it is valid to read from union field - // `func`. Each `StackJob` is constructed using field `func`, and only - // `StackJob::execute` writes to the union after construction. Since - // `StackJob::execute` is not called, it must still be valid to read - // from `func`. + // SAFETY: We have exclusive access to the active union field `func`. We + // take `self` by value, so no other `unwrap_*` method can also hold it, + // and the caller guarantees no `JobRef` exists, so `execute` cannot be + // aliasing `data` through a raw pointer. + // + // `func` is the active variant because every `StackJob` is constructed + // with `func`, only `execute` overwrites it, and the caller guarantees + // `execute` was never called. let func_ref = unsafe { &mut self.data.get_mut().func }; // SAFETY: The `StackJob` is dropped at the end of this block, so `data` // is never accessed again. @@ -357,28 +344,15 @@ where // Synchronize with the fence in `StackJob::execute`, establishing a // happens-after relationship with the following read. fence(Ordering::Acquire); - // SAFETY: For this access to be valid, we must first establish that - // we have exclusive access to `data`. Only three other functions access - // `data`, and none of them can race with this function: - // - // * Since `check` has returned `Ok`, and the latch is only set in - // `StackJob::execute`, `StackJob::execute` must have been called at - // least once. `StackJob::execute` may be called at most once, so it - // cannot be running now. - // - // * Since this function and `unwrap_func` both consume the `StackJob`, - // and a `StackJob` cannot be duplicated, `unwrap_func` cannot be - // running now. - // - // * Since this function and `unwrap_error` both consume the `StackJob`, - // and a `StackJob` cannot be duplicated, `unwrap_error` cannot be - // running now. - // - // Next, we must establish that it is valid to read from union field - // `output`. We know this because `check` returned `Ok`, which means - // `set` was called with a false `error_flag` within - // `StackJob::execute`. This always follows a write to union field - // `output`, after which the union is not written to again. + // 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`. + // + // `output` is the active variant because `check` returning `Ok` 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. let output_ref = unsafe { &mut self.data.get_mut().output }; // SAFETY: The `StackJob` is dropped at the end of this block, so `data` // is never accessed again. @@ -396,28 +370,15 @@ where // Synchronize with the fence in `StackJob::execute`, establishing a // happens-after relationship with the following read. fence(Ordering::Acquire); - // SAFETY: For this access to be valid, we must first establish that - // we have exclusive access to `data`. Only three other functions access - // `data`, and none of them can race with this function: - // - // * Since `check` has returned `Error`, and the latch is only set in - // `StackJob::execute`, `StackJob::execute` must have been called at - // least once. `StackJob::execute` may be called at most once, so it - // cannot be running now. - // - // * Since this function and `unwrap_func` both consume the `StackJob`, - // and a `StackJob` cannot be duplicated, `unwrap_func` cannot be - // running now. - // - // * Since this function and `unwrap_output` both consume the `StackJob`, - // and a `StackJob` cannot be duplicated, `unwrap_output` cannot be - // running now. - // - // Next, we must establish that it is valid to read from union field - // `error`. We know this because `check` returned `Error`, which means - // `set` was called with a true `error_flag` within `StackJob::execute`. - // This always follows a write to union field `error`, after which the - // union is not written to again. + // 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`. + // + // `error` is the active variant because `check` returning `Error` 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. let error_ref = unsafe { &mut self.data.get_mut().error }; // SAFETY: The `StackJob` is dropped at the end of this block, so `data` // is never accessed again. @@ -480,9 +441,10 @@ where } } }; - // This synchronizes with the `Acquire` fence within `return_value()`, - // establishing a happens-before relationship that makes the preceding - // `return_value` write visible to the reader. + // 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 is required because latches do not synchronize memory. fence(Ordering::Release); @@ -516,7 +478,8 @@ where /// Represents a job stored in the heap. Used to implement `scope` and `spawn`. /// -/// This is analogous to the rayon type `HeapJob`. There is no corresponding chili type. +/// This is analogous to the rayon type `HeapJob`. There is no corresponding +/// chili type. pub struct HeapJob { f: F, } @@ -555,10 +518,11 @@ where // We must now show that it is sound to call `HeapJob::execute` on // `job_ref` under these conditions, which in turn requires that: // - // * `job_pointer` is an aligned pointer to an initialized `Box`. + // * `job_pointer` is an aligned pointer to an initialized `HeapJob`. // - // We created it from a ref to `self`, which is a `Box`, so it - // must be. + // We created it with `Box::into_raw(self)`, which yields an aligned, + // non-null pointer to the initialized `HeapJob` formerly owned by the + // `Box`, so it must be. // // * `HeapJob::execute` is called at most once on any `HeapJob`. // diff --git a/src/latch.rs b/src/latch.rs index e03ee26..3c7cca9 100644 --- a/src/latch.rs +++ b/src/latch.rs @@ -31,11 +31,12 @@ const ASLEEP: u32 = 0b100; // Latch /// A Latch is a signaling mechanism used to indicate when an event has -/// occurred. The latch begins as *unset* (In the `LOCKED` state), and can later -/// be *set* by any thread (entering the `SIGNAL`) state. +/// occurred. /// -/// 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. +/// 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. /// /// The general idea and spirit for latches (as well as some of the /// documentation) is due to rayon. However the implementation is specific to @@ -141,15 +142,16 @@ impl Latch { // First we store a reference to the semaphore (which is 'static) so // that we can access it even if the latch pointer becomes dangling. // - // SAFETY: The caller guarantees the latch pointer is aligned and non-null. + // SAFETY: The caller guarantees the latch pointer is aligned and + // non-null. // // If Variant 1 is met, the latch cannot be dangling. // // If Variant 2 is met, the latch cannot become dangling so long as the - // state is `LOCKED` (because `check` will return `Pending`). Since there - // can have been no previous call to `set` since construction or the - // last `reset`, and there can be no racing calls to `set`, the state - // must be `LOCKED`. Therefore the latch cannot be dangling. + // state is `LOCKED` (because `check` will return `Pending`). Since + // there can have been no previous call to `set` since construction or + // the last `reset`, and there can be no racing calls to `set`, the + // state must be `LOCKED`. Therefore the latch cannot be dangling. // // Since this pointer is aligned, non-null, is not dangling, and the // latch is never accessed mutably, it is valid to access immutably. diff --git a/src/lib.rs b/src/lib.rs index 6675a7b..027bf56 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,8 +11,6 @@ //! //! * The ability to execute both closures and futures on the same pool. //! -//! * Hybrid scopes that can contain work distributed across multiple thread pools. -//! //! * A primitive for awaiting async work in non-async contexts without spinning. //! //! * An exposed unsafe api, built for low-level integration and customization. @@ -214,6 +212,7 @@ pub struct FutureMarker(); // ----------------------------------------------------------------------------- // Top-level exports +pub use latch::Latch; pub use scope::Scope; pub use scope::SpawnScoped; pub use thread_pool::Broadcast; @@ -221,7 +220,6 @@ pub use thread_pool::DEFAULT_POOL; pub use thread_pool::DefaultThreadPool; pub use thread_pool::Membership; pub use thread_pool::Spawn; -pub use thread_pool::SpawnLocal; pub use thread_pool::Task; pub use thread_pool::ThreadPool; pub use thread_pool::Worker; @@ -233,6 +231,7 @@ pub use thread_pool::num_members; pub use thread_pool::scope; pub use thread_pool::spawn; pub use thread_pool::spawn_broadcast; +pub use thread_pool::spawn_on; // ----------------------------------------------------------------------------- // Platform Support diff --git a/src/scope.rs b/src/scope.rs index 6f41229..7435843 100644 --- a/src/scope.rs +++ b/src/scope.rs @@ -28,6 +28,7 @@ use crate::job::JobRef; use crate::latch::Latch; use crate::platform::*; use crate::thread_pool::Broadcast; +use crate::thread_pool::Scheduler; use crate::thread_pool::Worker; use crate::unwind; use crate::unwind::AbortOnDrop; @@ -157,26 +158,59 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// * A `for<'worker> FnOnce(&'worker Worker)` closure, with no return type. /// /// * A `Future` future, with no return type. - pub fn spawn>(&'scope self, scoped_work: S) { - self.thread_pool.get_worker(|worker| { - scoped_work.spawn_scoped(self, worker); - }); + pub fn spawn(&'scope self, work: S) + where + S: SpawnScoped<'scope, M> + Send, + { + let thread_pool = self.thread_pool; + let scheduler = |job_ref, worker: Option<&Worker>| match worker { + Some(worker) => worker.fifo_queue.push_new(job_ref), + None => thread_pool.push_shared_job(job_ref), + }; + // SAFETY: `spawn_scoped` requires that `!Send` work only be scheduled + // to run on the current thread. Here `S: Send`, so that obligation is + // vacuous. + unsafe { work.spawn_scoped(self, scheduler) }; } - /// Runs a closure or future sometime before the scope completes. Valid - /// inputs to this method are: - /// - /// * A `for<'worker> FnOnce(&'worker Worker)` closure, with no return type. - /// - /// * A `Future` future, with no return type. - /// - /// Unlike [`Scope::spawn`], this accepts the current worker as a parameter. - pub fn spawn_on>( - &'scope self, - worker: &Worker, - scoped_work: S, - ) { - scoped_work.spawn_scoped(self, Some(worker)); + /// Like [`spawn`](Self::spawn), but routes the work to a specific pool + /// member so it runs on that member's thread. + pub fn spawn_on(&'scope self, member_index: usize, work: S) + where + S: SpawnScoped<'scope, M> + Send, + { + let member_data = self.thread_pool.get_member_data(); + let scheduler = |job_ref, _: Option<&Worker>| { + member_data.broadcasts[member_index].push(job_ref); + member_data.semaphores[member_index].signal(); + }; + // SAFETY: `spawn_scoped` requires that `!Send` work only be scheduled + // to run on the current thread. Here `S: Send`, so that obligation is + // vacuous. + unsafe { work.spawn_scoped(self, scheduler) } + } + + /// Like [`spawn`](Self::spawn), but accepts `!Send` work, which is confined + /// to (and only ever polled on) the current worker's thread. + pub fn spawn_local(&'scope self, worker: &Worker, work: S) + where + S: SpawnScoped<'scope, M>, + { + let queue = worker.nonsend_fifo_queue.clone(); + let semaphore = &worker.member_data.semaphores[worker.member_index]; + let scheduler = move |job_ref, _: Option<&Worker>| { + queue.push(job_ref); + semaphore.signal(); + }; + // SAFETY: `spawn_scoped` requires that `!Send` work only be scheduled + // to run on the current thread. + // + // The scheduler pushes the job onto this worker's `nonsend_fifo_queue`. + // That queue is only ever drained by this worker on this thread (see + // `Worker::find_local_work`). For a future, every reschedule re-pushes + // onto the same queue regardless of which thread issued the wakeup, so + // the work is confined to this thread. + unsafe { work.spawn_scoped(self, scheduler) } } /// Runs an operation across multiple threads. @@ -417,26 +451,33 @@ impl<'scope, 'env> Scope<'scope, 'env> { /// }); /// ``` /// Hopefully rustc will fix this type inference failure eventually. -pub trait SpawnScoped<'scope, M>: Send + 'scope { +pub trait SpawnScoped<'scope, M>: '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. - fn spawn_scoped<'env>( + /// + /// # Safety + /// + /// If `Self` is not `Send`, the `scheduler` must only ever schedule the + /// work to run on the current thread. For a future, this includes every + /// reschedule triggered by a wakeup, no matter which thread the wakeup + /// occurs on. + unsafe fn spawn_scoped<'env, S: for<'w> Scheduler<'w>>( self, scope: &'scope Scope<'scope, 'env>, - worker: Option<&Worker>, + scheduler: S, ); } impl<'scope, F> SpawnScoped<'scope, FnOnceMarker> for F where - F: FnOnce(&Worker) + Send + 'scope, + F: FnOnce(&Worker) + 'scope, { #[inline] - fn spawn_scoped<'env>( + unsafe fn spawn_scoped<'env, S: for<'w> Scheduler<'w>>( self, scope: &'scope Scope<'scope, 'env>, - worker: Option<&Worker>, + scheduler: S, ) { // Create a job to execute the spawned function in the scope. let scope_ptr = ScopePtr::new(scope); @@ -449,14 +490,6 @@ where drop(scope_ptr); }); - // SAFETY: We must ensure that the heap job does not outlive the data it - // closes over. In effect, this means it must not outlive `'scope`. - // - // This is ensured by the `scope_ptr` and the scope rules, which will - // keep the calling stack frame alive until this job completes, - // effectively extending the lifetime of `'scope` for as long as is - // necessary. - // SAFETY: `HeapJob::into_job_ref` requires: // // * The `JobRef` will not outlive any of the items closed over by @@ -478,29 +511,22 @@ where let job_ref = unsafe { job.into_job_ref() }; // Send the job to a queue to be executed. - match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => scope.thread_pool.push_shared_job(job_ref), - } + Worker::with_current(|worker| scheduler.schedule(job_ref, worker)); } } impl<'scope, Fut> SpawnScoped<'scope, FutureMarker> for Fut where - Fut: Future + Send + 'scope, + Fut: Future + 'scope, { #[inline] - fn spawn_scoped<'env>( + unsafe fn spawn_scoped<'env, S: for<'w> Scheduler<'w>>( self, scope: &'scope Scope<'scope, 'env>, - worker: Option<&Worker>, + scheduler: S, ) { - let poll_job = ScopeFutureJob::new(scope, self); - let job_ref = poll_job.into_job_ref(); - match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => scope.thread_pool.push_shared_job(job_ref), - } + let job = ScopeFutureJob::new(scope, scheduler, self); + Worker::with_current(|worker| job.schedule(worker)); } } @@ -563,7 +589,7 @@ const LOCKED: u32 = 2; /// The future is always queued on the thread on which it was woken. This allows /// the waker to be converted directly into a job-ref, and may mean that the /// cache is already primed. -struct ScopeFutureJob<'scope, 'env, Fut> { +struct ScopeFutureJob<'scope, 'env, S: for<'w> Scheduler<'w>, Fut> { /// The future that the job exists to poll. This is stored in an unsafe cell /// 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 @@ -574,11 +600,14 @@ struct ScopeFutureJob<'scope, 'env, Fut> { scope_ptr: ScopePtr<'scope, 'env>, /// The state of the job, which is either READY, WOKEN, or LOCKED. state: AtomicU32, + /// The scheduler used to queue this job whenever it needs to be polled. + scheduler: S, } -impl<'scope, 'env, Fut> ScopeFutureJob<'scope, 'env, Fut> +impl<'scope, 'env, S, Fut> ScopeFutureJob<'scope, 'env, S, Fut> where - Fut: Future + Send + 'scope, + S: for<'w> Scheduler<'w>, + Fut: Future + 'scope, { /// This vtable is part of what allows a `ScopeFutureJob` to act as an /// async task waker. @@ -591,7 +620,11 @@ where /// Creates a new `ScopedFutureJob` in an `Arc`. The caller is expected to /// immediately call `into_job_ref` and queue it on a worker to be polled. - fn new(scope: &Scope<'scope, 'env>, future: Fut) -> Arc { + fn new( + scope: &Scope<'scope, 'env>, + scheduler: S, + future: Fut, + ) -> Arc { let scope_ptr = ScopePtr::new(scope); Arc::new(Self { future: UnsafeCell::new(future), @@ -599,6 +632,7 @@ where // The job starts in the WOKEN state because we always queue it // after creating it. state: AtomicU32::new(WOKEN), + scheduler, }) } @@ -617,24 +651,43 @@ where // called on the resulting `JobRef` where it is sound to call `poll` on // `job_pointer`. // - // `Poll` has two preconditions: + // `Poll` has two obligations: // - // 1. `job_pointer` must have been produced by `Arc::into_raw` on an - // `Arc`. + // * `job_pointer` must have been produced by `Arc::into_raw` on an + // `Arc`. // - // We produced `job_pointer` this way just above. + // We produced `job_pointer` this way just above. // - // 2. We must hold ownership of exactly one strong reference count for - // the allocation. + // * We must hold ownership of exactly one strong reference count for + // the allocation. // - // We start with an `Arc`, so we must own a strong reference. - // Calling `Arc::into_raw` transfers the strong count of `self` onto - // `job_pointer` without decrementing it. Therefore, when `execute` - // is called, there will still be a strong reference for it to - // consume. + // We start with an `Arc`, so we must own a strong reference. + // Calling `Arc::into_raw` transfers the strong count of `self` onto + // `job_pointer` without decrementing it. Therefore, when `execute` + // is called, there will still be a strong reference for it to + // consume. unsafe { JobRef::new(job_pointer, Self::poll) } } + /// Schedules this job to be polled, by converting it into a `JobRef` and + /// 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. + // + // `into_job_ref` consumes `self` via `Arc::into_raw`, moving our strong + // reference into the `JobRef` *without* decrementing the count. The act + // of calling `schedule` enqueues that `JobRef`, after which another + // worker may execute and drop it before `schedule` returns. Holding + // `this` keeps an independent strong reference alive until the end of + // this function, so borrowing `this.scheduler` across the call cannot + // dangle. + let this = Arc::clone(&self); + let job_ref = self.into_job_ref(); + this.scheduler.schedule(job_ref, worker); + } + /// Polls the future. /// /// # Safety @@ -769,11 +822,13 @@ where // immediately. Conveniently, we know the state will already be // WOKEN, so we can leave it as it is. if rescheduled { - // This converts the local `Arc` into a job ref, - // preventing it from being dropped and potentially - // extending the job's lifetime. - let job_ref = this.into_job_ref(); - worker.fifo_queue.push_new(job_ref); + // The future was woken while we were polling it, so it is + // already back in the WOKEN state; queue it to be polled + // again, preferring this worker. + // + // NOTE: Eventually we may want to chose not to re-evaluate + // this sometimes, to break out of infinite async loops. + this.schedule(Some(worker)); } } // The job panicked. Store the panic in the scope so it can be @@ -815,13 +870,7 @@ where let this = unsafe { Arc::from_raw(this.cast::()) }; if this.state.swap(WOKEN, Ordering::Relaxed) == READY { - // Convert the waker into a job ref and queue it. - let thread_pool = this.scope_ptr.thread_pool(); - let job_ref = this.into_job_ref(); - thread_pool.get_worker(|worker| match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => thread_pool.push_shared_job(job_ref), - }); + Worker::with_current(|worker| this.schedule(worker)); } } @@ -843,12 +892,7 @@ where 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()); - let thread_pool = this.scope_ptr.thread_pool(); - let job_ref = this.into_job_ref(); - thread_pool.get_worker(|worker| match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => thread_pool.push_shared_job(job_ref), - }); + Worker::with_current(|worker| this.schedule(worker)); } } @@ -878,33 +922,16 @@ mod scope_ptr { use core::any::Any; use super::Scope; - use crate::ThreadPool; /// A reference-counted pointer to a scope. Used to capture a scope pointer /// in jobs without faking a lifetime. Holding a `ScopePtr` keeps the /// reference scope from being deallocated. pub struct ScopePtr<'scope, 'env>(*const Scope<'scope, 'env>); - // SAFETY: This is sound because: - // - // * `ScopePtr` is only used to call `add_reference`, `remove_reference`, - // and `store_panic`, all of which are designed to be called from multiple - // threads concurrently. - // - // * The `Scope` cannot be deallocated while any `ScopePtr` still points to - // it (due to reference counting), so the raw pointer is always valid. + // 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. unsafe impl Send for ScopePtr<'_, '_> {} - // SAFETY: This is sound because: - // - // * `ScopePtr` is only used to call `add_reference`, `remove_reference`, - // and `store_panic`, all of which are designed to be called from multiple - // threads concurrently. - // - // * The `Scope` cannot be deallocated while any `ScopePtr` still points to - // it (due to reference counting), so the raw pointer is always valid. - unsafe impl Sync for ScopePtr<'_, '_> {} - impl<'scope, 'env> ScopePtr<'scope, 'env> { /// Creates a new reference-counted scope pointer which can be sent to other /// threads. @@ -924,14 +951,6 @@ mod scope_ptr { let scope_ref = unsafe { &*self.0 }; scope_ref.store_panic(err); } - - pub fn thread_pool(&self) -> &'static ThreadPool { - // SAFETY: This was created using an immutable scope reference, and - // by the scope rules there can be no mutable references to this - // scope, nor can the scope have been moved or deallocated while the - // scope's counter remains incremented. - unsafe { (&*self.0).thread_pool } - } } impl Drop for ScopePtr<'_, '_> { @@ -1057,9 +1076,9 @@ mod tests { counter.fetch_add(1, Ordering::Relaxed); }); }); - scope.spawn(|worker: &Worker| { + scope.spawn(|_: &Worker| { counter.fetch_add(1, Ordering::Relaxed); - scope.spawn_on(worker, |_: &Worker| { + scope.spawn(|_: &Worker| { counter.fetch_add(1, Ordering::Relaxed); }); }); @@ -1163,10 +1182,10 @@ mod tests { let a = AtomicU8::new(0); let b = AtomicU8::new(0); - THREAD_POOL.with_worker(|worker| { + THREAD_POOL.with_worker(|_: &Worker| { scope(|scope| { for _ in 0..NUM_JOBS { - scope.spawn_on(worker, |_: &Worker| { + scope.spawn(|_: &Worker| { THREAD_POOL.join( |_| a.fetch_add(1, Ordering::Relaxed), |_| b.fetch_add(1, Ordering::Relaxed), @@ -1214,12 +1233,12 @@ mod tests { THREAD_POOL.with_worker(|worker| { worker.scope(|scope| { - scope.spawn_on(worker, |_: &Worker| { + scope.spawn(|_: &Worker| { // Creating a new worker instead of reusing the old one is // bad form, but we may as well test it. THREAD_POOL.with_worker(|worker| { worker.scope(|scope| { - scope.spawn_on(worker, |_: &Worker| { + scope.spawn(|_: &Worker| { completed = true; }); }); @@ -1243,8 +1262,8 @@ mod tests { let counter_p = &AtomicUsize::new(0); THREAD_POOL.with_worker(|worker| { worker.scope(|scope| { - scope.spawn(move |worker: &Worker| { - divide_and_conquer(worker, scope, counter_p, 1024) + scope.spawn(move |_: &Worker| { + divide_and_conquer(scope, counter_p, 1024) }) }); }); @@ -1260,17 +1279,16 @@ mod tests { } fn divide_and_conquer<'scope, 'env>( - worker: &Worker, scope: &'scope Scope<'scope, 'env>, counter: &'scope AtomicUsize, size: usize, ) { if size > 1 { - scope.spawn_on(worker, move |worker: &Worker| { - divide_and_conquer(worker, scope, counter, size / 2) + scope.spawn(move |_: &Worker| { + divide_and_conquer(scope, counter, size / 2) }); - scope.spawn_on(worker, move |worker: &Worker| { - divide_and_conquer(worker, scope, counter, size / 2) + scope.spawn(move |_: &Worker| { + divide_and_conquer(scope, counter, size / 2) }); } else { // count the leaves @@ -1337,9 +1355,9 @@ mod tests { OP: Fn(&mut T) + Sync, { let Tree { value, children } = self; - scope.spawn(move |worker: &Worker| { + scope.spawn(move |_: &Worker| { for child in children { - scope.spawn_on(worker, move |_: &Worker| { + scope.spawn(move |_: &Worker| { let child = child; child.update_in_scope(scope, op) }); diff --git a/src/thread_pool.rs b/src/thread_pool.rs index e50cfdb..2250813 100644 --- a/src/thread_pool.rs +++ b/src/thread_pool.rs @@ -127,11 +127,17 @@ impl MemberData { } } -// SAFETY: `Sharer` (aka `st3::Worker`) is `!Sync`. We allow it to be stored in -// this shared structure, but we only allow one thread to access it at a time -// (via the membership claiming logic). This is effectively like sending -// `st3:Worker` ownership between threads (although in practice it always -// occupies the same place on the heap). Luckily for us, it implements `Send`. +// 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. +// +// 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`. unsafe impl Sync for MemberData {} /// Represents a worker thread that is managed by the pool, as opposed to @@ -404,6 +410,24 @@ impl ThreadPool { } } +// ----------------------------------------------------------------------------- +// Schedulers + +pub trait Scheduler<'w>: Send + Sync { + // Is passed the result of `Worker::with_current` + fn schedule(&self, job_ref: JobRef, worker: Option<&'w Worker>); +} + +impl<'w, F> Scheduler<'w> for F +where + F: Fn(JobRef, Option<&'w Worker>) + Send + Sync, +{ + #[inline(always)] + fn schedule(&self, job_ref: JobRef, worker: Option<&'w Worker>) { + self(job_ref, worker); + } +} + // ----------------------------------------------------------------------------- // Spawn Trait @@ -439,49 +463,58 @@ pub use async_task::Task; /// THREAD_POOL.spawn(|_: &Worker| { }); /// ``` /// Hopefully rustc will fix this type inference failure eventually. -pub trait Spawn: Send + 'static { +pub trait Spawn: 'static { /// The handle returned when spawning this type. - type Output: Send + 'static; + type Output: 'static; - /// Spawns work onto the thread pool. - fn spawn( + /// Turns `Self` into a `JobRef` and passes it to the `Scheduler`. + /// + /// # Safety + /// + /// If `Self` is not `Send`, the `scheduler` must only ever schedule the + /// work to run on the current thread. For a future, this includes every + /// reschedule triggered by a wakeup, no matter which thread the wakeup + /// occurs on. + unsafe fn spawn( self, - thread_pool: &'static ThreadPool, + scheduler: S, worker: Option<&Worker>, - ) -> Self::Output; + ) -> Self::Output + where + S: for<'w> Scheduler<'w> + 'static; } impl Spawn for F where - F: for<'worker> FnOnce(&'worker Worker) + Send + 'static, + F: for<'worker> FnOnce(&'worker Worker) + 'static, { type Output = (); #[inline] - fn spawn(self, thread_pool: &'static ThreadPool, worker: Option<&Worker>) { + unsafe fn spawn(self, scheduler: S, worker: Option<&Worker>) + where + S: for<'w> Scheduler<'w> + 'static, + { // Allocate a new job on the heap to store the closure. let job = HeapJob::new(self); // Turn the job into an "owning" `JobRef` so it can be queued. // - // SAFETY: `HeapJob::into_job_ref` has two preconditions: + // SAFETY: `HeapJob::into_job_ref` requires: // - // * The `JobRef` must not outlive any of the items closed over by the - // function `f`. + // * The `JobRef` does not outlive any of the data closed over by the + // closure. `F: 'static`, so the closure captures only `'static` data, + // which outlives any `JobRef`. // - // Since `F: 'static`, the `JobRef` cannot outlive its captured data. + // * If the closure is `!Send`, `JobRef::execute` is only called on the + // thread where the `HeapJob` was constructed. // - // * If `F: !Send` then the `JobRef` must only be executed on this - // thread. - // - // `F` is `Send`, so this does not apply. + // The safety contract of `Spawn::spawn` makes the caller + // responsible for ensuring this. let job_ref = unsafe { job.into_job_ref() }; - // Queue the job for evaluation - match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => thread_pool.push_shared_job(job_ref), - } + // Schedule the job-ref. + scheduler.schedule(job_ref, worker); } } @@ -498,10 +531,11 @@ where /// called on the thread where the `Runnable` was created. #[inline(always)] unsafe fn execute_runnable(this: NonNull<()>, _worker: &Worker) { - // SAFETY: This pointer was created by `Runnable::into_raw` in the schedule - // closure. Jobs are executed exactly once, so `from_raw` is called at most - // once on this function (the next call to `schedule` will call `into_raw` - // again to get a "new" raw pointer to the same runnable). + // SAFETY: `Runnable::from_raw` must be given a pointer produced by + // `Runnable::into_raw` that has not already been consumed by a `from_raw`. + // + // The caller ensures this, according to the first clause of this function's + // safety contract. let runnable = unsafe { Runnable::<()>::from_raw(this) }; // Poll the task. This will drop the future if the task is // canceled or the future completes. @@ -510,188 +544,80 @@ unsafe fn execute_runnable(this: NonNull<()>, _worker: &Worker) { impl Spawn for Fut where - Fut: Future + Send + 'static, - T: Send + 'static, + Fut: Future + 'static, + T: 'static, { type Output = Task; #[inline] - fn spawn( - self, - thread_pool: &'static ThreadPool, - _worker: Option<&Worker>, - ) -> Task { - // Creates a schedule function that captures a reference to the - // thread-pool. - let schedule = |runnable: Runnable| { + unsafe fn spawn(self, scheduler: S, _worker: Option<&Worker>) -> Task + where + S: for<'w> Scheduler<'w> + 'static, + { + // Turn our `JobRef` scheduler into something that `async-task` knows how + // to use. + let schedule_task = move |runnable: Runnable| { // Temporarily turn the task into a raw pointer so that it can be // used as a job. We could also use `HeapJob` here, but since // `Runnable` is heap allocated this would result in a needless // second allocation. let job_pointer = runnable.into_raw(); - // SAFETY: `JobRef::new` requires us to show that `execute` will - // only be called on the returned `JobRef` when it is sound to call - // `execute_runnable` on `job_pointer`. The two preconditions to - // this are: - // - // * `job_pointer` must come from `Runnable::into_raw`. and must be - // called at most once per `into_raw`. + // SAFETY: `JobRef::new` requires us to show that `JobRef::execute` + // will only be called on the returned `JobRef` when it would be + // sound to call `execute_runnable` on `job_pointer`. This requires: // - // We produced `job_pointer` this way just above. The call to - // `execute` will consume the `JobRef`, so `execute_runnable` will - // be called at most once. + // * That `job_pointer` was produced by `Runnable::into_raw` and not + // yet consumed by `Runnable::from_raw`. We produced it with + // `into_raw` immediately above. The only call to `from_raw` is + // inside `execute_runnable`, which runs at most once because + // `JobRef::execute` consumes the `JobRef`. // - // * If the `Runnable` was created for a `!Send` future, the - // `JobRef` must only be executed on the thread where the - // `Runnable` was created. + // * If the future is `!Send`, `execute_runnable` is only called on + // the thread the future was spawned on. // - // The future is required to be `Send`, so this does not apply. + // The second clause of the `Spawn::spawn` safety contract + // requires the caller to ensure this. let job_ref = unsafe { JobRef::new(job_pointer, execute_runnable) }; - // Send this job off to be executed. - thread_pool.get_worker(|worker| match worker { - Some(worker) => worker.fifo_queue.push_new(job_ref), - None => thread_pool.push_shared_job(job_ref), - }); + // Schedule the job-ref, looking up the current worker so woken + // futures tend to run on the thread that woke them. + Worker::with_current(|worker| scheduler.schedule(job_ref, worker)); }; - // Create a runnable and add the thread pool as metadata. - let (runnable, task) = async_task::spawn(self, schedule); - - // Call the schedule function, pushing a `JobRef` for the future onto - // the local work queue. If the future doesn't complete, it can be - // woken and scheduled at a later point. + // Create a runnable for the future. // - // Because we always look up the local worker within the schedule - // function, woken futures will tend to run on the thread that wakes - // them. This is a desirable property, as typically the next thing a - // future is going to do after being woken up is read some data from the - // thread/task that woke it. + // SAFETY: `spawn_unchecked` has four obligations: // - // This is potentially more efficient than `Runnable::schedule`. - schedule(runnable); - - // Return the task. - task - } -} - -// ----------------------------------------------------------------------------- -// Local Spawn Trait - -/// A version of the [`Spawn`] trait without the `Send` bound. -/// -/// It is implemented for: -/// -/// * Closures that satisfy `for<'worker> FnOnce(&'worker Worker) + 'static`. -/// -/// * Futures that satisfy `Future + 'static` where `T: 'static`. -pub trait SpawnLocal: 'static { - /// The handled returned when spawning this type. - type Output: 'static; - - /// Spawns work that will run in the background on the current worker - /// thread. - fn spawn_local(self, worker: &Worker) -> Self::Output; -} - -impl SpawnLocal for F -where - F: for<'worker> FnOnce(&'worker Worker) + 'static, -{ - type Output = (); - - #[inline] - fn spawn_local(self, worker: &Worker) { - // Allocate a new job on the heap to store the closure. - let job = HeapJob::new(self); - - // Turn the job into an "owning" `JobRef` so it can be queued. + // * 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`. // - // SAFETY: `HeapJob::into_job_ref` has two preconditions: + // The second clause of the `Spawn::spawn` safety contract + // requires the caller to ensure this. // - // * The `JobRef` must not outlive any of the items closed over by the - // function `f`. + // * If `Self` is `!'static`, borrowed variables must outlive its + // `Runnable`. // - // Since `F: 'static`, the `JobRef` cannot outlive its captured data. + // Vacuously true, since `Self: 'static`. // - // * If `F: !Send` then the `JobRef` must only be executed on this - // thread. + // * If `schedule_task` is `!Send`/`!Sync`, all of the `Runnable`'s + // `Waker`s must be used and dropped on the original thread. // - // This `JobRef` is added to the `nonsend_fifo_queue` for this thread. - // No other thread ever pulls from this queue, and work is never - // shared from it, so it cannot be executed on any other thread. - let job_ref = unsafe { job.into_job_ref() }; - - // Push into the non-send queue, which can only be accessed from this - // thread. - worker.nonsend_fifo_queue.push(job_ref); - } -} - -impl SpawnLocal for Fut -where - Fut: Future + 'static, - T: 'static, -{ - type Output = Task; - - #[inline] - fn spawn_local(self, worker: &Worker) -> Task { - // Create a schedule function that will keep a copy of the local fifo - // queue arc and be able to wake the local worker up. - let queue = worker.nonsend_fifo_queue.clone(); - let member_index = worker.member_index; - let member_data = worker.member_data; - let schedule = move |runnable: Runnable| { - // Temporarily turn the task into a raw pointer so that it can be - // used as a job. We could also use `HeapJob` here, but since - // `Runnable` is heap allocated this would result in a needless - // second allocation. - let job_pointer = runnable.into_raw(); - - // SAFETY: `JobRef::new` requires us to show that `execute` will - // only be called on the returned `JobRef` when it is sound to call - // `execute_runnable` on `job_pointer`. The two preconditions to - // this are: - // - // * `job_pointer` must come from `Runnable::into_raw`. and must be - // called at most once per `into_raw`. - // - // We produced `job_pointer` this way just above. The call to - // `execute` will consume the `JobRef`, so `execute_runnable` will - // be called at most once. - // - // * If the `Runnable` was created for a `!Send` future, the - // `JobRef` must only be executed on the thread where the - // `Runnable` was created. - // - // This `JobRef` is added to the `nonsend_fifo_queue` for this - // thread. No other thread ever pulls from this queue, and work is - // never shared from it, so it cannot be executed on any other - // thread. - let job_ref = unsafe { JobRef::new(job_pointer, execute_runnable) }; - - // Send this job to the correct thread to be executed. - queue.push(job_ref); - - // Ensure that the worker is awake to execute this job. - member_data.semaphores[member_index].signal(); - }; - - // Create a runnable and add the thread pool as metadata. - let (runnable, task) = async_task::spawn_local(self, schedule); - - // Call the schedule function, pushing a `JobRef` for the future onto - // the local work queue. If the future doesn't complete, it can be - // woken and scheduled at a later point. + // Vacuously true, since `S: Scheduler` is `Send + Sync` so too + // must `schedule_task` be `Send + Sync`. + // + // * If `schedule_task` is `!'static`, borrowed variables must outlive + // all of the `Runnable`'s `Waker`s. // - // Because we always look up the local worker within the schedule - // function, woken futures will tend to run on the thread that wakes - // them. This is a desirable property, as typically the next thing a - // future is going to do after being woken up is read some data from the - // thread/task that woke it. + // Vacuously true, since `S: 'static`. + let (runnable, task) = + unsafe { async_task::spawn_unchecked(self, schedule_task) }; + + // Perform the initial schedule via the task's own stored schedule + // function, pushing a `JobRef` for the future onto the current worker's + // queue. If the future doesn't complete, it can be woken and scheduled + // again later. runnable.schedule(); // Return the task. @@ -722,8 +648,44 @@ impl ThreadPool { /// /// See also: [`Worker::spawn`] and [`spawn`]. #[inline(always)] - pub fn spawn>(&'static self, work: S) -> S::Output { - self.get_worker(|worker| work.spawn(self, worker)) + pub fn spawn(&'static self, work: S) -> S::Output + where + S: Spawn + Send, + S::Output: Send, + { + let scheduler = |job_ref, worker: Option<&Worker>| match worker { + Some(worker) => worker.fifo_queue.push_new(job_ref), + None => self.push_shared_job(job_ref), + }; + Worker::with_current(|worker| { + // SAFETY: `Spawn::spawn`'s contract requires that `!Send` work is + // only scheduled to run on this thread. The work is `Send` here, so + // this is vacuous. + unsafe { work.spawn(scheduler, worker) } + }) + } + + /// Tries to spawn a job on a specific member index. If no thread currently + /// holds that membership, the work will not be run. + #[inline(always)] + pub fn spawn_on( + &'static self, + member_index: usize, + work: S, + ) -> S::Output + where + S: Spawn + Send, + S::Output: Send, + { + let member_data = self.get_member_data(); + let scheduler = move |job_ref, _: Option<&Worker>| { + member_data.broadcasts[member_index].push(job_ref); + member_data.semaphores[member_index].signal(); + }; + // SAFETY: `Spawn::spawn`'s contract requires that `!Send` work is only + // scheduled to run on this thread. The work is `Send` here, so this is + // vacuous. + unsafe { work.spawn(scheduler, None) } } /// Blocks the thread waiting for a future to complete. @@ -807,9 +769,9 @@ pub struct Membership { thread_pool: &'static ThreadPool, /// Contains the index of a row in the `MembersData` table, if the worker /// has been granted membership on the thread-pool. - member_index: usize, + pub(crate) member_index: usize, /// A reference to the `MemberData` table. - member_data: &'static MemberData, + pub(crate) member_data: &'static MemberData, } impl ThreadPool { @@ -840,7 +802,7 @@ impl ThreadPool { pub fn try_enroll(&'static self) -> Option { loop { let available_bitmask = - !self.claimed_bitmask.load(Ordering::Relaxed); + !self.claimed_bitmask.load(Ordering::Acquire); if available_bitmask == 0 { return None; } @@ -868,7 +830,7 @@ impl ThreadPool { } let member_data = self.get_member_data(); loop { - let claimed_bitmask = self.claimed_bitmask.load(Ordering::Relaxed); + let claimed_bitmask = self.claimed_bitmask.load(Ordering::Acquire); if claimed_bitmask == u32::MAX { return Vec::new(); } @@ -1041,7 +1003,6 @@ impl Membership { } } - // Indicate we are no longer waiting to resign. worker .thread_pool .wants_to_resign @@ -1051,19 +1012,16 @@ impl Membership { let thread_pool = worker.thread_pool; drop(worker); - // Complete the resignation + // Complete the resignation. loop { let resignations = thread_pool.resignations.load(Ordering::Relaxed); - // Try to decrement the resignations count. This uses `Release` - // ordering so that the the `claimed_bitmask` store done as part of - // `drop(worker)` appears to any thread waiting for this resignation - // to complete. + // Try to decrement the resignations count. if thread_pool .resignations .compare_exchange_weak( resignations, resignations - (1 << 26), - Ordering::Release, + Ordering::Relaxed, Ordering::Relaxed, ) .is_ok() @@ -1073,7 +1031,7 @@ impl Membership { if (resignations >> 26) - 1 == 0 { atomic_wait::wake_all(&*thread_pool.resignations); } - // Exit the CAS loop + // Exit the CAS loop. break; } } @@ -1120,7 +1078,7 @@ impl Drop for Membership { // Release the claim on this membership. self.thread_pool .claimed_bitmask - .fetch_and(!(1 << self.member_index), Ordering::Relaxed); + .fetch_and(!(1 << self.member_index), Ordering::Release); // In case another thread is waiting for a membership slot to free // up, issue a wake on the bitmask. atomic_wait::wake_one(&*self.thread_pool.claimed_bitmask); @@ -1150,7 +1108,7 @@ impl Drop for Membership { pub struct Worker { /// Registers the worker as belonging to a specific thread pool, and /// potentially also grants "membership" on that thread-pool. - membership: Membership, + pub(crate) membership: Membership, /// A sequence of jobs waiting to be executed. Newer jobs are executed /// before older ones, allowing efficient depth-first execution. During /// promotion, the oldest job is shared. Populated by `join()`. @@ -1171,7 +1129,7 @@ pub struct Worker { /// that a `Future` that is `!Send` and has been spawned onto this thread /// can be woken on another thread (the other thread then sends this thread /// a job that polls the future). - nonsend_fifo_queue: Arc>, + pub(crate) nonsend_fifo_queue: Arc>, /// A local psudorandom number-generator. Used to spread out /// worker-to-worker operations evenly across the pool. rng: XorShift64Star, @@ -1339,11 +1297,15 @@ impl Worker { let batch_job = HeapJob::new(move |worker| { worker.fifo_queue.append(job_refs); }); - // SAFETY: `into_job_ref` requires that the data closed over by the - // `HeapJob` outlive the `JobRef`. + // SAFETY: `HeapJob::into_job_ref` requires: + // + // * The data closed over by the `HeapJob` outlives the `JobRef`. + // + // The closure captures `job_refs` (a `VecDeque`) by value, + // so it trivially outlives the newly created `JobRef`. // - // Here, the closure captures `job_refs` (a `VecDequeue`) by - // value, and so trivially outlives the newly created `JobRef`. + // * If the closure is `!Send`, `JobRef::execute` only runs on this + // thread. The closure is `Send`, so this obligation is vacuous. let batch_job_ref = unsafe { batch_job.into_job_ref() }; // Push the batch job into the steal queue so siblings can steal it. if let Err(job_ref) = self.sharing_queue().push(batch_job_ref) { @@ -1410,39 +1372,39 @@ impl Worker { } } - /// Finds a job to work on. This function is almost entirely local, but does - /// a small amount of synchronization to allow for !Send futures that must - /// be polled on this thread but are woken on a different thread. + /// Finds a job to work on. + /// + /// This function will not steal work from other workers or read from the + /// global injector queue. It prioritizes work that _must_ be run on this + /// thread (assuming no other workers are running). /// /// Work is prioritized as follows: /// 1. Pull from the LIFO queue (`join` calls) /// 2. Pull from the !Send FIFO queue (`spawn_local` calls) - /// 3. Pull from the regular FIFO queue (`spawn` calls) + /// 3. Pull from the broadcast queue (`broadcast` and `spawn_on` calls) + /// 4. Pull from the regular FIFO queue (`spawn` calls) + /// 5. Reclaim work shared by this worker #[inline(always)] fn find_local_work(&self) -> Option { (self.lifo_queue.pop_newest()) .or_else(|| self.nonsend_fifo_queue.pop()) + .or_else(|| self.broadcast_queue().pop()) .or_else(|| self.fifo_queue.pop_oldest()) + .or_else(|| self.sharing_queue().pop()) } /// Finds a job to work on. /// /// Work is prioritized as follows: - /// 1. Pull from the LIFO queue (`join` calls) - /// 2. Pull from the !Send FIFO queue (`spawn_local` calls) - /// 3. Pull from the regular FIFO queue (`spawn` calls) - /// 4. Pull from the broadcast queue (`broadcast` calls) - /// 5. Reclaim work shared by this worker - /// 6. Steal work shared from other workers - /// 7. Read from the global queue (external calls) + /// 1. Try [`find_local_work`][Worker::find_local_work] first + /// 2. Steal work shared from other workers + /// 3. Read from the global queue (external calls) /// /// If work is found in the last two cases, it is treated as having been /// "migrated" to this thread. #[inline(always)] fn find_work(&self) -> Option<(JobRef, bool)> { (self.find_local_work().map(|job| (job, false))) - .or_else(|| self.broadcast_queue().pop().map(|job| (job, false))) - .or_else(|| self.sharing_queue().pop().map(|job| (job, false))) .or_else(|| self.steal_from_siblings().map(|job| (job, true))) .or_else(|| { self.thread_pool.shared_queue.pop().map(|job| (job, true)) @@ -1454,6 +1416,7 @@ impl Worker { /// Iterates over occupied seats in a random order to avoid always hitting /// the same victim. Because stealers are pre-allocated and permanent, no /// lock or atomic load is needed to access them. + #[inline(always)] fn steal_from_siblings(&self) -> Option { let my_member_index = self.membership.member_index; let my_sharer = self.sharing_queue(); @@ -1470,7 +1433,6 @@ impl Worker { let shifted_idx = bits.trailing_zeros(); let idx = (shifted_idx + offset) % 32; bits &= bits - 1; - // The stealer is a permanent reference — no lock or atomic load needed. let stealer = &stealers[idx as usize]; // `steal_and_pop` returns one job directly and moves up to half the // remaining items into our steal queue for later use. @@ -1583,8 +1545,42 @@ impl Worker { /// * If a future panics, the [`Task`] will panic when awaited. /// #[inline] - pub fn spawn>(&self, work: S) -> S::Output { - work.spawn(self.thread_pool, Some(self)) + pub fn spawn(&self, work: S) -> S::Output + where + S: Spawn + Send, + S::Output: Send, + { + let thread_pool = self.thread_pool; + let scheduler = |job_ref, worker: Option<&Worker>| match worker { + Some(worker) => worker.fifo_queue.push_new(job_ref), + None => thread_pool.push_shared_job(job_ref), + }; + // SAFETY: `Spawn::spawn`'s contract requires that `!Send` work only be + // scheduled to run on this thread. The work is `Send` and this is + // vacuous. + unsafe { work.spawn(scheduler, Some(self)) } + } + + /// Runs work (a closure or future) on the background of a specific worker + /// thread. + /// + /// This is quite similar to [`spawn`](Worker::spawn), except that the work + /// will run on another worker instead of the current one. + /// + /// # 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. + #[inline] + pub fn spawn_on(&self, member_index: usize, work: S) -> S::Output + where + S: Spawn + Send, + S::Output: Send, + { + self.thread_pool.spawn_on(member_index, work) } /// Runs work (a closure or future) in the background of this thread. @@ -1602,8 +1598,25 @@ impl Worker { /// * If a future panics, the [`Task`] will panic when awaited. /// #[inline] - pub fn spawn_local>(&self, work: S) -> S::Output { - work.spawn_local(self) + pub fn spawn_local(&self, work: S) -> S::Output + where + S: Spawn, + { + let queue = self.nonsend_fifo_queue.clone(); + let semaphore = &self.member_data.semaphores[self.member_index]; + let scheduler = move |job_ref, _: Option<&Worker>| { + queue.push(job_ref); + semaphore.signal(); + }; + // SAFETY: `Spawn::spawn`'s contract requires that any `!Send` work be + // scheduled to run on this thread. + // + // The scheduler pushes the job onto this worker's `nonsend_fifo_queue`, + // which is only ever drained by this worker on this thread (see + // `find_local_work`). For a future, every reschedule re-pushes onto + // that same queue regardless of which thread issued the wakeup, so the + // work stays on this thread. + unsafe { work.spawn(scheduler, Some(self)) } } /// Polls a future to completion, then returns the outcome. @@ -1910,7 +1923,7 @@ impl Worker { /// let ok: Vec = vec![1, 2, 3]; /// forte::scope(|scope| { /// let bad: Vec = vec![4, 5, 6]; - /// scope.spawn_on(worker, |_: &Worker| { + /// scope.spawn(|_: &Worker| { /// // Transfer ownership of `bad` into a local variable (also named `bad`). /// // This will force the closure to take ownership of `bad` from the environment. /// let bad = bad; @@ -1918,7 +1931,7 @@ impl Worker { /// println!("bad: {:?}", bad); // refers to our local variable, above. /// }); /// - /// scope.spawn_on(worker, |_: &Worker| println!("ok: {:?}", ok)); // we too can borrow `ok` + /// scope.spawn(|_: &Worker| println!("ok: {:?}", ok)); // we too can borrow `ok` /// }); /// # }); /// ``` @@ -1937,14 +1950,14 @@ impl Worker { /// let ok: Vec = vec![1, 2, 3]; /// forte::scope(|scope| { /// let bad: Vec = vec![4, 5, 6]; - /// scope.spawn_on(worker, move |_: &Worker| { + /// scope.spawn(move |_: &Worker| { /// println!("ok: {:?}", ok); /// println!("bad: {:?}", bad); /// }); /// /// // That closure is fine, but now we can't use `ok` anywhere else, /// // since it is owned by the previous task: - /// // scope.spawn_on(worker, |_: &Worker| println!("ok: {:?}", ok)); + /// // scope.spawn(|_: &Worker| println!("ok: {:?}", ok)); /// }); /// # }); /// ``` @@ -1964,7 +1977,7 @@ impl Worker { /// forte::scope(|scope| { /// let bad: Vec = vec![4, 5, 6]; /// let ok: &Vec = &ok; // shadow the original `ok` - /// scope.spawn_on(worker, move |_: &Worker| { + /// scope.spawn(move |_: &Worker| { /// println!("ok: {:?}", ok); // captures the shadowed version /// println!("bad: {:?}", bad); /// }); @@ -1973,7 +1986,7 @@ impl Worker { /// // can be shared freely. Note that we need a `move` closure here though, /// // because otherwise we'd be trying to borrow the shadowed `ok`, /// // and that doesn't outlive `scope`. - /// scope.spawn_on(worker, move |_: &Worker| println!("ok: {:?}", ok)); + /// scope.spawn(move |_: &Worker| println!("ok: {:?}", ok)); /// }); /// # }); /// ``` @@ -1990,7 +2003,7 @@ impl Worker { /// let ok: Vec = vec![1, 2, 3]; /// forte::scope(|scope| { /// let bad: Vec = vec![4, 5, 6]; - /// scope.spawn_on(worker, |_: &Worker| { + /// scope.spawn(|_: &Worker| { /// // Transfer ownership of `bad` into a local variable (also named `bad`). /// // This will force the closure to take ownership of `bad` from the environment. /// let bad = bad; @@ -1998,7 +2011,7 @@ impl Worker { /// println!("bad: {:?}", bad); // refers to our local variable, above. /// }); /// - /// scope.spawn_on(worker, |_: &Worker| println!("ok: {:?}", ok)); // we too can borrow `ok` + /// scope.spawn(|_: &Worker| println!("ok: {:?}", ok)); // we too can borrow `ok` /// }); /// # }); /// ``` @@ -2034,10 +2047,10 @@ impl Worker { /// let mut counter = 0; /// let counter_ref = &mut counter; /// forte::scope(|scope| { - /// scope.spawn_on(worker, |worker: &Worker| { + /// scope.spawn(|worker: &Worker| { /// *counter_ref += 1; /// // Note: we borrow the scope again here. - /// scope.spawn_on(worker, move |_: &Worker| { + /// scope.spawn(move |_: &Worker| { /// *counter_ref += 1; /// }); /// }); @@ -2061,7 +2074,7 @@ impl Worker { /// // ^^^^^ ERROR: This creates a *static* job on the worker, /// // which may outlive the scope. /// - /// scope.spawn_on(worker, |_: &Worker| { }); + /// scope.spawn(|_: &Worker| { }); /// // ^^^^^ ERROR: This requires borrowing the scope within the /// // unscoped job, which isn't allowed by the compiler /// // because 'scope would have to to outlive 'static. @@ -2230,12 +2243,19 @@ impl Worker { return; } + // Wrap the operation in an `Arc` so each per-member job can own an + // independent, `'static` reference to it. We must *not* let a job + // capture a borrow of `f`: `spawn_broadcast` does not wait, so `f` + // would be dropped while jobs are still queued or running on other + // threads. + let f = Arc::new(f); + // Send the broadcast to each member, and wake them up. for (i, member_index) in members.iter_bits().enumerate() { - let func = &f; + let func = Arc::clone(&f); let op = move |worker: &Worker| { // Run the job - func(Broadcast { + (*func)(Broadcast { worker, index: i, participants, @@ -2247,15 +2267,17 @@ impl Worker { // SAFETY: `HeapJob::into_job_ref` has two preconditions: // // * The `JobRef` must not outlive any of the items closed over by - // the function `f`. + // `op`. // - // Since `F: 'static`, the `JobRef` cannot outlive its captured - // data. + // `op` owns an `Arc` clone (`func`). Since `F: 'static`, that + // `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. // - // The `op` is `Send`, so this does not apply. + // `op` is `Send` (`Arc: Send` because `F: Send + Sync`), so + // this does not apply. let job_ref = unsafe { job.into_job_ref() }; self.member_data.broadcasts[member_index].push(job_ref); self.member_data.semaphores[member_index].signal(); @@ -2302,7 +2324,7 @@ pub static DEFAULT_POOL: DefaultThreadPool = DefaultThreadPool { initialized: AtomicU32::new(0), }; -/// Runs the provided closure in the background. +/// Runs work in the background. /// /// When executed on a thread that is currently registered as a worker (i.e. the /// closure inside [`Membership::activate`], [`ThreadPool::with_worker`], or @@ -2314,13 +2336,38 @@ pub static DEFAULT_POOL: DefaultThreadPool = DefaultThreadPool { /// If you have a reference to a [`Worker`], it's better to use [`Worker::spawn`] /// instead. If you don't have a worker, but know which thread pool you want to /// use, [`ThreadPool::spawn`] is more appropriate. -pub fn spawn>(work: S) -> S::Output { +pub fn spawn(work: S) -> S::Output +where + S: Spawn + Send, + S::Output: Send, +{ Worker::with_current(|worker| match worker { Some(worker) => worker.spawn(work), None => DEFAULT_POOL.spawn(work), }) } +/// Runs work on the background on a specific worker thread. +/// +/// This is quite similar to [`spawn`], except that the work will run on another +/// worker instead of the current one. +/// +/// If not called within a thread pool, this uses the [`DEFAULT_POOL`]. +/// +/// If you have a reference to a [`Worker`], it's better to use +/// [`Worker::spawn_on`] instead. If you don't have a worker, but know which +/// thread pool you want to use, [`ThreadPool::spawn_on`] is more appropriate. +pub fn spawn_on(member_index: usize, work: S) -> S::Output +where + S: Spawn + Send, + S::Output: Send, +{ + Worker::with_current(|worker| match worker { + Some(worker) => worker.spawn_on(member_index, work), + None => DEFAULT_POOL.spawn_on(member_index, work), + }) +} + /// Waits for a future to complete. /// /// When executed on a thread that is currently registered as a worker (i.e. the diff --git a/src/time.rs b/src/time.rs index 451d5e7..b6879c1 100644 --- a/src/time.rs +++ b/src/time.rs @@ -25,8 +25,8 @@ 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 a general-purpose register and + // does not access Rust memory. unsafe { asm!( "rdtime {}", @@ -37,11 +37,15 @@ pub fn ticks() -> u64 { cnt } -/// Read from the real-time stamp counter on windows +/// Read from the real-time stamp counter on x86 / x86-64. #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] #[inline(always)] pub fn ticks() -> u64 { - // SAFETY: `_rdtsc` emits the CPU counter read instruction and has no Rust memory safety - // preconditions. - unsafe { core::arch::x86_64::_rdtsc() } + #[cfg(target_arch = "x86")] + 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. + unsafe { _rdtsc() } }