nuclei 0.5.4

//! How long a shard spins before it parks: [`Spin`].

use core::cell::Cell;
use std::time::{Duration, Instant};

/// The longest a shard spins before it parks, in nanoseconds.
const SPIN_MAX: u32 = 16_000;
/// The spin a fresh shard starts with.
const SPIN_START: u32 = 4_000;
/// The spin a shard whose park was cut short starts again with.
const SPIN_RESTART: u32 = 1_000;
/// A park that ended this soon would have been caught by a spin.
const SHORT_PARK: Duration = Duration::from_micros(50);
/// The longest gap between two looks of a spin, in pauses: each look reads
/// lines the waker writes, so looking seldom leaves it to run at full speed.
const SPIN_GAP_MAX: u32 = 64;

/// What a spin came to.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Spun {
    /// Work was found.
    Hit,
    /// The budget ran out, every look having failed (each asked the poller).
    Missed,
    /// The shard expected no work soon: no look was made.
    Cold,
}

/// How long a shard spins before it parks, adapted to what spins see.
///
/// Work that arrives within a spin costs neither a park nor the waker's
/// notification, and IO that becomes ready within it (each look asks the
/// poller too, without waiting) costs no kernel wake of the thread. A spin
/// that found work doubles the budget; one that ran out halves it, so a
/// shard whose work comes only after long gaps stops spinning; a park that
/// ended soon after it began sets it going again. The budget is time, read
/// from the clock after each gap of pauses (never every pause), and is a
/// hint: it only sizes the next spin.
///
/// A shard spins only while work may come soon: its caller says so (finite
/// operations of its ring out), or the shard has seen work come from another
/// thread within a spin or soon after a park (`hot`) and nothing since has
/// said otherwise. Else it parks at once and the budget stays. A job of its
/// runtime's blocking pool is not work that comes soon: it runs for as long
/// as it blocks, its end wakes the parked shard with one notification, and a
/// park it cuts short does not make the shard hot.
pub(crate) struct Spin {
    budget: Cell<u32>,
    hot: Cell<bool>,
}

impl Spin {
    pub(super) fn new() -> Self {
        Self {
            budget: Cell::new(SPIN_START),
            hot: Cell::new(false),
        }
    }

    /// Spends the budget, so the next spin does not run.
    #[cfg(test)]
    pub(super) fn spend(&self) {
        self.budget.set(0);
        self.hot.set(false);
    }

    /// With `expecting` work, spins until `found` says it is there or the
    /// budget is spent; without, makes no look. `expecting` is asked only of a shard that
    /// is not hot, the one that needs it. Inlined into the run loop, whichever
    /// codegen unit holds it.
    #[inline(always)]
    pub(super) fn run(
        &self,
        expecting: impl FnOnce() -> bool,
        mut found: impl FnMut() -> bool,
    ) -> Spun {
        let budget = self.budget.get();
        if budget == 0 || !(self.hot.get() || expecting()) {
            return Spun::Cold;
        }
        // The first look comes before the clock and the first pause: work already there costs
        // neither, and the budget runs from a failed look.
        let mut hit = found();
        if !hit {
            let limit = Duration::from_nanos(u64::from(budget));
            let start = Instant::now();
            let mut gap = 1;
            loop {
                for _ in 0..gap {
                    core::hint::spin_loop();
                }
                gap = (gap * 2).min(SPIN_GAP_MAX);
                hit = found();
                if hit || start.elapsed() >= limit {
                    break;
                }
            }
        }
        if hit {
            self.budget.set((budget * 2).clamp(SPIN_RESTART, SPIN_MAX));
            self.hot.set(true);
            return Spun::Hit;
        }
        self.budget.set(budget / 2);
        self.hot.set(false);
        Spun::Missed
    }

    /// Whether work came soon after a spin or a park, so the next spin runs
    /// unasked.
    #[cfg(test)]
    pub(super) fn is_hot(&self) -> bool {
        self.hot.get()
    }

    /// The shard slept for `slept`: a short sleep means work came right
    /// after the spin gave up.
    pub(super) fn parked_for(&self, slept: Duration) {
        if slept < SHORT_PARK {
            self.hot.set(true);
            self.budget.set(self.budget.get().max(SPIN_RESTART));
        }
    }
}

#[cfg(all(test, not(loom)))]
mod tests {
    use std::time::Duration;

    use super::{SPIN_MAX, SPIN_RESTART, SPIN_START, Spin, Spun};

    #[test]
    fn a_spin_that_finds_work_doubles_the_budget_and_one_that_runs_out_halves_it() {
        let spin = Spin::new();
        assert_eq!(spin.budget.get(), SPIN_START);
        assert_eq!(spin.run(|| true, || true), Spun::Hit);
        assert_eq!(spin.budget.get(), SPIN_START * 2);
        assert_ne!(spin.run(|| true, || false), Spun::Hit);
        assert_eq!(spin.budget.get(), SPIN_START);
        for _ in 0..20 {
            spin.run(|| true, || true);
        }
        assert_eq!(spin.budget.get(), SPIN_MAX, "capped");
        for _ in 0..20 {
            spin.run(|| true, || false);
        }
        assert_eq!(spin.budget.get(), 0, "an idle shard stops spinning");
        assert_eq!(
            spin.run(|| true, || true),
            Spun::Cold,
            "a spent budget does not spin"
        );
    }

    #[test]
    fn a_shard_expecting_nothing_parks_at_once_and_keeps_its_budget() {
        let spin = Spin::new();
        assert!(!spin.hot.get());
        let mut looks = 0;
        assert_eq!(
            spin.run(
                || false,
                || {
                    looks += 1;
                    true
                }
            ),
            Spun::Cold
        );
        assert_eq!(looks, 0, "it never looked");
        assert_eq!(spin.budget.get(), SPIN_START);
        spin.parked_for(Duration::from_micros(3));
        assert!(spin.hot.get(), "a short park: work comes soon after one");
        assert_eq!(
            spin.run(|| false, || true),
            Spun::Hit,
            "so it spins unasked"
        );
        assert_ne!(spin.run(|| false, || false), Spun::Hit);
        assert!(!spin.hot.get(), "a spin that ran out ends it");
        assert_ne!(spin.run(|| false, || true), Spun::Hit);
    }

    #[test]
    fn a_spin_ends_by_time_not_by_looks() {
        let spin = Spin::new();
        let start = std::time::Instant::now();
        assert_ne!(spin.run(|| true, || false), Spun::Hit);
        let spent = start.elapsed();
        assert!(
            spent >= Duration::from_nanos(u64::from(SPIN_START)),
            "{spent:?}"
        );
        assert!(spent < Duration::from_millis(50), "{spent:?}");
    }

    #[test]
    fn a_short_park_starts_the_spin_again() {
        let spin = Spin::new();
        for _ in 0..20 {
            spin.run(|| true, || false);
        }
        spin.parked_for(Duration::from_millis(5));
        assert_eq!(
            spin.budget.get(),
            0,
            "a long sleep: spinning would not have helped"
        );
        spin.parked_for(Duration::from_micros(3));
        assert_eq!(spin.budget.get(), SPIN_RESTART);
    }
}