最小工作窃取线程池

pub(crate) struct Shared {
  injector: Injector<JobRef>,
  stealers: Vec<Stealer<JobRef>>,
  pub(crate) sleep: Sleep,
  terminate: AtomicBool,
}

pub(crate) struct WorkerThread {
  worker: Worker<JobRef>,
  index: usize,
  pub(crate) shared: Arc<Shared>,
  rng: Cell<u64>,
}

thread_local! {
  static CURRENT: Cell<*const WorkerThread> = const { Cell::new(ptr::null()) };
}

fn find_work(&self) -> Option<JobRef> {
  self.worker.pop().or_else(|| self.steal()).or_else(|| loop {
    match self.shared.injector.steal() {
      Steal::Success(job) => return Some(job),
      Steal::Empty => return None,
      Steal::Retry => {}
    }
  })
}

fn steal(&self) -> Option<JobRef> {
  let n = self.shared.stealers.len();
  if n <= 1 {
    return None;
  }
  loop {
    let mut retry = false;
    let start = (self.next_random() % n as u64) as usize;
    for i in (start..n).chain(0..start) {
      if i == self.index {
        continue;
      }
      match self.shared.stealers[i].steal() {
        Steal::Success(job) => return Some(job),
        Steal::Retry => retry = true,
        Steal::Empty => {}
      }
    }
    if !retry {
      return None;
    }
  }
}

pub(crate) fn sleep(&self, still_idle: impl Fn() -> bool) {
  let guard = lock(&self.lock);
  self.sleepers.fetch_add(1, Ordering::SeqCst);
  // Pairs with the fence in wake_*: either the waker sees sleepers > 0,
  // or we see its job / latch in the re-check below.
  fence(Ordering::SeqCst);
  #[cfg(not(feature = "no_recheck"))]
  if !still_idle() {
    self.sleepers.fetch_sub(1, Ordering::SeqCst);
    return;
  }
  #[cfg(feature = "no_recheck")]
  let _ = &still_idle;
  let guard = self.cv.wait(guard).unwrap_or_else(|e| e.into_inner());
  self.sleepers.fetch_sub(1, Ordering::SeqCst);
  drop(guard);
}

pub(crate) fn wake_one(&self) {
  fence(Ordering::SeqCst);
  if self.sleepers.load(Ordering::SeqCst) > 0 {
    // Taking the lock means a sleeper is either not registered yet
    // (and will re-check) or already inside cv.wait.
    let _guard = lock(&self.lock);
    // The blocking reverse variants have sleepers that never look for
    // work; a single notify could land on one of them, so wake everyone.
    if cfg!(any(feature = "block_wait", feature = "no_steal_wait")) {
      self.cv.notify_all();
    } else {
      self.cv.notify_one();
    }
  }
}

pub(crate) fn in_worker<OP, R>(&self, op: OP) -> R
where
  OP: FnOnce(&WorkerThread, bool) -> R + Send,
  R: Send,
{
  let wt = WorkerThread::current();
  if !wt.is_null() {
    let wt = unsafe { &*wt };
    if Arc::ptr_eq(&wt.shared, &self.shared) {
      return op(wt, false);
    }
  }
  let job = StackJob::new(LockLatch::new(), |injected| {
    let wt = WorkerThread::current();
    if wt.is_null() {
      unreachable!("injected job runs on a worker thread");
    }
    op(unsafe { &*wt }, injected)
  });
  self.shared.inject(unsafe { JobRef::new(&job) });
  job.latch.wait();
  job.into_result()
}

STUCK: 2653 installs finished, then nothing for the rest of 10 s
STUCK: 136 installs finished, then nothing for the rest of 10 s
STUCK: 869 installs finished, then nothing for the rest of 10 s
STUCK: 312 installs finished, then nothing for the rest of 10 s
STUCK: 9929 installs finished, then nothing for the rest of 10 s
STUCK: 10825 installs finished, then nothing for the rest of 10 s
STUCK: 8144 installs finished, then nothing for the rest of 10 s
STUCK: 1671 installs finished, then nothing for the rest of 10 s
STUCK: 6656 installs finished, then nothing for the rest of 10 s
STUCK: 1014 installs finished, then nothing for the rest of 10 s