executor: rewrite the work-stealing thread pool (#1657)

This patch is a ground up rewrite of the existing work-stealing thread pool. The goal is to reduce overhead while simplifying code when possible. At a high level, the following architectural changes were made: - The local run queues were switched for bounded circle buffer queues. - Reduce cross-thread synchronization. - Refactor task constructs to use a single allocation and always include a join handle (#887). - Simplify logic around putting workers to sleep and waking them up. **Local run queues** Move away from crossbeam's implementation of the Chase-Lev deque. This implementation included unnecessary overhead as it supported capabilities that are not needed for the work-stealing thread pool. Instead, a fixed size circle buffer is used for the local queue. When the local queue is full, half of the tasks contained in it are moved to the global run queue. **Reduce cross-thread synchronization** This is done via many small improvements. Primarily, an upper bound is placed on the number of concurrent stealers. Limiting the number of stealers results in lower contention. Secondly, the rate at which workers are notified and woken up is throttled. This also reduces contention by preventing many threads from racing to steal work. **Refactor task structure** Now that Tokio is able to target a rust version that supports `std::alloc` as well as `std::task`, the pool is able to optimize how the task structure is laid out. Now, a single allocation per task is required and a join handle is always provided enabling the spawner to retrieve the result of the task (#887). **Simplifying logic** When possible, complexity is reduced in the implementation. This is done by using locks and other simpler constructs in cold paths. The set of sleeping workers is now represented as a `Mutex<VecDeque<usize>>`. Instead of optimizing access to this structure, we reduce the amount the pool must access this structure. Secondly, we have (temporarily) removed `threadpool::blocking`. This capability will come back later, but the original implementation was way more complicated than necessary. **Results** The thread pool benchmarks have improved significantly: Old thread pool: ``` test chained_spawn ... bench: 2,019,796 ns/iter (+/- 302,168) test ping_pong ... bench: 1,279,948 ns/iter (+/- 154,365) test spawn_many ... bench: 10,283,608 ns/iter (+/- 1,284,275) test yield_many ... bench: 21,450,748 ns/iter (+/- 1,201,337) ``` New thread pool: ``` test chained_spawn ... bench: 147,943 ns/iter (+/- 6,673) test ping_pong ... bench: 537,744 ns/iter (+/- 20,928) test spawn_many ... bench: 7,454,898 ns/iter (+/- 283,449) test yield_many ... bench: 16,771,113 ns/iter (+/- 733,424) ``` Real-world benchmarks improve significantly as well. This is testing the hyper hello world server using: `wrk -t1 -c50 -d10`: Old scheduler: ``` Running 10s test @ http://127.0.0.1:3000 1 threads and 50 connections Thread Stats Avg Stdev Max +/- Stdev Latency 371.53us 99.05us 1.97ms 60.53% Req/Sec 114.61k 8.45k 133.85k 67.00% 1139307 requests in 10.00s, 95.61MB read Requests/sec: 113923.19 Transfer/sec: 9.56MB ``` New scheduler: ``` Running 10s test @ http://127.0.0.1:3000 1 threads and 50 connections Thread Stats Avg Stdev Max +/- Stdev Latency 275.05us 69.81us 1.09ms 73.57% Req/Sec 153.17k 10.68k 171.51k 71.00% 1522671 requests in 10.00s, 127.79MB read Requests/sec: 152258.70 Transfer/sec: 12.78MB ```
2026-08-25 00:00:18 +02:00 · 2019-10-19 11:09:40 -07:00
parent 2a181320b7
commit ed5a94eb2d
100 changed files with 7409 additions and 6796 deletions
@@ -0,0 +1,502 @@
+use crate::loom::sync::atomic::AtomicUsize;
+
+use std::fmt;
+use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
+use std::usize;
+
+pub(super) struct State {
+    val: AtomicUsize,
+}
+
+/// Current state value
+#[derive(Copy, Clone)]
+pub(super) struct Snapshot(usize);
+
+/// The task is currently being run.
+const RUNNING: usize = 0b00_0001;
+
+/// The task has been notified by a waker.
+const NOTIFIED: usize = 0b00_0010;
+
+/// The task is complete.
+///
+/// Once this bit is set, it is never unset
+const COMPLETE: usize = 0b00_0100;
+
+/// The primary task handle has been dropped.
+const RELEASED: usize = 0b00_1000;
+
+/// The join handle is still around
+const JOIN_INTEREST: usize = 0b01_0000;
+
+/// A join handle waker has been set
+const JOIN_WAKER: usize = 0b10_0000;
+
+/// The task has been forcibly canceled.
+const CANCELLED: usize = 0b100_0000;
+
+/// All bits
+const LIFECYCLE_MASK: usize =
+    RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED;
+
+/// Bits used by the waker ref count portion of the state.
+///
+/// Ref counts only cover **wakers**. Other handles are tracked with other state
+/// bits.
+const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK;
+
+/// Number of positions to shift the ref count
+const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize;
+
+/// One ref count
+const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT;
+
+/// Initial state
+const INITIAL_STATE: usize = NOTIFIED;
+
+/// All transitions are performed via RMW operations. This establishes an
+/// unambiguous modification order.
+impl State {
+    /// Starts with a ref count of 1
+    pub(super) fn new_background() -> State {
+        State {
+            val: AtomicUsize::new(INITIAL_STATE),
+        }
+    }
+
+    /// Starts with a ref count of 2
+    pub(super) fn new_joinable() -> State {
+        State {
+            val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST),
+        }
+    }
+
+    /// Load the current state, establishes `Acquire` ordering.
+    pub(super) fn load(&self) -> Snapshot {
+        Snapshot(self.val.load(Acquire))
+    }
+
+    /// Transition a task to the `Running` state.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn transition_to_running(&self) -> Snapshot {
+        const DELTA: usize = RUNNING | NOTIFIED;
+
+        let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire));
+        debug_assert!(prev.is_notified());
+
+        if prev.is_running() {
+            // We were signalled to cancel
+            //
+            // Apply the state
+            let prev = self.val.fetch_or(CANCELLED, AcqRel);
+            return Snapshot(prev | CANCELLED);
+        }
+
+        debug_assert!(!prev.is_running());
+
+        let next = Snapshot(prev.0 ^ DELTA);
+
+        debug_assert!(next.is_running());
+        debug_assert!(!next.is_notified());
+
+        next
+    }
+
+    /// Transition the task from `Running` -> `Idle`.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn transition_to_idle(&self) -> Snapshot {
+        const DELTA: usize = RUNNING;
+
+        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
+
+        if !prev.is_running() {
+            // We were signaled to cancel.
+            //
+            // Apply the state
+            let prev = self.val.fetch_or(CANCELLED, AcqRel);
+            return Snapshot(prev | CANCELLED);
+        }
+
+        let next = Snapshot(prev.0 ^ DELTA);
+
+        debug_assert!(!next.is_running());
+
+        next
+    }
+
+    /// Transition the task from `Running` -> `Complete`.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn transition_to_complete(&self) -> Snapshot {
+        const DELTA: usize = RUNNING | COMPLETE;
+
+        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
+
+        debug_assert!(!prev.is_complete());
+
+        let next = Snapshot(prev.0 ^ DELTA);
+
+        debug_assert!(next.is_complete());
+
+        next
+    }
+
+    /// Transition the task from `Running` -> `Released`.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn transition_to_released(&self) -> Snapshot {
+        const DELTA: usize = RUNNING | COMPLETE | RELEASED;
+
+        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
+
+        debug_assert!(prev.is_running());
+        debug_assert!(!prev.is_complete());
+        debug_assert!(!prev.is_released());
+
+        let next = Snapshot(prev.0 ^ DELTA);
+
+        debug_assert!(!next.is_running());
+        debug_assert!(next.is_complete());
+        debug_assert!(next.is_released());
+
+        next
+    }
+
+    /// Transition the task to the canceled state.
+    ///
+    /// Returns the snapshot of the state **after** the transition **if** the
+    /// transition was made successfully
+    ///
+    /// # States
+    ///
+    /// - Notifed: task may be in a queue, caller must not release.
+    /// - Running: cannot drop. The poll handle will handle releasing.
+    /// - Other prior states do not require cancellation.
+    ///
+    /// If the task has been notified, then it may still be in a queue. The
+    /// caller must not release the task.
+    pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot {
+        let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel));
+
+        debug_assert!(!prev.is_complete());
+        debug_assert!(!prev.is_running() || prev.is_notified());
+
+        Snapshot(prev.0 | CANCELLED)
+    }
+
+    pub(super) fn transition_to_canceled_from_list(&self) -> Option<Snapshot> {
+        let mut prev = self.load();
+
+        loop {
+            if !prev.is_active() {
+                return None;
+            }
+
+            let mut next = prev;
+
+            // Use the running flag to signal cancellation
+            if prev.is_running() {
+                next.0 -= RUNNING;
+            } else if prev.is_notified() {
+                next.0 += RUNNING;
+            } else {
+                next.0 |= CANCELLED;
+            }
+
+            let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire);
+
+            match res {
+                Ok(_) if next.is_canceled() => return Some(next),
+                Ok(_) => return None,
+                Err(actual) => {
+                    prev = Snapshot(actual);
+                }
+            }
+        }
+    }
+
+    /// Final transition to `Released`. Called when primary task handle is
+    /// dropped. This is roughly a "ref decrement" operation.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn release_task(&self) -> Snapshot {
+        use crate::loom::sync::atomic;
+
+        const DELTA: usize = RELEASED;
+
+        let prev = Snapshot(self.val.fetch_or(DELTA, Release));
+
+        debug_assert!(!prev.is_released());
+        debug_assert!(prev.is_terminal(), "state = {:?}", prev);
+
+        let next = Snapshot(prev.0 | DELTA);
+
+        debug_assert!(next.is_released());
+
+        if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) {
+            // The final reference to the task was dropped, the caller must free the
+            // memory. Establish an acquire ordering.
+            atomic::fence(Acquire);
+        }
+
+        next
+    }
+
+    /// Transition the state to `Scheduled`.
+    ///
+    /// Returns `true` if the task needs to be submitted to the pool for
+    /// execution
+    pub(super) fn transition_to_notified(&self) -> bool {
+        const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED;
+
+        let prev = self.val.fetch_or(NOTIFIED, Release);
+        prev & MASK == 0
+    }
+
+    /// Optimistically try to swap the state assuming the join handle is
+    /// __immediately__ dropped on spawn
+    pub(super) fn drop_join_handle_fast(&self) -> bool {
+        use std::sync::atomic::Ordering::Relaxed;
+
+        // Relaxed is acceptable as if this function is called and succeeds,
+        // then nothing has been done w/ the join handle.
+        //
+        // The moment the join handle is used (polled), the `JOIN_WAKER` flag is
+        // set, at which point the CAS will fail.
+        //
+        // Given this, there is no risk if this operation is reordered.
+        self.val
+            .compare_exchange_weak(
+                INITIAL_STATE | JOIN_INTEREST,
+                INITIAL_STATE,
+                Relaxed,
+                Relaxed,
+            )
+            .is_ok()
+    }
+
+    /// The join handle has completed by reading the output
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn complete_join_handle(&self) -> Snapshot {
+        use crate::loom::sync::atomic;
+
+        const DELTA: usize = JOIN_INTEREST;
+
+        let prev = Snapshot(self.val.fetch_sub(DELTA, Release));
+
+        debug_assert!(prev.is_join_interested());
+
+        let next = Snapshot(prev.0 - DELTA);
+
+        if !next.is_final_ref() {
+            return next;
+        }
+
+        atomic::fence(Acquire);
+
+        next
+    }
+
+    /// The join handle is being dropped, this fails if the task has been
+    /// completed and the output must be dropped first then
+    /// `complete_join_handle` should be called.
+    ///
+    /// Returns a snapshot of the state **after** the transition.
+    pub(super) fn drop_join_handle_slow(&self) -> Result<Snapshot, Snapshot> {
+        const MASK: usize = COMPLETE | CANCELLED;
+
+        let mut prev = self.val.load(Acquire);
+
+        loop {
+            // Once the complete bit is set, it is never unset.
+            if prev & MASK != 0 {
+                return Err(Snapshot(prev));
+            }
+
+            debug_assert!(prev & JOIN_INTEREST == JOIN_INTEREST);
+
+            let next = (prev - JOIN_INTEREST) & !JOIN_WAKER;
+
+            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);
+
+            match res {
+                Ok(_) => {
+                    return Ok(Snapshot(next));
+                }
+                Err(actual) => {
+                    prev = actual;
+                }
+            }
+        }
+    }
+
+    /// Store the join waker.
+    pub(super) fn store_join_waker(&self) -> Snapshot {
+        use crate::loom::sync::atomic;
+
+        const DELTA: usize = JOIN_WAKER;
+
+        let prev = Snapshot(self.val.fetch_xor(DELTA, Release));
+
+        debug_assert!(!prev.has_join_waker());
+
+        let next = Snapshot(prev.0 ^ DELTA);
+
+        debug_assert!(next.has_join_waker());
+
+        if next.is_complete() {
+            atomic::fence(Acquire);
+        }
+
+        next
+    }
+
+    pub(super) fn unset_waker(&self) -> Snapshot {
+        const MASK: usize = COMPLETE | CANCELLED;
+
+        let mut prev = self.val.load(Acquire);
+
+        loop {
+            // Once the `COMPLETE` bit is set, it is never unset
+            if prev & MASK != 0 {
+                return Snapshot(prev);
+            }
+
+            debug_assert!(Snapshot(prev).has_join_waker());
+
+            let next = prev - JOIN_WAKER;
+
+            let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);
+
+            match res {
+                Ok(_) => return Snapshot(next),
+                Err(actual) => {
+                    prev = actual;
+                }
+            }
+        }
+    }
+
+    pub(super) fn ref_inc(&self) {
+        use std::process;
+        use std::sync::atomic::Ordering::Relaxed;
+
+        // Using a relaxed ordering is alright here, as knowledge of the
+        // original reference prevents other threads from erroneously deleting
+        // the object.
+        //
+        // As explained in the [Boost documentation][1], Increasing the
+        // reference counter can always be done with memory_order_relaxed: New
+        // references to an object can only be formed from an existing
+        // reference, and passing an existing reference from one thread to
+        // another must already provide any required synchronization.
+        //
+        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
+        let prev = self.val.fetch_add(WAKER_ONE, Relaxed);
+
+        // If the reference count overflowed, abort.
+        if prev > isize::max_value() as usize {
+            process::abort();
+        }
+    }
+
+    /// Returns `true` if the task should be released.
+    pub(super) fn ref_dec(&self) -> bool {
+        use crate::loom::sync::atomic;
+
+        let prev = self.val.fetch_sub(WAKER_ONE, Release);
+        let next = Snapshot(prev - WAKER_ONE);
+
+        if next.is_final_ref() {
+            atomic::fence(Acquire);
+        }
+
+        next.is_final_ref()
+    }
+}
+
+impl Snapshot {
+    pub(super) fn is_running(self) -> bool {
+        self.0 & RUNNING == RUNNING
+    }
+
+    pub(super) fn is_notified(self) -> bool {
+        self.0 & NOTIFIED == NOTIFIED
+    }
+
+    pub(super) fn is_released(self) -> bool {
+        self.0 & RELEASED == RELEASED
+    }
+
+    pub(super) fn is_complete(self) -> bool {
+        self.0 & COMPLETE == COMPLETE
+    }
+
+    pub(super) fn is_canceled(self) -> bool {
+        self.0 & CANCELLED == CANCELLED
+    }
+
+    /// Used during normal runtime.
+    pub(super) fn is_active(self) -> bool {
+        self.0 & (COMPLETE | CANCELLED) == 0
+    }
+
+    /// Used before dropping the task
+    pub(super) fn is_terminal(self) -> bool {
+        // When both the notified & running flags are set, the task was canceled
+        // after being notified, before it was run.
+        //
+        // There is a race where:
+        // - The task state transitions to notified
+        // - The global queue is shutdown
+        // - The waker attempts to push into the global queue and fails.
+        // - The waker holds the last reference to the task, thus drops it.
+        //
+        // In this scenario, the cancelled bit will never get set.
+        !self.is_active() || (self.is_notified() && self.is_running())
+    }
+
+    pub(super) fn is_join_interested(self) -> bool {
+        self.0 & JOIN_INTEREST == JOIN_INTEREST
+    }
+
+    pub(super) fn has_join_waker(self) -> bool {
+        self.0 & JOIN_WAKER == JOIN_WAKER
+    }
+
+    pub(super) fn is_final_ref(self) -> bool {
+        const MASK: usize = WAKER_COUNT_MASK | RELEASED | JOIN_INTEREST;
+
+        (self.0 & MASK) == RELEASED
+    }
+}
+
+impl fmt::Debug for State {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use std::sync::atomic::Ordering::SeqCst;
+
+        let snapshot = Snapshot(self.val.load(SeqCst));
+
+        fmt.debug_struct("State")
+            .field("snapshot", &snapshot)
+            .finish()
+    }
+}
+
+impl fmt::Debug for Snapshot {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("Snapshot")
+            .field("is_running", &self.is_running())
+            .field("is_notified", &self.is_notified())
+            .field("is_released", &self.is_released())
+            .field("is_complete", &self.is_complete())
+            .field("is_canceled", &self.is_canceled())
+            .field("is_join_interested", &self.is_join_interested())
+            .field("has_join_waker", &self.has_join_waker())
+            .field("is_final_ref", &self.is_final_ref())
+            .finish()
+    }
+}