任务窃取与空闲

pub(super) struct Idle {
    /// Tracks both the number of searching workers and the number of unparked
    /// workers.
    ///
    /// Used as a fast-path to avoid acquiring the lock when needed.
    state: AtomicUsize,

    /// Total number of workers.
    num_workers: usize,
}

/// Data synchronized by the scheduler mutex
pub(super) struct Synced {
    /// Sleeping workers
    sleepers: Vec<usize>,
}

const UNPARK_SHIFT: usize = 16;
const UNPARK_MASK: usize = !SEARCH_MASK;
const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1;

pub(super) fn transition_worker_to_searching(&self) -> bool {
    let state = State::load(&self.state, SeqCst);
    if 2 * state.num_searching() >= self.num_workers {
        return false;
    }

    // It is possible for this routine to allow more than 50% of the workers
    // to search. That is OK. Limiting searchers is only an optimization to
    // prevent too much contention.
    State::inc_num_searching(&self.state, SeqCst);
    true
}

fn steal_work(&mut self, worker: &Worker) -> Option<Notified> {
    if !self.transition_to_searching(worker) {
        return None;
    }

    let num = worker.handle.shared.remotes.len();
    // Start from a random worker
    let start = self.rand.fastrand_n(num as u32) as usize;

    for i in 0..num {
        let i = (start + i) % num;

        // Don't steal from ourself! We know we don't have work.
        if i == worker.index {
            continue;
        }

        let target = &worker.handle.shared.remotes[i];
        if let Some(task) = target
            .steal
            .steal_into(&mut self.run_queue, &mut self.stats)
        {
            return Some(task);
        }
    }

    // Fallback on checking the global queue
    worker.handle.next_remote_task()
}

/// Returns `true` if another parked worker was notified, `false` otherwise.
fn transition_worker_from_searching(&self) -> bool {
    if self.shared.idle.transition_worker_from_searching() {
        // We are the final searching worker. Because work was found, we
        // need to notify another worker.
        self.notify_parked_local()
    } else {
        false
    }
}

pub(super) fn worker_to_notify(&self, shared: &Shared) -> Option<usize> {
    // If at least one worker is spinning, work being notified will
    // eventually be found. A searching thread will find **some** work and
    // notify another worker, eventually leading to our work being found.
    //
    // For this to happen, this load must happen before the thread
    // transitioning `num_searching` to zero. Acquire / Release does not
    // provide sufficient guarantees, so this load is done with `SeqCst` and
    // will pair with the `fetch_sub(1)` when transitioning out of
    // searching.
    if !self.notify_should_wakeup() {
        return None;
    }

    // Acquire the lock
    let mut lock = shared.synced.lock();

    // Check again, now that the lock is acquired
    if !self.notify_should_wakeup() {
        return None;
    }

    // A worker should be woken up, atomically increment the number of
    // searching workers as well as the number of unparked workers.
    State::unpark_one(&self.state, 1);

    // Get the worker to unpark
    let ret = lock.idle.sleepers.pop();
    debug_assert!(ret.is_some());

    ret
}

fn notify_should_wakeup(&self) -> bool {
    // This must be a `SeqCst` RMW rather than a load: it is what makes
    // the caller's preceding inject-queue push visible to a parking
    // worker's subsequent unlocked queue-emptiness check.
    let state = State(self.state.fetch_add(0, SeqCst));
    state.num_searching() == 0 && state.num_unparked() < self.num_workers
}

fn transition_to_parked(&mut self, worker: &Worker) -> bool {
    // Workers should not park if they have work to do or are about to be traced
    if self.has_tasks() || self.is_traced {
        return false;
    }

    // When the final worker transitions **out** of searching to parked, it
    // must check all the queues one last time in case work materialized
    // between the last work scan and transitioning out of searching.
    let is_last_searcher = worker.handle.shared.idle.transition_worker_to_parked(
        &worker.handle.shared,
        worker.index,
        self.is_searching,
    );

    // The worker is no longer searching. Setting this is the local cache
    // only.
    self.is_searching = false;

    if is_last_searcher {
        worker.handle.notify_if_work_pending();
    }

    true
}

fn transition_from_parked(&mut self, worker: &Worker) -> bool {
    // If a task is in the lifo slot/run queue, then we must unpark regardless of
    // being notified. Same when woken to be traced: a dump unparks worker
    // threads without going through `Idle`, so the worker does it here.
    if self.has_tasks() || self.is_traced {
        // When a worker wakes, it should only transition to the "searching"
        // state when the wake originates from another worker *or* a new task
        // is pushed. We do *not* want the worker to transition to "searching"
        // when it wakes when the I/O driver receives new events.
        self.is_searching = !worker
            .handle
            .shared
            .idle
            .unpark_worker_by_id(&worker.handle.shared, worker.index);
        return true;
    }

    if worker
        .handle
        .shared
        .idle
        .is_parked(&worker.handle.shared, worker.index)
    {
        return false;
    }

    // When unparked, the worker is in the searching state.
    self.is_searching = true;
    true
}

const EMPTY: usize = 0;
const PARKED_CONDVAR: usize = 1;
const PARKED_DRIVER: usize = 2;
const NOTIFIED: usize = 3;

fn park(&self, handle: &driver::Handle) -> HadDriver {
    // If we were previously notified then we consume this notification and
    // return quickly.
    if self
        .state
        .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
        .is_ok()
    {
        return HadDriver::No;
    }

    if let Some(mut driver) = self.shared.driver.try_lock() {
        self.park_driver(&mut driver, handle, None)
    } else {
        self.park_condvar(None);
        HadDriver::No
    }
}

match self.state.swap(NOTIFIED, SeqCst) {
    EMPTY => {}    // no one was waiting
    NOTIFIED => {} // already unparked
    PARKED_CONDVAR => self.unpark_condvar(),
    PARKED_DRIVER => driver.unpark(),
    actual => panic!("inconsistent state in unpark; actual = {actual}"),
}

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

#[tokio::main(flavor = "multi_thread", worker_threads = 4)]
async fn main() {
  let seen = Arc::new(Mutex::new(HashMap::new()));
  let s = seen.clone();
  tokio::spawn(async move {
    let mut hs = Vec::new();
    for _ in 0..64 {
      let s = s.clone();
      hs.push(tokio::spawn(async move {
        let t = Instant::now();
        while t.elapsed() < Duration::from_millis(5) {}
        let name = format!("{:?}", std::thread::current().id());
        *s.lock().unwrap().entry(name).or_insert(0) += 1;
      }));
    }
    for h in hs {
      h.await.unwrap();
    }
  })
  .await
  .unwrap();
  let mut v: Vec<_> = seen.lock().unwrap().clone().into_iter().collect();
  v.sort();
  println!("{} threads ran the 64 tasks: {:?}", v.len(), v);
}

4 threads ran the 64 tasks: [("ThreadId(2)", 16), ("ThreadId(3)", 16), ("ThreadId(4)", 16), ("ThreadId(5)", 16)]