Latch 家族

pub(super) trait Latch {
    /// Set the latch, signalling others.
    ///
    /// # WARNING
    ///
    /// Setting a latch triggers other threads to wake up and (in some
    /// cases) complete. This may, in turn, cause memory to be
    /// deallocated and so forth. One must be very careful about this,
    /// and it's typically better to read all the fields you will need
    /// to access *before* a latch is set!
    ///
    /// This function operates on `*const Self` instead of `&self` to allow it
    /// to become dangling during this call. The caller must ensure that the
    /// pointer is valid upon entry, and not invalidated during the call by any
    /// actions other than `set` itself.
    unsafe fn set(this: *const Self);
}

pub(super) trait AsCoreLatch {
    fn as_core_latch(&self) -> &CoreLatch;
}

/// Latch is not set, owning thread is awake
const UNSET: usize = 0;
/// Latch is not set, owning thread is going to sleep on this latch
/// (but has not yet fallen asleep).
const SLEEPY: usize = 1;
/// Latch is not set, owning thread is asleep on this latch and
/// must be awoken.
const SLEEPING: usize = 2;
/// Latch is set.
const SET: usize = 3;

pub(super) struct CoreLatch {
    state: AtomicUsize,
}

impl CoreLatch {
    pub(super) fn get_sleepy(&self) -> bool {
        self.state
            .compare_exchange(UNSET, SLEEPY, Ordering::SeqCst, Ordering::Relaxed)
            .is_ok()
    }

    pub(super) fn fall_asleep(&self) -> bool {
        self.state
            .compare_exchange(SLEEPY, SLEEPING, Ordering::SeqCst, Ordering::Relaxed)
            .is_ok()
    }

    pub(super) fn wake_up(&self) {
        if !self.probe() {
            let _ =
                self.state
                    .compare_exchange(SLEEPING, UNSET, Ordering::SeqCst, Ordering::Relaxed);
        }
    }

    unsafe fn set(this: *const Self) -> bool {
        let old_state = (*this).state.swap(SET, Ordering::AcqRel);
        old_state == SLEEPING
    }

    pub(super) fn probe(&self) -> bool {
        self.state.load(Ordering::Acquire) == SET
    }
}

pub(super) struct SpinLatch<'r> {
    core_latch: CoreLatch,
    registry: &'r Arc<Registry>,
    target_worker_index: usize,
    cross: bool,
}

impl Latch for SpinLatch<'_> {
    #[inline]
    unsafe fn set(this: *const Self) {
        let registry: &Registry = if (*this).cross {
            // Ensure the registry stays alive while we notify it.
            // Otherwise, it would be possible that we set the spin
            // latch and the other thread sees it and exits, causing
            // the registry to be deallocated, all before we get a
            // chance to invoke `registry.notify_worker_latch_is_set`.
            &Arc::clone((*this).registry)
        } else {
            // If this is not a "cross-registry" spin-latch, then the
            // thread which is performing `set` is itself ensuring
            // that the registry stays alive. However, that doesn't
            // include this *particular* `Arc` handle if the waiting
            // thread then exits, so we must completely dereference it.
            (*this).registry
        };
        let target_worker_index = (*this).target_worker_index;

        // NOTE: Once we `set`, the target may proceed and invalidate `this`!
        if CoreLatch::set(&(*this).core_latch) {
            // Subtle: at this point, we can no longer read from
            // `self`, because the thread owning this spin latch may
            // have awoken and deallocated the latch. Therefore, we
            // only use fields whose values we already read.
            registry.notify_worker_latch_is_set(target_worker_index);
        }
    }
}

pub(super) struct LockLatch {
    m: Mutex<bool>,
    v: Condvar,
}

impl LockLatch {
    pub(super) const fn new() -> LockLatch { /* ... */ }

    /// Block until latch is set, then resets this lock latch so it can be reused again.
    pub(super) fn wait_and_reset(&self) {
        let mut guard = self.m.lock().unwrap();
        while !*guard {
            guard = self.v.wait(guard).unwrap();
        }
        *guard = false;
    }

    /// Block until latch is set.
    pub(super) fn wait(&self) { /* 同上，不 reset */ }
}

impl Latch for LockLatch {
    unsafe fn set(this: *const Self) {
        let mut guard = (*this).m.lock().unwrap();
        *guard = true;
        (*this).v.notify_all();
    }
}

pub(super) unsafe fn set_and_tickle_one(
    this: *const Self,
    registry: &Registry,
    target_worker_index: usize,
) {
    if CoreLatch::set(&(*this).core_latch) {
        registry.notify_worker_latch_is_set(target_worker_index);
    }
}

pub(super) struct CountLatch {
    counter: AtomicUsize,
    kind: CountLatchKind,
}

enum CountLatchKind {
    /// A latch for scopes created on a rayon thread which will participate in work
    /// stealing while it waits for completion. This thread is not necessarily part
    /// of the same registry as the scope itself!
    Stealing {
        latch: CoreLatch,
        registry: Arc<Registry>,
        worker_index: usize,
    },

    /// A latch for scopes created on a non-rayon thread which will block to wait.
    Blocking { latch: LockLatch },
}

impl CountLatch {
    pub(super) fn increment(&self) {
        let old_counter = self.counter.fetch_add(1, Ordering::Relaxed);
        debug_assert!(old_counter != 0);
    }
}

impl Latch for CountLatch {
    unsafe fn set(this: *const Self) {
        if (*this).counter.fetch_sub(1, Ordering::SeqCst) == 1 {
            // NOTE: Once we call `set` on the internal `latch`,
            // the target may proceed and invalidate `this`!
            match (*this).kind {
                CountLatchKind::Stealing {
                    ref latch,
                    ref registry,
                    worker_index,
                } => {
                    let registry = Arc::clone(registry);
                    if CoreLatch::set(latch) {
                        registry.notify_worker_latch_is_set(worker_index);
                    }
                }
                CountLatchKind::Blocking { ref latch } => LockLatch::set(latch),
            }
        }
    }
}

/// `&L` without any implication of `dereferenceable` for `Latch::set`
pub(super) struct LatchRef<'a, L> {
    inner: *const L,
    marker: PhantomData<&'a L>,
}

impl<L: Latch> Latch for LatchRef<'_, L> {
    unsafe fn set(this: *const Self) {
        L::set((*this).inner);
    }
}