最小 join

pub(crate) struct StackJob<L, F, R> {
  pub(crate) latch: L,
  func: UnsafeCell<Option<F>>,
  result: UnsafeCell<Option<R>>,
}

impl<L: Latch, F: FnOnce(bool) -> R, R> Job for StackJob<L, F, R> {
  unsafe fn execute(this: *const ()) {
    let this = this as *const Self;
    let func = match unsafe { (*(*this).func.get()).take() } {
      Some(func) => func,
      None => unreachable!("job executed twice"),
    };
    let abort = AbortIfPanic;
    // Reached through a JobRef: stolen or injected, so "migrated".
    let r = func(true);
    mem::forget(abort);
    unsafe { *(*this).result.get() = Some(r) };
    // Last touch of `*this`: after set() the owner may pop the frame.
    unsafe { L::set(&(*this).latch) };
  }
}

pub(crate) trait Latch {
  /// Takes a raw pointer, not `&self`: once the flag is set the waiter may
  /// return and free the latch, so `set` must not touch it afterwards.
  unsafe fn set(this: *const Self);
}

impl Latch for SpinLatch<'_> {
  unsafe fn set(this: *const Self) {
    // Read everything we need first. `Shared` itself stays alive: the setter
    // is a worker of the same pool and holds an Arc to it.
    let shared: &Shared = unsafe { (*this).shared };
    unsafe { (*this).done.store(true, Ordering::Release) };
    // From here on `this` may dangle.
    shared.sleep.wake_all();
  }
}

fn join_in_worker<A, B, RA, RB>(wt: &WorkerThread, injected: bool, a: A, b: B) -> (RA, RB)
where
  A: FnOnce(bool) -> RA + Send,
  B: FnOnce(bool) -> RB + Send,
  RA: Send,
  RB: Send,
{
  let job_b = StackJob::new(SpinLatch::new(&wt.shared), b);
  let job_b_ref = unsafe { JobRef::new(&job_b) };
  wt.push(job_b_ref);

  // If A panics, B may be running on another thread and still point into
  // this frame. rayon waits for B and then resumes the panic; we abort.
  let guard = AbortIfPanic;
  let ra = a(injected);

  #[cfg(not(feature = "block_wait"))]
  while !job_b.latch.probe() {
    match wt.pop() {
      Some(job) if job == job_b_ref => {
        // Nobody stole B: run it here, no latch involved.
        mem::forget(guard);
        let rb = job_b.run_inline(injected);
        return (ra, rb);
      }
      // Something A left behind (can't happen with join alone, can with spawn).
      Some(job) => unsafe { wt.execute(job) },
      None => {
        // B was stolen: wait, executing other jobs meanwhile.
        wt.wait_until(&job_b.latch);
        break;
      }
    }
  }
  #[cfg(feature = "block_wait")]
  wt.block_until(&job_b.latch);

  mem::forget(guard);
  (ra, job_b.into_result())
}

pub(crate) fn wait_until(&self, latch: &SpinLatch<'_>) {
  while !latch.probe() {
    if let Some(job) = self.find_work() {
      unsafe { job.execute() };
      continue;
    }
    let shared = &self.shared;
    shared.sleep.sleep(|| !latch.probe() && !shared.has_work());
  }
}