ktstr/workload/

affinity.rs

//! CPU affinity intent + resolution for worker tasks.
//!
//! [`AffinityIntent`] expresses the test author's request before
//! topology resolution; [`ResolvedAffinity`] is the post-resolution
//! shape the spawn pipeline consumes. [`resolve_affinity`] walks the
//! `ResolvedAffinity` to a concrete CPU set; [`set_thread_affinity`]
//! issues the `sched_setaffinity` syscall; [`sched_getcpu`] is the
//! inverse query (which CPU is the current task on right now).
//!
//! Re-exported from the parent module via `pub use affinity::*` so
//! `crate::workload::AffinityIntent` etc. stay valid.

use std::collections::BTreeSet;

use anyhow::{Context, Result};

/// Scenario-level affinity intent for a group of workers.
///
/// Resolved to a concrete [`ResolvedAffinity`] at runtime based on the
/// cgroup's effective cpuset and the VM's topology. When attached to
/// a [`WorkSpec`](crate::workload::WorkSpec), determines per-worker `sched_setaffinity` masks.
///
/// Resolution uses [`resolve_affinity_for_cgroup()`](crate::scenario::resolve_affinity_for_cgroup).
///
/// # Naming pattern (Intent vs Resolved)
///
/// [`AffinityIntent`] and [`ResolvedAffinity`] form a pre/post-resolution
/// pair. Variant names line up where the same shape exists on both
/// sides; payload differences encode the intent → concrete-CPU-set
/// distinction:
///
/// | [`AffinityIntent`]                       | [`ResolvedAffinity`]              |
/// |------------------------------------------|-----------------------------------|
/// | `Inherit` (no payload)                   | `None`                            |
/// | `Exact(BTreeSet<usize>)`                 | `Fixed(BTreeSet<usize>)`          |
/// | `RandomSubset { from, count }`           | `Random { from, count }`          |
/// | `SingleCpu` (no payload)                 | `SingleCpu(usize)`                |
/// | `LlcAligned` / `CrossCgroup`             | `Fixed(...)` (resolver expands)   |
/// | `SmtSiblingPair` (no payload)            | `Fixed({sibling_a, sibling_b})`   |
///
/// Constructor helpers: [`AffinityIntent::exact`] takes any
/// `IntoIterator<Item = usize>` for the `Exact` set;
/// [`AffinityIntent::random_subset`] takes the same iterator shape
/// for the `RandomSubset` pool plus a sample-count argument.
///
/// The `SingleCpu` pair specifically: [`AffinityIntent::SingleCpu`]
/// expresses "pin to one CPU; resolver picks which based on cgroup
/// state and worker index", and [`ResolvedAffinity::SingleCpu`]
/// records the concrete CPU id chosen. Reusing the variant name keeps
/// the pre/post mapping lexically obvious — payload presence
/// distinguishes intent from resolution without renaming the variant.
///
/// [`AffinityIntent::RandomSubset`] carries the resolved pool
/// (`from`) and sample size (`count`) — sampling itself is deferred
/// to spawn time so each worker gets an independent draw. The
/// scenario engine's `resolve_affinity_for_cgroup` materialises the
/// pool from cgroup cpuset / topology before constructing this
/// variant; spawn-time `resolve_affinity` samples per-worker.
/// Construct directly via [`AffinityIntent::random_subset`].
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AffinityIntent {
    /// No affinity constraint -- inherit from parent cgroup.
    #[default]
    Inherit,
    /// Pin each worker to a random subset of `from`, sampling `count`
    /// CPUs per worker. Sampling is deferred to spawn time so each
    /// worker gets an independent draw — mirrors
    /// [`ResolvedAffinity::Random`] semantics. Construct with the
    /// resolved pool already materialised; the scenario engine pre-
    /// resolves topology-aware "pick from cgroup state" intent
    /// before building this variant.
    RandomSubset { from: BTreeSet<usize>, count: usize },
    /// Pin to the CPUs in the worker's LLC.
    LlcAligned,
    /// Pin to all CPUs (crosses cgroup boundaries).
    CrossCgroup,
    /// Pin to a single CPU.
    SingleCpu,
    /// Pin to an exact set of CPUs.
    Exact(BTreeSet<usize>),
    /// Pin all workers in the group to the two SMT siblings of one
    /// physical core. Tests how the scheduler handles two
    /// compute-bound tasks placed on SMT siblings — both threads
    /// contend for the core's shared front-end / execution
    /// resources, exposing scheduler decisions about co-running
    /// vs. spreading compute load across cores.
    ///
    /// Designed for `WorkType::SmtSiblingSpin` and other
    /// `worker_group_size = 2` variants
    /// (`WorkType::FutexPingPong`, `WorkType::AsymmetricWaker`,
    /// `WorkType::SignalStorm`, etc.) where both workers in a
    /// group are intended to run on a sibling pair. The variant
    /// has no payload — the resolver picks an SMT-sibling pair
    /// from the cgroup's effective cpuset (or the full topology
    /// when no cpuset is active).
    ///
    /// Resolution is performed by the scenario engine's
    /// `resolve_affinity_for_cgroup` (topology-aware, not
    /// available at the bare `WorkloadHandle::spawn` gate). The
    /// resolver searches the cpuset for a physical core with at
    /// least two thread siblings present and resolves to
    /// [`ResolvedAffinity::Fixed`] containing those two CPU IDs.
    /// All workers in the group get pinned to that 2-CPU set;
    /// when `num_workers == 2` the kernel runs one worker on each
    /// sibling, which is the contention pattern this intent
    /// targets.
    ///
    /// Returns an error from the resolver — NOT a silent
    /// fallback — when no SMT-sibling pair is available
    /// (`threads_per_core == 1`, or the cpuset isolates each
    /// sibling onto a different CPU set). Callers must handle
    /// the error; running `WorkType::SmtSiblingSpin` without
    /// SMT siblings would produce a misleading result.
    SmtSiblingPair,
}

impl AffinityIntent {
    /// Construct an `Exact` affinity from any iterator of CPU indices.
    ///
    /// Accepts arrays, ranges, `Vec`, `BTreeSet`, or any `IntoIterator<Item = usize>`.
    pub fn exact(cpus: impl IntoIterator<Item = usize>) -> Self {
        AffinityIntent::Exact(cpus.into_iter().collect())
    }

    /// Construct a `RandomSubset` from a pool iterator and a sample
    /// size. Mirrors the [`Self::exact`] constructor's iterator
    /// flexibility — accepts arrays, `Vec`, `BTreeSet`, ranges, or
    /// any `IntoIterator<Item = usize>` for the pool.
    ///
    /// Sampling is deferred to spawn time; each worker gets an
    /// independent `count`-sized draw from `from`. `count > from.len()`
    /// is clamped to `from.len()` at sample time (topology fact, not
    /// caller error). `count == 0` and empty `from` are rejected at
    /// the spawn-time affinity gate with an actionable diagnostic —
    /// use [`AffinityIntent::Inherit`] for no affinity constraint.
    pub fn random_subset(from: impl IntoIterator<Item = usize>, count: usize) -> Self {
        AffinityIntent::RandomSubset {
            from: from.into_iter().collect(),
            count,
        }
    }
}

/// Resolved CPU affinity for a worker process.
///
/// Created from [`AffinityIntent`] at runtime based on topology and
/// cpuset assignments. Variant names track [`AffinityIntent`] where the
/// same shape exists pre/post-resolution; payload presence
/// distinguishes intent from concrete CPU id(s). See the
/// [`AffinityIntent`] type doc for the full pre/post mapping table.
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ResolvedAffinity {
    /// No affinity constraint.
    #[default]
    None,
    /// Pin to a specific set of CPUs.
    Fixed(BTreeSet<usize>),
    /// Pin to `count` randomly-chosen CPUs from `from`.
    ///
    /// - `count` must be `> 0`; zero is rejected at resolve time
    ///   (previously it coerced silently to 1 and masked caller bugs).
    /// - `count > from.len()` is clamped to `from.len()` — asking for
    ///   more CPUs than the pool contains is a topology fact, not a
    ///   caller error.
    /// - `from` empty with `count > 0` is a caller bug: `resolve_affinity`
    ///   bails (an unsatisfiable sample request would otherwise produce an
    ///   empty `sched_setaffinity` mask that the kernel rejects with
    ///   `EINVAL`). The resolution step that produces this variant — see
    ///   [`crate::scenario::resolve_affinity_for_cgroup`] — bails on
    ///   empty pools before construction; no silent fallback. Direct
    ///   constructor callers (e.g. test fixtures) must do the same: use
    ///   [`ResolvedAffinity::None`] for "no affinity constraint", never
    ///   `Random { from: empty, count: > 0 }`.
    Random { from: BTreeSet<usize>, count: usize },
    /// Pin to a single CPU.
    SingleCpu(usize),
}

impl ResolvedAffinity {
    /// Construct a [`ResolvedAffinity::Fixed`] from any iterator over
    /// CPU ids. Mirrors [`AffinityIntent::exact`].
    pub fn fixed(cpus: impl IntoIterator<Item = usize>) -> Self {
        ResolvedAffinity::Fixed(cpus.into_iter().collect())
    }

    /// Construct a [`ResolvedAffinity::Random`] from a pool iterator
    /// and a sample count. Mirrors [`AffinityIntent::random_subset`].
    pub fn random(from: impl IntoIterator<Item = usize>, count: usize) -> Self {
        ResolvedAffinity::Random {
            from: from.into_iter().collect(),
            count,
        }
    }

    /// Construct a [`ResolvedAffinity::SingleCpu`].
    pub const fn single_cpu(cpu: usize) -> Self {
        ResolvedAffinity::SingleCpu(cpu)
    }
}

/// Resolve a [`ResolvedAffinity`] into the concrete CPU set the
/// spawn pipeline writes into the worker's `sched_setaffinity` mask.
///
/// `Random` samples `count` CPUs from `from` per call (each worker
/// gets an independent draw at spawn time). Empty `from` is a caller
/// bug and bails — the upstream resolver
/// [`crate::scenario::resolve_affinity_for_cgroup`] itself bails
/// (rather than degrading to [`ResolvedAffinity::None`]) on every
/// path that would produce an empty pool — empty cpuset intersection
/// against `RandomSubset.from`, `count == 0`, or any other
/// unsatisfiable shape. Reaching this fn with an empty `Random.from`
/// therefore indicates a caller that bypassed the resolver. Invalid
/// input must fail loudly, never silently degrade to "no affinity
/// applied". `count == 0` likewise bails.
pub(crate) fn resolve_affinity(mode: &ResolvedAffinity) -> Result<Option<BTreeSet<usize>>> {
    match mode {
        ResolvedAffinity::None => Ok(None),
        ResolvedAffinity::Fixed(cpus) => Ok(Some(cpus.clone())),
        ResolvedAffinity::SingleCpu(cpu) => Ok(Some([*cpu].into_iter().collect())),
        ResolvedAffinity::Random { from, count } => {
            use rand::seq::IndexedRandom;
            if *count == 0 {
                anyhow::bail!(
                    "ResolvedAffinity::Random.count must be > 0; a zero count \
                     previously silently coerced to 1, masking caller bugs"
                );
            }
            if from.is_empty() {
                anyhow::bail!(
                    "ResolvedAffinity::Random.from is empty with count={count}; \
                     a worker cannot be pinned to an empty CPU pool. The \
                     resolution step that produced this Random must reject \
                     the empty set up-front (e.g. via the bail paths in \
                     `crate::scenario::resolve_affinity_for_cgroup`) — \
                     forwarding an unsatisfiable sample request would \
                     silently drop the affinity constraint",
                    count = count,
                );
            }
            let pool: Vec<usize> = from.iter().copied().collect();
            // Clamp count down to the pool size (user asked for more
            // CPUs than exist). Silent clamp is fine here: the pool
            // upper bound is a topology fact, not a caller bug.
            let count = (*count).min(pool.len());
            Ok(Some(
                pool.sample(&mut rand::rng(), count).copied().collect(),
            ))
        }
    }
}

/// Return the CPU the calling task is currently running on.
///
/// Falls back to `0` on syscall failure (rare; would mean
/// `getcpu(2)` is unavailable, which is not the case on any
/// supported kernel). Wraps [`nix::sched::sched_getcpu`].
pub(crate) fn sched_getcpu() -> usize {
    nix::sched::sched_getcpu().unwrap_or(0)
}

/// Set per-thread CPU affinity via `sched_setaffinity(2)`.
///
/// `pid` must be `> 0` — `pid <= 0` has broadcast semantics at the
/// syscall level (target the calling task or every task in a tgid
/// depending on layer) and is rejected up-front so no caller passes
/// an unchecked `0` through.
pub fn set_thread_affinity(pid: libc::pid_t, cpus: &BTreeSet<usize>) -> Result<()> {
    use nix::sched::{CpuSet, sched_setaffinity};
    use nix::unistd::{Pid, SysconfVar, sysconf};
    // See `set_sched_policy` for the rationale — pid <= 0 has
    // broadcast semantics at the syscall and must not be passed
    // through unchecked.
    if pid <= 0 {
        anyhow::bail!("sched_setaffinity: invalid pid {pid} (must be > 0)");
    }
    // Snapshot the host's online-CPU count and the cpu_set bitmap
    // width before the loop so the diagnostic on overflow carries
    // both numbers without re-syscalling per offending CPU. Render
    // sysconf failure as "unavailable" rather than 0 to disambiguate
    // a degenerate sysconf result from a legitimately zero count.
    let online_cpus_str: std::borrow::Cow<'static, str> =
        match sysconf(SysconfVar::_NPROCESSORS_ONLN).ok().flatten() {
            Some(n) => format!("{n}").into(),
            None => "unavailable".into(),
        };
    let cpuset_bitmap_width: usize = libc::CPU_SETSIZE as usize;
    let mut cpu_set = CpuSet::new();
    for &cpu in cpus {
        cpu_set.set(cpu).with_context(|| {
            format!(
                "CPU {cpu} out of range: cpu_set bitmap holds CPU IDs \
                 0..{cpuset_bitmap_width} (libc CPU_SETSIZE) and host \
                 reports {online_cpus_str} online CPUs (sysconf \
                 _SC_NPROCESSORS_ONLN). Either the cpuset spec was \
                 resolved against a stale topology or the bitmap cap \
                 needs raising on this build."
            )
        })?;
    }
    sched_setaffinity(Pid::from_raw(pid), &cpu_set)
        .with_context(|| format!("sched_setaffinity pid={pid}"))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeSet;

    #[test]
    fn resolve_affinity_none() {
        let r = resolve_affinity(&ResolvedAffinity::None).unwrap();
        assert!(r.is_none());
    }
    #[test]
    fn resolve_affinity_fixed() {
        let cpus: BTreeSet<usize> = [0, 1, 2].into_iter().collect();
        let r = resolve_affinity(&ResolvedAffinity::Fixed(cpus.clone())).unwrap();
        assert_eq!(r, Some(cpus));
    }
    #[test]
    fn resolve_affinity_single_cpu() {
        let r = resolve_affinity(&ResolvedAffinity::SingleCpu(5)).unwrap();
        assert_eq!(r, Some([5].into_iter().collect()));
    }
    /// `ResolvedAffinity` derives `Debug`; the `SingleCpu` variant
    /// must render its variant name and the embedded CPU id so
    /// failure-dump and tracing output are diagnosable. Pins the
    /// derive against accidental removal.
    #[test]
    fn resolved_affinity_single_cpu_debug_format() {
        let dbg = format!("{:?}", ResolvedAffinity::SingleCpu(7));
        assert!(
            dbg.contains("SingleCpu"),
            "Debug output must name the variant, got: {dbg}"
        );
        assert!(
            dbg.contains('7'),
            "Debug output must include the CPU id payload, got: {dbg}"
        );
    }
    #[test]
    fn resolve_affinity_random() {
        let from: BTreeSet<usize> = (0..8).collect();
        let r = resolve_affinity(&ResolvedAffinity::Random { from, count: 3 }).unwrap();
        let cpus = r.unwrap();
        assert_eq!(cpus.len(), 3);
        assert!(cpus.iter().all(|c| *c < 8));
    }
    #[test]
    fn resolve_affinity_random_clamps_count() {
        let from: BTreeSet<usize> = [0, 1].into_iter().collect();
        let r = resolve_affinity(&ResolvedAffinity::Random { from, count: 10 }).unwrap();
        assert_eq!(r.unwrap().len(), 2);
    }
    #[test]
    fn resolve_affinity_random_single_cpu_pool() {
        let from: BTreeSet<usize> = [7].into_iter().collect();
        let r = resolve_affinity(&ResolvedAffinity::Random { from, count: 1 }).unwrap();
        assert_eq!(r.unwrap(), [7].into_iter().collect());
    }
    #[test]
    fn affinity_mode_debug_shows_cpus() {
        let a = ResolvedAffinity::Fixed([0, 1, 7].into_iter().collect());
        let s = format!("{:?}", a);
        assert!(s.contains("0"), "must show CPU 0");
        assert!(s.contains("1"), "must show CPU 1");
        assert!(s.contains("7"), "must show CPU 7");
        // Different CPU sets produce different output.
        let b = ResolvedAffinity::Fixed([3, 4].into_iter().collect());
        let s2 = format!("{:?}", b);
        assert!(s2.contains("3"), "must show CPU 3");
        assert_ne!(
            s, s2,
            "different CPU sets must produce different debug output"
        );
    }
    #[test]
    fn affinity_mode_clone_preserves_cpus() {
        let cpus: BTreeSet<usize> = [2, 5, 7].into_iter().collect();
        let a = ResolvedAffinity::Random {
            from: cpus.clone(),
            count: 2,
        };
        let b = a.clone();
        match b {
            ResolvedAffinity::Random { from, count } => {
                assert_eq!(from, cpus, "cloned from set must match original");
                assert_eq!(count, 2, "cloned count must match original");
            }
            _ => panic!("clone must preserve variant"),
        }
    }
    // -- resolve_affinity edge cases --

    #[test]
    fn resolve_affinity_random_zero_count_rejected() {
        // Regression: count=0 previously coerced silently to 1, masking
        // caller bugs. Now it must return an Err.
        let from: BTreeSet<usize> = (0..4).collect();
        let err = resolve_affinity(&ResolvedAffinity::Random { from, count: 0 }).unwrap_err();
        let msg = format!("{err}");
        assert!(
            msg.contains("count") && msg.contains("> 0"),
            "error must name the field: {msg}"
        );
    }
    #[test]
    fn resolve_affinity_random_empty_pool_bails() {
        // Empty Random.from with count > 0 is unsatisfiable: a worker
        // cannot be pinned to an empty CPU pool, and the prior
        // silent-degrade-to-Ok(None) was a silent-drop bug. The bail
        // forces the bug to surface and points the caller at the
        // upstream resolver as the place where empty pools should be
        // rejected up-front.
        let from: BTreeSet<usize> = BTreeSet::new();
        let err = resolve_affinity(&ResolvedAffinity::Random { from, count: 1 }).unwrap_err();
        let msg = format!("{err}");
        assert!(
            msg.contains("empty") && msg.contains("count=1"),
            "diagnostic must name the empty pool and the count: got {msg}",
        );
        // Also pin the actionable suggestion — names the upstream
        // resolver as the place callers should reject the empty set
        // so a caller hitting this error learns where to plug the hole.
        assert!(
            msg.contains("resolve_affinity_for_cgroup"),
            "diagnostic must point to the upstream resolver so callers \
             learn where the empty pool should have been rejected: got {msg}",
        );
    }

    #[test]
    fn sched_getcpu_valid() {
        let cpu = sched_getcpu();
        let max = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);
        assert!(cpu < max, "cpu {cpu} >= max {max}");
    }

    #[test]
    fn set_thread_affinity_cpu_zero() {
        let pid: libc::pid_t = unsafe { libc::getpid() };
        let cpus: BTreeSet<usize> = [0].into_iter().collect();
        let result = set_thread_affinity(pid, &cpus);
        assert!(result.is_ok(), "pinning to CPU 0 should succeed");
    }

    /// GAP 7: pin that the three constructors produce the same
    /// values a direct variant construction yields. A regression
    /// where `fixed(iter)` started normalising the input would
    /// silently shift downstream semantics.
    #[test]
    fn resolved_affinity_constructors_match_direct_variants() {
        let from_ctor = ResolvedAffinity::fixed([0_usize, 1, 2]);
        let from_variant = ResolvedAffinity::Fixed([0_usize, 1, 2].into_iter().collect());
        assert_eq!(from_ctor, from_variant);

        let from_ctor = ResolvedAffinity::random([0_usize, 1, 2, 3], 2);
        let from_variant = ResolvedAffinity::Random {
            from: [0_usize, 1, 2, 3].into_iter().collect(),
            count: 2,
        };
        assert_eq!(from_ctor, from_variant);

        let from_ctor = ResolvedAffinity::single_cpu(5);
        let from_variant = ResolvedAffinity::SingleCpu(5);
        assert_eq!(from_ctor, from_variant);
    }

    // Compile-time pin: `ResolvedAffinity::single_cpu` is `pub const
    // fn`. A regression that drops `const` (e.g. switches to a body
    // requiring runtime allocation) would silently break the
    // const-context use case. The `const _` binding fails to
    // type-check if `single_cpu` is no longer const-evaluable.
    const _: ResolvedAffinity = ResolvedAffinity::single_cpu(7);

    /// GAP 8: pin that `ResolvedAffinity::default()` is `None`
    /// and that every variant roundtrips through serde unchanged,
    /// including empty-payload edge cases for `Fixed` and `Random`.
    /// Serde drift on any variant breaks failure-dump replay and
    /// captured-workload reproduction — the persisted JSON would
    /// deserialize into a different variant or lose payload data.
    /// A regression that landed `#[serde(skip_serializing_if =
    /// "BTreeSet::is_empty")]` on the inner field would silently
    /// drop the variant tag on the empty case; pinning empty
    /// payloads catches that.
    #[test]
    fn resolved_affinity_default_is_none_and_serde_roundtrip_per_variant() {
        let d: ResolvedAffinity = Default::default();
        assert_eq!(d, ResolvedAffinity::None);

        let variants = [
            ResolvedAffinity::None,
            ResolvedAffinity::Fixed([0_usize, 1, 5].into_iter().collect()),
            ResolvedAffinity::Fixed(BTreeSet::new()),
            ResolvedAffinity::Random {
                from: [0_usize, 1, 2, 3, 4].into_iter().collect(),
                count: 2,
            },
            ResolvedAffinity::Random {
                from: BTreeSet::new(),
                count: 0,
            },
            ResolvedAffinity::SingleCpu(7),
        ];
        for original in &variants {
            let bytes = serde_json::to_vec(original).expect("serialize");
            let restored: ResolvedAffinity = serde_json::from_slice(&bytes).expect("deserialize");
            assert_eq!(restored, *original, "roundtrip drift for {original:?}");
        }
    }
}