scx_mitosis/

main.rs

// Copyright (c) Meta Platforms, Inc. and affiliates.

// This software may be used and distributed according to the terms of the
// GNU General Public License version 2.

#[allow(clippy::unwrap_used)]
mod bpf_skel;
pub use bpf_skel::*;
pub mod bpf_intf;
mod cell_manager;
mod config;
mod stats;
mod topology;
mod undefok_flags;

use cell_manager::{CellManager, CpuAssignment, CpuRecipient};
use config::{ConfiguredCells, SubcellMatch};

use std::cmp::max;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fmt::Display;
use std::mem::MaybeUninit;
use std::os::fd::AsFd;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;

use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use clap::Parser;
use libbpf_rs::MapCore as _;
use libbpf_rs::OpenObject;
use libbpf_rs::ProgramInput;
use nix::sys::epoll::{Epoll, EpollCreateFlags, EpollEvent, EpollFlags, EpollTimeout};
use nix::sys::eventfd::EventFd;
use scx_stats::prelude::*;
use scx_utils::Cpumask;
use scx_utils::NR_CPUS_POSSIBLE;
use scx_utils::Topology;
use scx_utils::TopologyArgs;
use scx_utils::UserExitInfo;
use scx_utils::build_id;
use scx_utils::compat;
use scx_utils::init_libbpf_logging;
use scx_utils::libbpf_clap_opts::LibbpfOpts;
use scx_utils::scx_enums;
use scx_utils::scx_ops_attach;
use scx_utils::scx_ops_load;
use scx_utils::scx_ops_open;
use scx_utils::uei_exited;
use scx_utils::uei_report;
use tracing::{debug, info, trace, warn};
use tracing_subscriber::filter::EnvFilter;

use stats::CellMetrics;
use stats::Metrics;
use topology::MitosisTopology;

const SCHEDULER_NAME: &str = "scx_mitosis";
const MAX_CELLS: usize = bpf_intf::consts_MAX_CELLS as usize;
const MAX_SUBCELLS_PER_CELL: usize = bpf_intf::consts_MAX_SUBCELLS_PER_CELL as usize;
const NR_CSTATS: usize = bpf_intf::cell_stat_idx_NR_CSTATS as usize;
const SHARE_SIBLING_SUBCELL_CPUS: bool = true;
/// Fixed-point 1.0 of the per-task demand accumulated by BPF.
const DEMAND_ONE: f64 = (1u64 << bpf_intf::consts_DEMAND_SHIFT) as f64;
/// Epoll token for inotify events (cgroup creation/destruction)
const INOTIFY_TOKEN: u64 = 1;
/// Epoll token for stats request wakeups
const STATS_TOKEN: u64 = 2;

fn parse_ewma_factor(s: &str) -> Result<f64, String> {
    let v: f64 = s.parse().map_err(|e| format!("{e}"))?;
    if !(0.0..=1.0).contains(&v) {
        return Err(format!("value {v} not in range 0.0..=1.0"));
    }
    Ok(v)
}

/// scx_mitosis: A dynamic affinity scheduler
///
/// Cgroups are assigned to a dynamic number of Cells which are assigned to a
/// dynamic set of CPUs. The BPF part does simple vtime scheduling for each subcell.
///
/// Userspace makes the dynamic decisions of which Cells should be merged or
/// split and which CPUs they should be assigned to.
#[derive(Debug, Parser)]
#[command(after_long_help = r#"VIRTUAL LLC CONFIGURATION:
    --virt-llc requires --enable-llc-awareness to affect DSQ creation.
    MIN-MAX specifies the target range of physical cores per virtual LLC.
    For each physical LLC and core type, the smallest exact divisor in the
    range is preferred; otherwise, the size with the smallest remainder is
    used, favoring smaller sizes. Remaining cores join the last partition.

    Example: On a 32-core physical LLC, --virt-llc=8-8 creates four virtual
    LLCs, and therefore four LLC DSQs per cell."#)]
struct Opts {
    /// Deprecated, noop, use RUST_LOG or --log-level instead.
    #[clap(short = 'v', long, action = clap::ArgAction::Count)]
    verbose: u8,

    /// Specify the logging level. Accepts rust's envfilter syntax for modular
    /// logging: https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#example-syntax. Examples: ["info", "warn,tokio=info"]
    #[clap(long, default_value = "info")]
    log_level: String,

    /// Exit debug dump buffer length. 0 indicates default.
    #[clap(long, default_value = "0")]
    exit_dump_len: u32,

    /// Interval to report monitoring information
    #[clap(long, default_value = "1")]
    monitor_interval_s: u64,

    /// Launch the scheduler and monitor its stats in-process with the
    /// specified interval.
    #[clap(long, conflicts_with = "monitor")]
    monitor_sync: Option<f64>,

    /// Run in stats monitoring mode with the specified interval. Scheduler
    /// is not launched.
    #[clap(long, conflicts_with = "monitor_sync")]
    monitor: Option<f64>,

    /// Print scheduler version and exit.
    #[clap(short = 'V', long, action = clap::ArgAction::SetTrue)]
    version: bool,

    /// Optional run ID for tracking scheduler instances.
    #[clap(long)]
    run_id: Option<u64>,

    /// Ignore the listed long flags if present. Accepts a comma-separated list
    /// of names without the leading `--`, for example
    /// `--undefok=reconfiguration-interval-s,rebalance-cpus-interval-s`.
    #[clap(long, value_delimiter = ',')]
    undefok: Vec<String>,

    /// Enable workaround for exiting tasks with offline cgroups during scheduler load.
    /// This works around a kernel bug where tasks can be initialized with cgroups that
    /// were never initialized. Disable this once the kernel bug is fixed.
    #[clap(long, default_value = "true", action = clap::ArgAction::Set)]
    exiting_task_workaround: bool,

    /// Deprecated compatibility flag. Mitosis always uses tracepoint cgroup tracking.
    #[allow(dead_code)]
    #[clap(long, action = clap::ArgAction::SetTrue)]
    cpu_controller_disabled: bool,

    /// Reject tasks with multi-CPU pinning that doesn't cover the entire cell.
    /// By default, these tasks are allowed but may have degraded performance.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    reject_multicpu_pinning: bool,

    /// Enable LLC-awareness. This will populate the scheduler's LLC maps and cause it
    /// to use LLC-aware scheduling. Required for --virt-llc to affect DSQ creation.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    enable_llc_awareness: bool,

    /// Deprecated, noop. LLC-aware mode always scans sibling LLC DSQs.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    enable_work_stealing: bool,

    /// Parent cgroup path whose direct children become cells.
    /// Scheduler startup requires this unless running in --monitor or --version mode.
    /// Example: --cell-parent-cgroup /workloads
    #[clap(
        long,
        required_unless_present_any = ["monitor", "version", "cell_config"],
        conflicts_with = "cell_config"
    )]
    cell_parent_cgroup: Option<String>,

    /// JSON cell configuration file. Each matching cgroup becomes a cell and
    /// inherits the subcell layout from its matching cell spec.
    #[clap(long, value_name = "PATH")]
    cell_config: Option<PathBuf>,

    /// Interval for rescanning cgroups matched by --cell-config.
    #[clap(long, default_value_t = 10)]
    reconfiguration_interval_s: u64,

    /// Exact directory name of a direct child cgroup to exclude from cell creation
    /// (excluded cgroups remain in cell 0). Matched against the directory basename,
    /// not the full path. Can be specified multiple times. Requires --cell-parent-cgroup.
    /// Example: --cell-exclude systemd-workaround.service
    #[clap(long)]
    cell_exclude: Vec<String>,

    /// Reserve up to this many CPUs for cell 0 (the root/catch-all cell)
    /// before child cpusets are applied (the holdout). A non-zero value keeps
    /// cell 0 from being starved to zero when child cgroups' cpusets cover
    /// every CPU. The reservation never takes a child cell's last CPU, so cell 0
    /// may receive fewer than requested. 0 (the default) disables the holdout;
    /// the assignment then bails if cell 0 would receive no CPUs.
    #[clap(long, default_value_t = 0)]
    cell0_min_cpus: usize,

    /// Enable cross-cell CPU borrowing. Subcells always share idle CPUs with
    /// sibling subcells inside their own cell.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    enable_borrowing: bool,

    /// Use lockless scx_bpf_dsq_peek() instead of the default iterator-based peek.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    use_lockless_peek: bool,

    /// Enable demand-based CPU rebalancing between cells.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    enable_rebalancing: bool,

    /// Utilization spread (max - min) that triggers rebalancing (default: 20%)
    #[clap(long, default_value = "20.0")]
    rebalance_threshold: f64,

    /// Minimum seconds between rebalancing events (default: 5)
    #[clap(long, default_value = "5")]
    rebalance_cooldown_s: u64,

    /// EWMA smoothing factor for demand tracking. Higher = more responsive (default: 0.3)
    #[clap(long, default_value = "0.3", value_parser = parse_ewma_factor)]
    demand_smoothing: f64,

    /// Dynamically reassign multi-CPU affinitized tasks on each wake: prefer an
    /// idle CPU within the mask, fall back to random. Redistribute at enqueue if
    /// the target CPU already has queued work.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    dynamic_affinity_cpu_selection: bool,

    /// Enable slice shrinking for CPU-pinned tasks. Uses per-task EWMA
    /// runtime to shrink the running task's slice when pinned waiters are queued.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    enable_slice_shrinking: bool,

    /// Upper bound for shrink limit (us). Used when the proportional
    /// value (avg_runtime * K) exceeds it.
    #[clap(long, default_value = "4000")]
    slice_shrink_max_us: u64,

    /// Minimum shrink limit (us). Slices are never shrunk below this value.
    /// In practice, the resolution here is determined by the kernel's
    /// tick period.
    #[clap(long, default_value = "500")]
    slice_shrink_min_us: u64,

    #[clap(flatten, next_help_heading = "Topology Options")]
    topology: TopologyArgs,

    #[clap(flatten, next_help_heading = "Libbpf Options")]
    pub libbpf: LibbpfOpts,
}

// The subset of cstats we care about.
// Local + Default + Hi + Lo = Total Decisions
// Affinity violations are not queue decisions, but
// will be calculated separately and reported as a percent of the total
const QUEUE_STATS_IDX: [bpf_intf::cell_stat_idx; 4] = [
    bpf_intf::cell_stat_idx_CSTAT_LOCAL,
    bpf_intf::cell_stat_idx_CSTAT_CPU_DSQ,
    bpf_intf::cell_stat_idx_CSTAT_CELL_DSQ,
    bpf_intf::cell_stat_idx_CSTAT_BORROWED,
];

// Per cell book-keeping
#[derive(Debug)]
struct Cell {
    cpus: Cpumask,
    subcells: Vec<Subcell>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct Subcell {
    id: u32,
    name: String,
    primary: Cpumask,
    borrowable: Option<Cpumask>,
    matches: Vec<Vec<SubcellMatch>>,
}

impl Cell {
    fn new() -> Self {
        Self {
            cpus: Cpumask::new(),
            subcells: vec![Subcell::new(0)],
        }
    }
}

impl Subcell {
    fn new(id: u32) -> Self {
        Self {
            id,
            name: default_subcell_name(id),
            primary: Cpumask::new(),
            borrowable: None,
            matches: Vec::new(),
        }
    }
}

/// Where the scheduler's cells come from.
enum CellSource {
    /// Direct children of the cgroup passed via --cell-parent-cgroup,
    /// tracked through inotify.
    Parent(CellManager),
    /// Cgroups matched by the --cell-config rules, rescanned periodically.
    Configured(ConfiguredCells),
}

struct Scheduler<'a> {
    skel: BpfSkel<'a>,
    monitor_interval: Duration,
    cells: HashMap<u32, Cell>,
    // These are the per-cell cstats.
    // Note these are accumulated across all CPUs.
    prev_cell_stats: [[u64; NR_CSTATS]; MAX_CELLS],
    // Per-cell running_ns tracking for demand metrics
    prev_cell_running_ns: [u64; MAX_CELLS],
    prev_cell_own_ns: [u64; MAX_CELLS],
    prev_cell_lent_ns: [u64; MAX_CELLS],
    /// When the demand metrics were last sampled, the base of the next deltas
    demand_sampled_at: Option<Instant>,
    prev_subcell_running_ns: [[u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    prev_subcell_own_ns: [[u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    prev_subcell_lent_ns: [[u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    prev_subcell_queued_ns: [[u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    prev_subcell_demand_sum: [[u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    metrics: Metrics,
    stats_server: Option<StatsServer<(), Metrics>>,
    last_configuration_seq: Option<u32>,
    /// Last observed cpuset_seq for cpuset change detection
    last_cpuset_seq: u32,
    /// Where the scheduler's cells come from.
    cell_source: CellSource,
    /// CPU-to-LLC mapping used to steer configured-cell holdout selection.
    cpu_to_llc: HashMap<usize, usize>,
    /// Best-effort number of CPUs reserved for cell 0.
    cell0_min_cpus: usize,
    /// Sticky configured-mode indication that the holdout displaced a claim.
    configured_enforced_holdout: bool,
    /// Whether CPU borrowing is enabled
    enable_borrowing: bool,
    /// Whether demand-based rebalancing is enabled
    enable_rebalancing: bool,
    /// Utilization spread threshold for triggering rebalancing
    rebalance_threshold: f64,
    /// Minimum duration between rebalancing events
    rebalance_cooldown: Duration,
    /// EWMA smoothing factor for demand tracking
    demand_smoothing: f64,
    /// EWMA-smoothed utilization per cell
    smoothed_util: [f64; MAX_CELLS],
    /// EWMA-smoothed utilization per subcell
    smoothed_subcell_util: [[f64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    /// EWMA-smoothed load (average runnable tasks) per subcell
    smoothed_subcell_load: [[f64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    /// EWMA-smoothed demand (CPUs) per subcell
    smoothed_subcell_demand: [[f64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
    /// Last time rebalancing was performed
    last_rebalance: Instant,
    /// Number of rebalancing events
    rebalance_count: u64,
    /// Epoll instance for waiting on multiple fds (inotify, stats wakeup)
    epoll: Epoll,
    /// EventFd to wake up main loop when stats are requested
    stats_waker: EventFd,
    /// Interval for rescanning configured cgroups
    reconfiguration_interval: Duration,
    /// Last time configured cgroups were rescanned
    last_config_refresh: Instant,
}

struct DistributionStats {
    total_decisions: u64,
    share_of_decisions_pct: f64,
    local_q_pct: f64,
    cpu_q_pct: f64,
    cell_q_pct: f64,
    borrowed_pct: f64,
    affn_viol_pct: f64,
    steal_pct: f64,
    pin_skip_pct: f64,

    // for formatting
    global_queue_decisions: u64,
}

impl Display for DistributionStats {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // This makes the output easier to read by improving column alignment. First, it guarantees that within a
        // given logging interval, the global and cell queueing decision counts print at the same width.
        // Second, it reduces variance in column width between logging intervals. 5 is simply a heuristic.
        const MIN_DECISIONS_WIDTH: usize = 5;
        let descisions_width = if self.global_queue_decisions > 0 {
            max(
                MIN_DECISIONS_WIDTH,
                (self.global_queue_decisions as f64).log10().ceil() as usize,
            )
        } else {
            MIN_DECISIONS_WIDTH
        };
        write!(
            f,
            "{:width$} {:5.1}% | Local:{:4.1}% From: CPU:{:4.1}% Cell:{:4.1}% Borrow:{:4.1}% | V:{:4.1}% S:{:4.1}% PS:{:4.1}%",
            self.total_decisions,
            self.share_of_decisions_pct,
            self.local_q_pct,
            self.cpu_q_pct,
            self.cell_q_pct,
            self.borrowed_pct,
            self.affn_viol_pct,
            self.steal_pct,
            self.pin_skip_pct,
            width = descisions_width,
        )
    }
}

impl<'a> Scheduler<'a> {
    fn validate_args(opts: &Opts) -> Result<()> {
        if opts.cell_parent_cgroup.is_some() && opts.cell_config.is_some() {
            bail!("--cell-parent-cgroup and --cell-config are mutually exclusive");
        }
        if opts.cell_parent_cgroup.is_none() && opts.cell_config.is_none() {
            bail!("either --cell-parent-cgroup or --cell-config is required");
        }
        if !opts.cell_exclude.is_empty() && opts.cell_parent_cgroup.is_none() {
            bail!("--cell-exclude requires --cell-parent-cgroup");
        }
        Ok(())
    }

    fn init(opts: &Opts, open_object: &'a mut MaybeUninit<OpenObject>) -> Result<Self> {
        Self::validate_args(opts).context("validating scheduler options")?;

        let topology = Topology::with_args(&opts.topology).context("detecting system topology")?;

        let nr_llc = topology.all_llcs.len().max(1);

        let mut skel_builder = BpfSkelBuilder::default();
        skel_builder
            .obj_builder
            .debug(opts.log_level.contains("trace"));
        init_libbpf_logging(None);
        info!(
            "Running scx_mitosis (build ID: {})",
            build_id::full_version(env!("CARGO_PKG_VERSION"))
        );

        let open_opts = opts.libbpf.clone().into_bpf_open_opts();
        let mut skel = scx_ops_open!(skel_builder, open_object, mitosis, open_opts)
            .context("opening BPF skeleton")?;

        skel.struct_ops.mitosis_mut().exit_dump_len = opts.exit_dump_len;

        let rodata = skel
            .maps
            .rodata_data
            .as_mut()
            .expect("BUG: rodata_data missing after skel open");

        rodata.slice_ns = scx_enums.SCX_SLICE_DFL;
        rodata.exiting_task_workaround_enabled = opts.exiting_task_workaround;
        rodata.dynamic_affinity_cpu_selection = opts.dynamic_affinity_cpu_selection;

        // Slice shrinking configuration
        if opts.slice_shrink_min_us >= opts.slice_shrink_max_us {
            bail!(
                "--slice-shrink-min-us ({}) must be less than --slice-shrink-max-us ({})",
                opts.slice_shrink_min_us,
                opts.slice_shrink_max_us
            );
        }
        rodata.enable_slice_shrinking = opts.enable_slice_shrinking;
        rodata.slice_shrink_max_ns = opts.slice_shrink_max_us * 1_000;
        // K=2: in the proportional region, a pinned task waits at most 2x its historical runtime
        rodata.slice_shrink_multiplier = 2;
        rodata.slice_shrink_min_ns = opts.slice_shrink_min_us * 1_000;

        rodata.nr_possible_cpus = *NR_CPUS_POSSIBLE as u32;
        for cpu in topology.all_cpus.keys() {
            rodata.all_cpus[cpu / 8] |= 1 << (cpu % 8);
        }

        rodata.reject_multicpu_pinning = opts.reject_multicpu_pinning;

        // Set nr_llc in rodata
        rodata.nr_llc = nr_llc as u32;
        rodata.enable_llc_awareness = opts.enable_llc_awareness;

        rodata.enable_borrowing = opts.enable_borrowing;
        rodata.use_lockless_peek = opts.use_lockless_peek;

        match *compat::SCX_OPS_ALLOW_QUEUED_WAKEUP {
            0 => info!("Kernel does not support queued wakeup optimization."),
            v => skel.struct_ops.mitosis_mut().flags |= v,
        }

        let mitosis_topology = MitosisTopology::new(&topology);
        mitosis_topology
            .apply(&mut skel)
            .context("mitosis_topology.apply")?;

        let skel = scx_ops_load!(skel, mitosis, uei).context("loading BPF skeleton")?;

        let stats_server = StatsServer::new(stats::server_data())
            .launch()
            .context("launching stats server")?;

        let cpu_to_llc: HashMap<usize, usize> = mitosis_topology.cpu_to_llc.into_iter().collect();
        let cell_source = if let Some(ref parent_cgroup) = opts.cell_parent_cgroup {
            let exclude: HashSet<String> = opts.cell_exclude.iter().cloned().collect();
            CellSource::Parent(
                CellManager::new(
                    parent_cgroup,
                    MAX_CELLS as u32,
                    topology.span.clone(),
                    exclude,
                    opts.cell0_min_cpus,
                    cpu_to_llc.clone(),
                )
                .with_context(|| {
                    format!("initializing cell manager for cgroup {}", parent_cgroup)
                })?,
            )
        } else if let Some(ref cell_config) = opts.cell_config {
            CellSource::Configured(
                ConfiguredCells::load(cell_config, MAX_CELLS as u32, topology.span.clone())
                    .with_context(|| {
                        format!(
                            "initializing configured cells from {}",
                            cell_config.display()
                        )
                    })?,
            )
        } else {
            bail!("either --cell-parent-cgroup or --cell-config is required");
        };

        // Create epoll instance for event-driven main loop
        let epoll = Epoll::new(EpollCreateFlags::empty()).context("creating epoll instance")?;

        // Create eventfd for stats wakeup (non-blocking, semaphore mode)
        let stats_waker = EventFd::from_value_and_flags(
            0,
            nix::sys::eventfd::EfdFlags::EFD_NONBLOCK | nix::sys::eventfd::EfdFlags::EFD_SEMAPHORE,
        )
        .context("creating stats-waker eventfd")?;

        // Register stats_waker with epoll
        epoll
            .add(
                &stats_waker,
                EpollEvent::new(EpollFlags::EPOLLIN, STATS_TOKEN),
            )
            .context("registering stats-waker with epoll")?;

        match &cell_source {
            CellSource::Parent(cell_manager) => {
                epoll
                    .add(
                        cell_manager,
                        EpollEvent::new(EpollFlags::EPOLLIN, INOTIFY_TOKEN),
                    )
                    .context("registering cell manager inotify with epoll")?;
            }
            CellSource::Configured(_) => {}
        }

        Ok(Self {
            skel,
            monitor_interval: Duration::from_secs(opts.monitor_interval_s),
            cells: HashMap::new(),
            prev_cell_stats: [[0; NR_CSTATS]; MAX_CELLS],
            prev_cell_running_ns: [0; MAX_CELLS],
            prev_cell_own_ns: [0; MAX_CELLS],
            prev_cell_lent_ns: [0; MAX_CELLS],
            demand_sampled_at: None,
            prev_subcell_running_ns: [[0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            prev_subcell_own_ns: [[0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            prev_subcell_lent_ns: [[0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            prev_subcell_queued_ns: [[0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            prev_subcell_demand_sum: [[0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            metrics: Metrics::default(),
            stats_server: Some(stats_server),
            last_configuration_seq: None,
            last_cpuset_seq: 0,
            cell_source,
            cpu_to_llc,
            cell0_min_cpus: opts.cell0_min_cpus,
            configured_enforced_holdout: false,
            enable_borrowing: opts.enable_borrowing,
            enable_rebalancing: opts.enable_rebalancing,
            rebalance_threshold: opts.rebalance_threshold,
            rebalance_cooldown: Duration::from_secs(opts.rebalance_cooldown_s),
            demand_smoothing: opts.demand_smoothing,
            smoothed_util: [0.0; MAX_CELLS],
            smoothed_subcell_util: [[0.0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            smoothed_subcell_load: [[0.0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            smoothed_subcell_demand: [[0.0; MAX_SUBCELLS_PER_CELL]; MAX_CELLS],
            last_rebalance: Instant::now(),
            rebalance_count: 0,
            epoll,
            stats_waker,
            reconfiguration_interval: Duration::from_secs(opts.reconfiguration_interval_s),
            last_config_refresh: Instant::now(),
        })
    }

    fn run(&mut self, shutdown: Arc<AtomicBool>) -> Result<UserExitInfo> {
        let struct_ops = scx_ops_attach!(self.skel, mitosis).context("attaching BPF scheduler")?;

        info!("Mitosis Scheduler Attached. Run `scx_mitosis --monitor` for metrics.");

        // Apply initial cell configuration if CellManager is active
        self.apply_initial_cells()
            .context("applying initial cell configuration")?;

        let (res_ch, req_ch) = self
            .stats_server
            .as_ref()
            .expect("BUG: stats_server missing after init")
            .channels();

        // Spawn thread to bridge stats requests to eventfd.
        // The thread exits when the channel closes (stats_server dropped).
        // Clone the eventfd so the thread owns its own handle to the same kernel object.
        let stats_waker_fd = self
            .stats_waker
            .as_fd()
            .try_clone_to_owned()
            .context("cloning stats-waker fd for bridge thread")?;
        let stats_waker = unsafe { EventFd::from_owned_fd(stats_waker_fd) };
        let stats_bridge = std::thread::spawn(move || {
            while req_ch.recv().is_ok() {
                // Wake up main loop via eventfd
                let _ = stats_waker.write(1);
            }
        });

        while !shutdown.load(Ordering::Relaxed) && !uei_exited!(&self.skel, uei) {
            let mut events = [EpollEvent::empty(); 1];
            let timeout = EpollTimeout::try_from(self.monitor_interval).with_context(|| {
                format!(
                    "monitor_interval {:?} exceeds maximum epoll timeout",
                    self.monitor_interval,
                )
            })?;

            match self.epoll.wait(&mut events, timeout) {
                Ok(n) => {
                    for event in &events[..n] {
                        match event.data() {
                            INOTIFY_TOKEN => {
                                // Cgroup event - process immediately
                                self.process_cell_events()
                                    .context("processing cell manager events")?;
                            }
                            STATS_TOKEN => {
                                // Stats request - drain eventfd and send metrics
                                let _ = self.stats_waker.read();
                                res_ch
                                    .send(self.get_metrics())
                                    .context("sending metrics response")?;
                            }
                            _ => {}
                        }
                    }
                }
                Err(nix::errno::Errno::EINTR) => continue,
                Err(e) => return Err(e.into()),
            }

            // Periodic work on every iteration
            self.refresh_bpf_cells()
                .context("refreshing BPF cell state")?;
            self.check_cpuset_changes()
                .context("checking cpuset changes")?;
            self.maybe_refresh_configured_cells()
                .context("refreshing configured cells")?;
            self.collect_metrics().context("collecting metrics")?;

            if self.enable_rebalancing {
                self.maybe_rebalance().context("running rebalance check")?;
            }
        }

        drop(struct_ops);
        // Drop stats_server to close the channel, allowing stats_bridge to exit
        drop(self.stats_server.take());
        let _ = stats_bridge.join();
        info!("Unregister {SCHEDULER_NAME} scheduler");
        uei_report!(&self.skel, uei)
    }

    /// Cell manager backing parent-cgroup mode.
    fn cell_manager(&self) -> Result<&CellManager> {
        match &self.cell_source {
            CellSource::Parent(cell_manager) => Ok(cell_manager),
            CellSource::Configured(_) => {
                bail!("BUG: cell manager requested in --cell-config mode")
            }
        }
    }

    fn cell_manager_mut(&mut self) -> Result<&mut CellManager> {
        match &mut self.cell_source {
            CellSource::Parent(cell_manager) => Ok(cell_manager),
            CellSource::Configured(_) => {
                bail!("BUG: cell manager requested in --cell-config mode")
            }
        }
    }

    /// Configured cells backing --cell-config mode.
    fn configured_cells(&self) -> Result<&ConfiguredCells> {
        match &self.cell_source {
            CellSource::Parent(_) => {
                bail!("BUG: configured cells requested in --cell-parent-cgroup mode")
            }
            CellSource::Configured(configured_cells) => Ok(configured_cells),
        }
    }

    fn configured_cells_mut(&mut self) -> Result<&mut ConfiguredCells> {
        match &mut self.cell_source {
            CellSource::Parent(_) => {
                bail!("BUG: configured cells requested in --cell-parent-cgroup mode")
            }
            CellSource::Configured(configured_cells) => Ok(configured_cells),
        }
    }

    /// Apply initial cell assignments discovered at startup
    fn apply_initial_cells(&mut self) -> Result<()> {
        match self.cell_source {
            CellSource::Parent(_) => self.apply_initial_parent_cells(),
            CellSource::Configured(_) => self.apply_initial_configured_cells(),
        }
    }

    fn apply_initial_parent_cells(&mut self) -> Result<()> {
        let cpu_assignments = self
            .compute_and_apply_cell_config(&[])
            .context("computing initial cell configuration")?;

        info!(
            "Applied initial cell configuration: {}",
            self.cell_manager()?.format_cell_config(&cpu_assignments)
        );

        Ok(())
    }

    fn apply_initial_configured_cells(&mut self) -> Result<()> {
        let cpu_assignments = self
            .compute_and_apply_configured_cell_config()
            .context("computing initial configured cell configuration")?;
        self.last_config_refresh = Instant::now();

        info!(
            "Applied initial configured cell configuration: {}",
            self.configured_cells()?
                .format_cell_config(&cpu_assignments)
        );

        Ok(())
    }

    /// Process cell manager events (new/destroyed cgroups)
    fn process_cell_events(&mut self) -> Result<()> {
        let (num_new, num_destroyed, new_cell_ids, destroyed_cell_ids) = {
            let (new_cells, destroyed_cells) = self
                .cell_manager_mut()?
                .process_events()
                .context("processing inotify events")?;

            if new_cells.is_empty() && destroyed_cells.is_empty() {
                return Ok(());
            }

            let new_ids: Vec<u32> = new_cells.iter().map(|(_, cell_id)| *cell_id).collect();
            (
                new_cells.len(),
                destroyed_cells.len(),
                new_ids,
                destroyed_cells,
            )
        };

        // Clear smoothed_util for destroyed cells so stale data doesn't
        // leak if the cell ID is reused later.
        for &cell_id in &destroyed_cell_ids {
            self.smoothed_util[cell_id as usize] = 0.0;
            self.smoothed_subcell_util[cell_id as usize] = [0.0; MAX_SUBCELLS_PER_CELL];
            self.smoothed_subcell_load[cell_id as usize] = [0.0; MAX_SUBCELLS_PER_CELL];
            self.smoothed_subcell_demand[cell_id as usize] = [0.0; MAX_SUBCELLS_PER_CELL];
        }

        let cpu_assignments = self
            .compute_and_apply_cell_config(&new_cell_ids)
            .context("recomputing cell configuration for new cgroups")?;

        info!(
            "Cell config updated ({} new, {} destroyed): {}",
            num_new,
            num_destroyed,
            self.cell_manager()?.format_cell_config(&cpu_assignments)
        );

        Ok(())
    }

    /// Compute cell configuration from CellManager and apply it to BPF.
    ///
    /// When rebalancing is enabled and there is existing utilization data,
    /// uses demand-weighted CPU assignment instead of equal-weight. New cells
    /// (listed in `new_cell_ids`) are seeded to the average smoothed_util of
    /// existing cells so they start with a fair share rather than zero.
    ///
    /// Returns the CPU assignments for use with `format_cell_config`.
    fn compute_and_apply_cell_config(
        &mut self,
        new_cell_ids: &[u32],
    ) -> Result<Vec<CpuAssignment>> {
        let (cell_assignments, cpu_assignments) = {
            let active_cell_ids: Vec<u32> = self
                .cell_manager()?
                .get_cell_assignments()
                .iter()
                .map(|(_, cell_id)| *cell_id)
                .collect();
            // Cell 0 is always active
            let all_cell_ids: Vec<u32> = std::iter::once(0)
                .chain(active_cell_ids.iter().copied())
                .collect();

            let cpu_assignments = if self.enable_rebalancing {
                // Check if any existing (non-new) cell has utilization data
                let new_set: HashSet<u32> = new_cell_ids.iter().copied().collect();
                let existing_utils: Vec<f64> = all_cell_ids
                    .iter()
                    .filter(|id| !new_set.contains(id))
                    .map(|&id| self.smoothed_util[id as usize])
                    .collect();

                let has_data = existing_utils.iter().any(|&u| u > 0.0);

                if has_data {
                    // Seed new cells to the average utilization of existing cells
                    let avg_util: f64 =
                        existing_utils.iter().sum::<f64>() / existing_utils.len().max(1) as f64;
                    for &id in new_cell_ids {
                        self.smoothed_util[id as usize] = avg_util;
                        info!(
                            "Seeded new cell {} smoothed_util to average {:.1}%",
                            id, avg_util
                        );
                    }

                    // Build demand map from smoothed_util for all active cells
                    let cell_demands: HashMap<u32, f64> = all_cell_ids
                        .iter()
                        .map(|&id| (id, self.smoothed_util[id as usize]))
                        .collect();

                    self.cell_manager()?
                        .compute_demand_cpu_assignments(&cell_demands, self.enable_borrowing)
                        .context("computing demand-weighted CPU assignments")?
                } else {
                    // No utilization data yet (e.g., initial startup) — equal weight
                    self.cell_manager()?
                        .compute_cpu_assignments(self.enable_borrowing)
                        .context("computing equal-weight CPU assignments (no utilization data)")?
                }
            } else {
                self.cell_manager()?
                    .compute_cpu_assignments(self.enable_borrowing)
                    .context("computing equal-weight CPU assignments (rebalancing disabled)")?
            };

            (self.cell_manager()?.get_cell_assignments(), cpu_assignments)
        };

        let subcell_assignments = self.compute_subcell_assignments(&cpu_assignments)?;

        self.apply_cell_config(&cell_assignments, &cpu_assignments, &subcell_assignments)
            .context("applying cell configuration to BPF")?;

        Ok(cpu_assignments)
    }

    /// Compute cell configuration from --cell-config and apply it to BPF.
    fn compute_and_apply_configured_cell_config(&mut self) -> Result<Vec<CpuAssignment>> {
        let cell_demands = if self.enable_rebalancing {
            Some(self.active_cell_demands())
        } else {
            None
        };
        let cell0_min_cpus = self.cell0_min_cpus;

        let (resolution, all_cpus) = {
            let configured_cells = self.configured_cells_mut()?;
            let resolution = configured_cells
                .resolve(cell_demands.as_ref(), cell0_min_cpus)
                .context("resolving configured cells")?;
            (resolution, configured_cells.all_cpus().clone())
        };

        for configured_cell in &resolution.cells {
            let cell = self
                .cells
                .entry(configured_cell.id)
                .or_insert_with(Cell::new);
            // Keep the CPU masks BPF last reported for subcells that persist
            // so the cell never appears to have zero capacity between this
            // apply and the next read-back. New subcells start empty.
            let mut existing: HashMap<u32, Subcell> = std::mem::take(&mut cell.subcells)
                .into_iter()
                .map(|subcell| (subcell.id, subcell))
                .collect();
            cell.subcells = configured_cell
                .subcells
                .iter()
                .map(|subcell| {
                    let (primary, borrowable) = existing
                        .remove(&subcell.id)
                        .map(|old| (old.primary, old.borrowable))
                        .unwrap_or((Cpumask::new(), None));
                    Subcell {
                        id: subcell.id,
                        name: subcell.name.clone(),
                        primary,
                        borrowable,
                        matches: subcell.matches.clone(),
                    }
                })
                .collect();
        }

        let (cpu_assignments, displaced_claim) =
            CellManager::compute_configured_cell_cpu_assignments(
                &all_cpus,
                &resolution.cell_recipients,
                self.enable_borrowing,
                &self.cpu_to_llc,
            )
            .context("computing configured cell CPU assignments")?;
        self.configured_enforced_holdout |= displaced_claim;
        let subcell_assignments = self.compute_subcell_assignments(&cpu_assignments)?;

        let cell_assignments = resolution.cell_assignments;

        self.apply_cell_config(&cell_assignments, &cpu_assignments, &subcell_assignments)
            .context("applying configured cell configuration to BPF")?;

        Ok(cpu_assignments)
    }

    fn active_cell_demands(&self) -> HashMap<u32, f64> {
        self.cells
            .keys()
            .map(|&id| (id, self.smoothed_util[id as usize]))
            .collect()
    }

    /// Check if rebalancing should be triggered and apply demand-weighted CPU assignments.
    fn maybe_rebalance(&mut self) -> Result<()> {
        // Check cooldown
        if self.last_rebalance.elapsed() < self.rebalance_cooldown {
            return Ok(());
        }

        // Compute min/max smoothed utilization across active cells
        let active_cells: Vec<u32> = self.cells.keys().copied().collect();
        if active_cells.len() < 2 {
            return Ok(());
        }

        let mut min_util = f64::MAX;
        let mut max_util = f64::MIN;
        for &cell_id in &active_cells {
            let util = self.smoothed_util[cell_id as usize];
            if util < min_util {
                min_util = util;
            }
            if util > max_util {
                max_util = util;
            }
        }

        let spread = max_util - min_util;
        if spread < self.rebalance_threshold {
            return Ok(());
        }

        match self.cell_source {
            CellSource::Parent(_) => self.rebalance_parent_cells(spread, &active_cells),
            CellSource::Configured(_) => self.rebalance_configured_cells(spread),
        }
    }

    /// Recompute demand-weighted CPU assignments for parent-cgroup cells.
    fn rebalance_parent_cells(&mut self, spread: f64, active_cells: &[u32]) -> Result<()> {
        // Build demand map from smoothed utilization
        let cell_demands: HashMap<u32, f64> = active_cells
            .iter()
            .map(|&cell_id| (cell_id, self.smoothed_util[cell_id as usize]))
            .collect();

        // Compute new assignments and check if they differ from current
        let (cell_assignments, cpu_assignments) = {
            let cpu_assignments = self
                .cell_manager()?
                .compute_demand_cpu_assignments(&cell_demands, self.enable_borrowing)
                .context("computing demand-weighted CPU assignments for rebalance")?;

            let changed = cpu_assignments.iter().any(|a| {
                self.cells
                    .get(&a.id)
                    .is_none_or(|cell| cell.cpus != a.primary)
            });

            // TODO(kkd): Need logic to check changed demand assignments for
            // subcells as well once demand weighted allocations work.

            if !changed {
                return Ok(());
            }

            (self.cell_manager()?.get_cell_assignments(), cpu_assignments)
        };

        let subcell_assignments = self.compute_subcell_assignments(&cpu_assignments)?;

        self.apply_cell_config(&cell_assignments, &cpu_assignments, &subcell_assignments)
            .context("applying rebalanced cell configuration to BPF")?;

        self.last_rebalance = Instant::now();
        self.rebalance_count += 1;
        self.metrics.rebalance_count = self.rebalance_count;

        info!(
            "Rebalanced CPUs (spread={:.1}%, count={}): {}",
            spread,
            self.rebalance_count,
            self.cell_manager()?.format_cell_config(&cpu_assignments)
        );

        Ok(())
    }

    /// Re-resolve configured cells with current demand and apply the result.
    fn rebalance_configured_cells(&mut self, spread: f64) -> Result<()> {
        let cpu_assignments = self
            .compute_and_apply_configured_cell_config()
            .context("recomputing configured cell configuration for rebalance")?;
        self.last_rebalance = Instant::now();
        self.last_config_refresh = Instant::now();
        self.rebalance_count += 1;
        self.metrics.rebalance_count = self.rebalance_count;

        info!(
            "Rebalanced configured cells (spread={:.1}%, count={}): {}",
            spread,
            self.rebalance_count,
            self.configured_cells()?
                .format_cell_config(&cpu_assignments)
        );

        Ok(())
    }

    fn compute_subcell_assignments(
        &self,
        cell_cpu_assignments: &[CpuAssignment],
    ) -> Result<Vec<Vec<CpuAssignment>>> {
        cell_cpu_assignments
            .iter()
            .map(|cell_assignment| {
                let recipients: Vec<CpuRecipient> = self
                    .cells
                    .get(&cell_assignment.id)
                    .map(|cell| {
                        let use_demand = self.enable_rebalancing
                            && (cell_assignment.id as usize) < MAX_CELLS
                            && cell.subcells.iter().any(|subcell| {
                                (subcell.id as usize) < MAX_SUBCELLS_PER_CELL
                                    && self.smoothed_subcell_demand[cell_assignment.id as usize]
                                        [subcell.id as usize]
                                        > 0.0
                            });
                        cell.subcells
                            .iter()
                            .map(|subcell| {
                                let weight = if use_demand
                                    && (subcell.id as usize) < MAX_SUBCELLS_PER_CELL
                                {
                                    self.smoothed_subcell_demand[cell_assignment.id as usize]
                                        [subcell.id as usize]
                                } else {
                                    1.0
                                };
                                CpuRecipient::unpinned(subcell.id, weight, &cell_assignment.primary)
                            })
                            .collect()
                    })
                    .unwrap_or_else(|| {
                        vec![CpuRecipient::unpinned(0, 1.0, &cell_assignment.primary)]
                    });

                CellManager::compute_subcell_cpu_assignments(
                    &cell_assignment.primary,
                    &recipients,
                    SHARE_SIBLING_SUBCELL_CPUS,
                    &self.cpu_to_llc,
                )
            })
            .collect()
    }

    fn refresh_bpf_subcells(&mut self, active_cells: &HashSet<u32>) -> Result<()> {
        for cell_id in active_cells {
            let bpf_cell = read_bpf_cell(&self.skel, *cell_id)?;
            let existing_matches: HashMap<u32, Vec<Vec<SubcellMatch>>> = self
                .cells
                .get(cell_id)
                .map(|cell| {
                    cell.subcells
                        .iter()
                        .map(|subcell| (subcell.id, subcell.matches.clone()))
                        .collect()
                })
                .unwrap_or_default();
            let existing_names: HashMap<u32, String> = self
                .cells
                .get(cell_id)
                .map(|cell| {
                    cell.subcells
                        .iter()
                        .map(|subcell| (subcell.id, subcell.name.clone()))
                        .collect()
                })
                .unwrap_or_default();
            let subcells: Vec<Subcell> = bpf_cell
                .subcells
                .iter()
                .filter(|subcell| subcell.in_use != 0)
                .map(|subcell| {
                    Ok(Subcell {
                        id: subcell.id,
                        name: existing_names
                            .get(&subcell.id)
                            .cloned()
                            .unwrap_or_else(|| default_subcell_name(subcell.id)),
                        primary: read_cpumask_from_bytes(&subcell.primary.mask)?,
                        borrowable: Some(read_cpumask_from_bytes(&subcell.borrowable.mask)?),
                        matches: existing_matches
                            .get(&subcell.id)
                            .cloned()
                            .unwrap_or_default(),
                    })
                })
                .collect::<Result<_>>()?;

            if subcells.is_empty() {
                bail!("Cell {} has no active subcells in BPF state", cell_id);
            }

            let cell = self.cells.get_mut(cell_id).ok_or_else(|| {
                anyhow::anyhow!("Cell {} missing during subcell refresh", cell_id)
            })?;
            cell.subcells = subcells;
        }

        Ok(())
    }

    /// Apply cell configuration to BPF.
    ///
    /// Writes the cell and CPU assignments to the BPF config struct and triggers
    /// the BPF program to apply the configuration.
    fn apply_cell_config(
        &mut self,
        cell_assignments: &[(u64, u32)],
        cpu_assignments: &[CpuAssignment],
        subcell_assignments: &[Vec<CpuAssignment>],
    ) -> Result<()> {
        let subcell_matches: HashMap<(u32, u32), Vec<Vec<SubcellMatch>>> = self
            .cells
            .iter()
            .flat_map(|(&cell_id, cell)| {
                cell.subcells
                    .iter()
                    .map(move |subcell| ((cell_id, subcell.id), subcell.matches.clone()))
            })
            .collect();

        let bss_data = self
            .skel
            .maps
            .bss_data
            .as_mut()
            .expect("bss_data must be available after scheduler load");

        let config = &mut bss_data.cell_config;

        // Zero out the config struct. This is necessary because:
        // 1. Cell IDs can be sparse (e.g., cells 0, 2, 3 if cell 1 was destroyed)
        // 2. We only write cpumasks for active cells, leaving gaps unwritten
        // 3. BPF iterates 0..num_cells and applies each cpumask
        // 4. Without zeroing, a gap (e.g., cell 1) would have a stale cpumask,
        //    causing CPUs to be assigned to an unused cell
        // Safety: cell_config is a plain data struct with no Drop impl
        unsafe {
            std::ptr::write_bytes(
                config as *mut _ as *mut u8,
                0,
                std::mem::size_of_val(config),
            );
        }

        if subcell_assignments.len() != cpu_assignments.len() {
            bail!(
                "Mismatched subcell assignments: {} cell assignments vs {} subcell groups",
                cpu_assignments.len(),
                subcell_assignments.len()
            );
        }

        if cell_assignments.len() > bpf_intf::consts_MAX_CELLS as usize {
            bail!(
                "Too many cell assignments: {} > MAX_CELLS ({})",
                cell_assignments.len(),
                bpf_intf::consts_MAX_CELLS
            );
        }
        config.num_cell_assignments = cell_assignments.len() as u32;

        for cell_id in 0..MAX_CELLS {
            for subcell_id in 0..MAX_SUBCELLS_PER_CELL {
                config.subcells[cell_id][subcell_id].id = subcell_id as u32;
            }
        }

        for (i, (cgid, cell_id)) in cell_assignments.iter().enumerate() {
            config.assignments[i].cgid = *cgid;
            config.assignments[i].cell_id = *cell_id;
        }

        // Set cell cpumasks and borrowable cpumasks
        let mut max_cell_id: u32 = 0;
        for (a, cell_subcells) in cpu_assignments.iter().zip(subcell_assignments.iter()) {
            validate_subcell_assignments(a, cell_subcells, SHARE_SIBLING_SUBCELL_CPUS)?;

            if a.id >= bpf_intf::consts_MAX_CELLS {
                bail!(
                    "Cell ID {} exceeds MAX_CELLS ({})",
                    a.id,
                    bpf_intf::consts_MAX_CELLS
                );
            }
            max_cell_id = max_cell_id.max(a.id + 1);

            write_cpumask_to_config(&a.primary, &mut config.cpumasks[a.id as usize].mask);

            if let Some(ref borrowable) = a.borrowable {
                write_cpumask_to_config(
                    borrowable,
                    &mut config.borrowable_cpumasks[a.id as usize].mask,
                );
            }

            for subcell in cell_subcells {
                let slot = subcell.id as usize;
                config.subcells[a.id as usize][slot].id = subcell.id;
                config.subcells[a.id as usize][slot].in_use = 1;
                let matches = subcell_matches
                    .get(&(a.id, subcell.id))
                    .map_or(&[][..], |matches| matches.as_slice());
                write_subcell_matches(matches, &mut config.subcells[a.id as usize][slot])
                    .with_context(|| {
                        format!(
                            "writing subcell matches for cell {} subcell {}",
                            a.id, subcell.id
                        )
                    })?;
                write_cpumask_to_config(
                    &subcell.primary,
                    &mut config.subcells[a.id as usize][slot].primary.mask,
                );
                if let Some(ref borrowable) = subcell.borrowable {
                    write_cpumask_to_config(
                        borrowable,
                        &mut config.subcells[a.id as usize][slot].borrowable.mask,
                    );
                }
            }
        }
        config.num_cells = max_cell_id;

        // Trigger the BPF program to apply the configuration
        let prog = &mut self.skel.progs.apply_cell_config;
        let out = prog
            .test_run(ProgramInput::default())
            .context("Failed to run apply_cell_config BPF program")?;
        if out.return_value != 0 {
            bail!(
                "apply_cell_config BPF program returned error {} (num_assignments={}, num_cells={})",
                out.return_value as i32,
                cell_assignments.len(),
                cpu_assignments.len()
            );
        }

        Ok(())
    }

    fn get_metrics(&self) -> Metrics {
        self.metrics.clone()
    }

    fn calculate_distribution_stats(
        &self,
        queue_counts: &[u64; QUEUE_STATS_IDX.len()],
        global_queue_decisions: u64,
        scope_queue_decisions: u64,
        scope_affn_viols: u64,
        scope_steals: u64,
        scope_pin_skips: u64,
    ) -> Result<DistributionStats> {
        // First % on the line: share of global work
        // We know global_queue_decisions is non-zero.
        let share_of_global =
            100.0 * (scope_queue_decisions as f64) / (global_queue_decisions as f64);

        // Each queue's % of the scope total
        let queue_pct = if scope_queue_decisions == 0 {
            debug!("No queue decisions in scope, zeroing out queue distribution");
            [0.0; QUEUE_STATS_IDX.len()]
        } else {
            core::array::from_fn(|i| {
                100.0 * (queue_counts[i] as f64) / (scope_queue_decisions as f64)
            })
        };

        // These are summed differently for the global and per-cell totals.
        let affinity_violations_percent = if scope_queue_decisions == 0 {
            debug!("No queue decisions in scope, zeroing out affinity violations");
            0.0
        } else {
            100.0 * (scope_affn_viols as f64) / (scope_queue_decisions as f64)
        };

        let steal_pct = if scope_queue_decisions == 0 {
            0.0
        } else {
            100.0 * (scope_steals as f64) / (scope_queue_decisions as f64)
        };

        let pin_skip_pct = if scope_queue_decisions == 0 {
            0.0
        } else {
            100.0 * (scope_pin_skips as f64) / (scope_queue_decisions as f64)
        };

        const EXPECTED_QUEUES: usize = 4;
        if queue_pct.len() != EXPECTED_QUEUES {
            bail!(
                "Expected {} queues, got {}",
                EXPECTED_QUEUES,
                queue_pct.len()
            );
        }

        Ok(DistributionStats {
            total_decisions: scope_queue_decisions,
            share_of_decisions_pct: share_of_global,
            local_q_pct: queue_pct[0],
            cpu_q_pct: queue_pct[1],
            cell_q_pct: queue_pct[2],
            borrowed_pct: queue_pct[3],
            affn_viol_pct: affinity_violations_percent,
            steal_pct,
            pin_skip_pct,
            global_queue_decisions,
        })
    }

    // Queue stats for the whole node
    fn update_and_log_global_queue_stats(
        &mut self,
        global_queue_decisions: u64,
        cell_stats_delta: &[[u64; NR_CSTATS]; MAX_CELLS],
    ) -> Result<()> {
        // Get total of each queue summed over all cells
        let mut queue_counts = [0; QUEUE_STATS_IDX.len()];
        for cells in 0..MAX_CELLS {
            for (i, stat) in QUEUE_STATS_IDX.iter().enumerate() {
                queue_counts[i] += cell_stats_delta[cells][*stat as usize];
            }
        }

        let prefix = "Total Decisions:";

        // Here we want to sum the affinity violations over all cells.
        let scope_affn_viols: u64 = cell_stats_delta
            .iter()
            .map(|&cell| cell[bpf_intf::cell_stat_idx_CSTAT_AFFN_VIOL as usize])
            .sum::<u64>();

        // Sum steals over all cells
        let scope_steals: u64 = cell_stats_delta
            .iter()
            .map(|&cell| cell[bpf_intf::cell_stat_idx_CSTAT_STEAL as usize])
            .sum::<u64>();

        // Sum pin skips over all cells
        let scope_pin_skips: u64 = cell_stats_delta
            .iter()
            .map(|&cell| cell[bpf_intf::cell_stat_idx_CSTAT_PIN_SKIP as usize])
            .sum::<u64>();

        // Special case where the number of scope decisions == number global decisions
        let stats = self
            .calculate_distribution_stats(
                &queue_counts,
                global_queue_decisions,
                global_queue_decisions,
                scope_affn_viols,
                scope_steals,
                scope_pin_skips,
            )
            .context("calculating global queue distribution stats")?;

        self.metrics.update(&stats);

        // Slice shrink stats bypass DistributionStats — they're raw event counts
        let sum = |idx: usize| -> u64 { cell_stats_delta.iter().map(|c| c[idx]).sum() };
        self.metrics.drain_cnt = sum(bpf_intf::cell_stat_idx_CSTAT_DRAIN_CNT as usize);
        self.metrics.drain_affn_cnt = sum(bpf_intf::cell_stat_idx_CSTAT_DRAIN_AFFN_CNT as usize);
        self.metrics.slice_shrink_max =
            sum(bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_MAX as usize);
        self.metrics.slice_shrink_proportional =
            sum(bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_PROPORTIONAL as usize);
        self.metrics.slice_shrink_min =
            sum(bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_MIN as usize);
        self.metrics.slice_shrink = self.metrics.slice_shrink_max
            + self.metrics.slice_shrink_proportional
            + self.metrics.slice_shrink_min;

        trace!("{} {}", prefix, stats);

        Ok(())
    }

    // Print out the per-cell stats
    fn update_and_log_cell_queue_stats(
        &mut self,
        global_queue_decisions: u64,
        cell_stats_delta: &[[u64; NR_CSTATS]; MAX_CELLS],
    ) -> Result<()> {
        for cell in 0..MAX_CELLS {
            let cell_queue_decisions = QUEUE_STATS_IDX
                .iter()
                .map(|&stat| cell_stats_delta[cell][stat as usize])
                .sum::<u64>();

            // FIXME: This should really query if the cell is enabled or not.
            if cell_queue_decisions == 0 {
                continue;
            }

            let mut queue_counts = [0; QUEUE_STATS_IDX.len()];
            for (i, &stat) in QUEUE_STATS_IDX.iter().enumerate() {
                queue_counts[i] = cell_stats_delta[cell][stat as usize];
            }

            const MIN_CELL_WIDTH: usize = 2;
            let cell_width: usize = max(MIN_CELL_WIDTH, (MAX_CELLS as f64).log10().ceil() as usize);

            let prefix = format!("        Cell {:width$}:", cell, width = cell_width);

            // Sum affinity violations for this cell
            let scope_affn_viols: u64 =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_AFFN_VIOL as usize];

            // Steals for this cell
            let scope_steals: u64 =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_STEAL as usize];

            // Pin skips for this cell
            let scope_pin_skips: u64 =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_PIN_SKIP as usize];

            let stats = self
                .calculate_distribution_stats(
                    &queue_counts,
                    global_queue_decisions,
                    cell_queue_decisions,
                    scope_affn_viols,
                    scope_steals,
                    scope_pin_skips,
                )
                .with_context(|| {
                    format!("calculating queue distribution stats for cell {}", cell)
                })?;

            let cell_metrics = self.metrics.cells.entry(cell as u32).or_default();
            cell_metrics.update(&stats);

            // Raw event counts bypass DistributionStats.
            cell_metrics.drain_cnt =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_DRAIN_CNT as usize];
            cell_metrics.drain_affn_cnt =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_DRAIN_AFFN_CNT as usize];
            cell_metrics.slice_shrink_max =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_MAX as usize];
            cell_metrics.slice_shrink_proportional = cell_stats_delta[cell]
                [bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_PROPORTIONAL as usize];
            cell_metrics.slice_shrink_min =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_SLICE_SHRINK_MIN as usize];
            cell_metrics.slice_shrink = cell_metrics.slice_shrink_max
                + cell_metrics.slice_shrink_proportional
                + cell_metrics.slice_shrink_min;

            trace!("{} {}", prefix, stats);
        }
        Ok(())
    }

    fn update_drain_metrics(&mut self, cell_stats_delta: &[[u64; NR_CSTATS]; MAX_CELLS]) {
        let mut total = 0;
        let mut affn_total = 0;

        for cell in 0..MAX_CELLS {
            let drain_cnt =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_DRAIN_CNT as usize];
            let drain_affn_cnt =
                cell_stats_delta[cell][bpf_intf::cell_stat_idx_CSTAT_DRAIN_AFFN_CNT as usize];
            total += drain_cnt;
            affn_total += drain_affn_cnt;

            if let Some(cell_metrics) = self.metrics.cells.get_mut(&(cell as u32)) {
                cell_metrics.drain_cnt = drain_cnt;
                cell_metrics.drain_affn_cnt = drain_affn_cnt;
            } else if drain_cnt > 0 || drain_affn_cnt > 0 {
                let cell_metrics = self.metrics.cells.entry(cell as u32).or_default();
                cell_metrics.drain_cnt = drain_cnt;
                cell_metrics.drain_affn_cnt = drain_affn_cnt;
            }
        }

        self.metrics.drain_cnt = total;
        self.metrics.drain_affn_cnt = affn_total;
    }

    fn log_all_queue_stats(
        &mut self,
        cell_stats_delta: &[[u64; NR_CSTATS]; MAX_CELLS],
    ) -> Result<()> {
        // Get total decisions
        let global_queue_decisions: u64 = cell_stats_delta
            .iter()
            .flat_map(|cell| QUEUE_STATS_IDX.iter().map(|&idx| cell[idx as usize]))
            .sum();

        self.update_drain_metrics(cell_stats_delta);

        if global_queue_decisions == 0 {
            if self.metrics.drain_cnt == 0 {
                warn!("No queueing decisions made globally");
            }
            return Ok(());
        }

        self.update_and_log_global_queue_stats(global_queue_decisions, cell_stats_delta)
            .context("updating global queue stats")?;

        self.update_and_log_cell_queue_stats(global_queue_decisions, cell_stats_delta)
            .context("updating per-cell queue stats")?;

        Ok(())
    }

    fn calculate_cell_stat_delta(
        &mut self,
        cpu_ctxs: &[bpf_intf::cpu_ctx],
    ) -> Result<[[u64; NR_CSTATS]; MAX_CELLS]> {
        let mut cell_stats_delta = [[0_u64; NR_CSTATS]; MAX_CELLS];

        // Loop over cells and stats first, then CPU contexts
        // TODO: We should loop over the in_use cells only.
        for cell in 0..MAX_CELLS {
            for stat in 0..NR_CSTATS {
                let mut cur_cell_stat = 0;

                // Accumulate stats from all CPUs
                for cpu_ctx in cpu_ctxs.iter() {
                    cur_cell_stat += cpu_ctx.cstats[cell][stat];
                }

                // Calculate delta and update previous stat
                cell_stats_delta[cell][stat] = cur_cell_stat - self.prev_cell_stats[cell][stat];
                self.prev_cell_stats[cell][stat] = cur_cell_stat;
            }
        }
        Ok(cell_stats_delta)
    }

    /// Cgroup path backing `cell_id` in the active cell source.
    fn cgroup_path_for_cell(&self, cell_id: u32) -> String {
        match &self.cell_source {
            CellSource::Parent(cell_manager) => cell_manager.cgroup_path_for_cell(cell_id),
            CellSource::Configured(configured_cells) => {
                configured_cells.cgroup_path_for_cell(cell_id)
            }
        }
    }

    /// Collect metrics and out various debugging data like per cell stats, per-cpu stats, etc.
    fn collect_metrics(&mut self) -> Result<()> {
        // the demand deltas are measured up to here, see collect_demand_metrics()
        let sampled_at = Instant::now();
        let cpu_ctxs = read_cpu_ctxs(&self.skel).context("reading per-CPU contexts for metrics")?;

        let cell_stats_delta = self
            .calculate_cell_stat_delta(&cpu_ctxs)
            .context("calculating cell stat deltas")?;

        self.log_all_queue_stats(&cell_stats_delta)
            .context("logging queue stats")?;

        // Mirror the sticky holdout flag on every collection, independent of the
        // zero-decisions early return inside log_all_queue_stats above.
        let enforced_holdout = match &self.cell_source {
            CellSource::Parent(cell_manager) => cell_manager.enforced_holdout(),
            CellSource::Configured(_) => self.configured_enforced_holdout,
        };
        self.metrics.enforced_holdout = enforced_holdout as u64;

        self.collect_demand_metrics(&cpu_ctxs, sampled_at)
            .context("collecting demand metrics")?;
        self.collect_subcell_demand_metrics(&cpu_ctxs)
            .context("collecting subcell demand metrics")?;

        for (cell_id, cell) in &self.cells {
            trace!("CELL[{}]: {}", cell_id, cell.cpus);
        }

        for (cell_id, cell) in self.cells.iter() {
            let cgroup_path = self.cgroup_path_for_cell(*cell_id);
            // Assume we have a CellMetrics entry if we have a known cell
            self.metrics
                .cells
                .entry(*cell_id)
                .and_modify(|cell_metrics| {
                    cell_metrics.num_cpus = cell.cpus.weight() as u32;
                    cell_metrics.cgroup_path = cgroup_path;
                    for subcell in &cell.subcells {
                        cell_metrics
                            .subcells
                            .entry(subcell.id)
                            .or_default()
                            .update_name_and_cpus(&subcell.name, subcell.primary.weight() as u32);
                    }
                });
        }
        self.metrics.num_cells = self.cells.len() as u32;

        Ok(())
    }

    /// Compute per-cell demand metrics (utilization, borrowed, lent) from BPF running_ns counters.
    fn collect_demand_metrics(
        &mut self,
        cpu_ctxs: &[bpf_intf::cpu_ctx],
        sampled_at: Instant,
    ) -> Result<()> {
        // Deltas over the time that really passed: an iteration can take well
        // over the nominal interval and events wake the loop early. Let half
        // an interval build up before sampling again. sampled_at is taken right
        // before the counters are read.
        let half_interval = self.monitor_interval / 2;
        let interval_ns = match self.demand_sampled_at {
            Some(prev) if sampled_at.duration_since(prev) < half_interval => return Ok(()),
            Some(prev) => Some(sampled_at.duration_since(prev).as_nanos() as u64),
            None => None,
        };
        self.demand_sampled_at = Some(sampled_at);

        // Per-cell cumulative counters derived from BPF per-CPU running_ns:
        //   total_running_ns[c] = total time tasks in cell c ran (on any CPU)
        //   on_own_ns[c]        = time tasks in cell c ran on CPUs owned by cell c
        //   lent_ns[c]          = time foreign tasks ran on CPUs owned by cell c
        let mut total_running_ns = [0u64; MAX_CELLS];
        let mut on_own_ns = [0u64; MAX_CELLS];
        let mut lent_ns = [0u64; MAX_CELLS];

        for cpu_ctx in cpu_ctxs.iter() {
            let owner = cpu_ctx.cell as usize;
            for cell in 0..MAX_CELLS {
                let ns = cpu_ctx.running_ns[cell];
                total_running_ns[cell] += ns;
                if owner == cell {
                    on_own_ns[cell] += ns;
                }
            }
            if owner >= MAX_CELLS {
                bail!(
                    "CPU has invalid cell assignment {} (MAX_CELLS={})",
                    owner,
                    MAX_CELLS
                );
            }
            // Lent time: non-owner cell tasks running on this CPU
            let total_on_cpu: u64 = cpu_ctx.running_ns.iter().sum();
            let owner_on_cpu = cpu_ctx.running_ns[owner];
            lent_ns[owner] += total_on_cpu.saturating_sub(owner_on_cpu);
        }

        // the first call only seeds the deltas
        let Some(interval_ns) = interval_ns else {
            self.prev_cell_running_ns = total_running_ns;
            self.prev_cell_own_ns = on_own_ns;
            self.prev_cell_lent_ns = lent_ns;
            return Ok(());
        };

        let mut global_running_delta = 0u64;
        let mut global_borrowed_delta = 0u64;
        let mut global_lent_delta = 0u64;
        let mut global_capacity = 0u64;

        for cell in 0..MAX_CELLS {
            let delta_running =
                total_running_ns[cell].saturating_sub(self.prev_cell_running_ns[cell]);
            let delta_on_own = on_own_ns[cell].saturating_sub(self.prev_cell_own_ns[cell]);
            let delta_lent = lent_ns[cell].saturating_sub(self.prev_cell_lent_ns[cell]);

            self.prev_cell_running_ns[cell] = total_running_ns[cell];
            self.prev_cell_own_ns[cell] = on_own_ns[cell];
            self.prev_cell_lent_ns[cell] = lent_ns[cell];

            // Borrowed = ran somewhere other than own CPUs
            let delta_borrowed = delta_running.saturating_sub(delta_on_own);

            // After a cell is destroyed, BPF may still report residual
            // running_ns from the previous interval. Skip stale cells.
            let Some(cell_info) = self.cells.get(&(cell as u32)) else {
                continue;
            };

            let nr_cpus = cell_info.cpus.weight() as u64;
            if nr_cpus == 0 {
                // Transitional: the applied masks have not been read back yet.
                // Skip the sample but keep the smoothed history.
                continue;
            }

            // capacity = total available CPU-time this interval
            let capacity = nr_cpus * interval_ns;
            // util: fraction of own capacity consumed by own tasks
            let util_pct = 100.0 * (delta_running as f64) / (capacity as f64);
            // demand_borrow: fraction of running time that was borrowed from other cells
            let demand_borrow_pct = if delta_running > 0 {
                100.0 * (delta_borrowed as f64) / (delta_running as f64)
            } else {
                0.0
            };
            // lent: fraction of own capacity used by foreign tasks
            let lent_pct = 100.0 * (delta_lent as f64) / (capacity as f64);

            // Update EWMA-smoothed utilization
            if self.enable_rebalancing {
                self.smoothed_util[cell] = self.demand_smoothing * util_pct
                    + (1.0 - self.demand_smoothing) * self.smoothed_util[cell];
            }

            self.metrics
                .cells
                .entry(cell as u32)
                .or_default()
                .update_demand(util_pct, demand_borrow_pct, lent_pct);

            // Update smoothed_util_pct in metrics
            if self.enable_rebalancing {
                self.metrics
                    .cells
                    .entry(cell as u32)
                    .or_default()
                    .smoothed_util_pct = self.smoothed_util[cell];
            }

            global_running_delta = global_running_delta.saturating_add(delta_running);
            global_borrowed_delta = global_borrowed_delta.saturating_add(delta_borrowed);
            global_lent_delta = global_lent_delta.saturating_add(delta_lent);
            global_capacity = global_capacity.saturating_add(capacity);
        }

        let global_util_pct = if global_capacity > 0 {
            100.0 * (global_running_delta as f64) / (global_capacity as f64)
        } else {
            0.0
        };
        let global_borrow_pct = if global_running_delta > 0 {
            100.0 * (global_borrowed_delta as f64) / (global_running_delta as f64)
        } else {
            0.0
        };
        let global_lent_pct = if global_capacity > 0 {
            100.0 * (global_lent_delta as f64) / (global_capacity as f64)
        } else {
            0.0
        };

        self.metrics
            .update_demand(global_util_pct, global_borrow_pct, global_lent_pct);

        Ok(())
    }

    /// Compute per-subcell demand metrics from BPF running_ns counters.
    fn collect_subcell_demand_metrics(&mut self, cpu_ctxs: &[bpf_intf::cpu_ctx]) -> Result<()> {
        let mut total_running_ns = [[0u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS];
        let mut on_own_ns = [[0u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS];
        let mut lent_ns = [[0u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS];
        let mut queued_ns = [[0u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS];
        let mut demand_sum = [[0u64; MAX_SUBCELLS_PER_CELL]; MAX_CELLS];
        let active_subcells: HashMap<(usize, usize), (usize, String)> = self
            .cells
            .iter()
            .flat_map(|(&cell_id, cell)| {
                cell.subcells.iter().map(move |subcell| {
                    (
                        (cell_id as usize, subcell.id as usize),
                        (subcell.primary.weight(), subcell.name.clone()),
                    )
                })
            })
            .collect();
        let subcell_accounts = read_subcell_accounts(&self.skel, active_subcells.keys().copied())
            .context("reading subcell accounts")?;

        for (cpu, cpu_ctx) in cpu_ctxs.iter().enumerate() {
            let owner_cell = cpu_ctx.cell as usize;
            let owner_subcell = cpu_ctx.subcell as usize;
            if owner_cell >= MAX_CELLS || owner_subcell >= MAX_SUBCELLS_PER_CELL {
                bail!(
                    "CPU has invalid subcell assignment {}:{} (max {}:{})",
                    owner_cell,
                    owner_subcell,
                    MAX_CELLS,
                    MAX_SUBCELLS_PER_CELL
                );
            }

            let mut total_on_cpu = 0u64;
            for (&(cell, subcell), per_cpu_accounts) in &subcell_accounts {
                let ns = per_cpu_accounts[cpu].running_ns;
                total_running_ns[cell][subcell] =
                    total_running_ns[cell][subcell].saturating_add(ns);
                queued_ns[cell][subcell] =
                    queued_ns[cell][subcell].saturating_add(per_cpu_accounts[cpu].queued_ns);
                demand_sum[cell][subcell] =
                    demand_sum[cell][subcell].wrapping_add(per_cpu_accounts[cpu].demand_sum);
                total_on_cpu = total_on_cpu.saturating_add(ns);
                if owner_cell == cell && owner_subcell == subcell {
                    on_own_ns[cell][subcell] = on_own_ns[cell][subcell].saturating_add(ns);
                }
            }

            let owner_on_cpu = subcell_accounts
                .get(&(owner_cell, owner_subcell))
                .map_or(0, |per_cpu_accounts| per_cpu_accounts[cpu].running_ns);
            lent_ns[owner_cell][owner_subcell] = lent_ns[owner_cell][owner_subcell]
                .saturating_add(total_on_cpu.saturating_sub(owner_on_cpu));
        }

        let interval_ns = self.monitor_interval.as_nanos() as u64;

        for cell in 0..MAX_CELLS {
            for subcell in 0..MAX_SUBCELLS_PER_CELL {
                let delta_running = total_running_ns[cell][subcell]
                    .saturating_sub(self.prev_subcell_running_ns[cell][subcell]);
                let delta_on_own = on_own_ns[cell][subcell]
                    .saturating_sub(self.prev_subcell_own_ns[cell][subcell]);
                let delta_lent =
                    lent_ns[cell][subcell].saturating_sub(self.prev_subcell_lent_ns[cell][subcell]);
                let delta_queued = queued_ns[cell][subcell]
                    .saturating_sub(self.prev_subcell_queued_ns[cell][subcell]);
                // The demand sum is allowed to wrap; a wrapping delta stays exact.
                let delta_demand = demand_sum[cell][subcell]
                    .wrapping_sub(self.prev_subcell_demand_sum[cell][subcell]);

                self.prev_subcell_running_ns[cell][subcell] = total_running_ns[cell][subcell];
                self.prev_subcell_own_ns[cell][subcell] = on_own_ns[cell][subcell];
                self.prev_subcell_lent_ns[cell][subcell] = lent_ns[cell][subcell];
                self.prev_subcell_queued_ns[cell][subcell] = queued_ns[cell][subcell];
                self.prev_subcell_demand_sum[cell][subcell] = demand_sum[cell][subcell];

                if delta_running == 0 && delta_lent == 0 && delta_queued == 0 && delta_demand == 0 {
                    continue;
                }

                let Some((nr_cpus, name)) = active_subcells.get(&(cell, subcell)) else {
                    continue;
                };
                if *nr_cpus == 0 {
                    // Transitional: the applied masks have not been read back
                    // yet. Skip the sample but keep the smoothed history.
                    continue;
                }

                let capacity = (*nr_cpus as u64) * interval_ns;
                let delta_borrowed = delta_running.saturating_sub(delta_on_own);
                let util_pct = 100.0 * (delta_running as f64) / (capacity as f64);
                let demand_borrow_pct = if delta_running > 0 {
                    100.0 * (delta_borrowed as f64) / (delta_running as f64)
                } else {
                    0.0
                };
                let lent_pct = 100.0 * (delta_lent as f64) / (capacity as f64);
                // Time-averaged runnable task count: running plus queued time
                // over the interval. Unlike util_pct it is not bounded by the
                // subcell's current capacity.
                let load = delta_running.saturating_add(delta_queued) as f64 / interval_ns as f64;
                // Sum of the tasks' unconstrained demand, in CPUs.
                let demand = delta_demand as f64 / (DEMAND_ONE * interval_ns as f64);

                if self.enable_rebalancing {
                    self.smoothed_subcell_util[cell][subcell] = self.demand_smoothing * util_pct
                        + (1.0 - self.demand_smoothing) * self.smoothed_subcell_util[cell][subcell];
                    self.smoothed_subcell_load[cell][subcell] = self.demand_smoothing * load
                        + (1.0 - self.demand_smoothing) * self.smoothed_subcell_load[cell][subcell];
                    self.smoothed_subcell_demand[cell][subcell] = self.demand_smoothing * demand
                        + (1.0 - self.demand_smoothing)
                            * self.smoothed_subcell_demand[cell][subcell];
                }

                let subcell_metrics = self
                    .metrics
                    .cells
                    .entry(cell as u32)
                    .or_default()
                    .subcells
                    .entry(subcell as u32)
                    .or_default();
                subcell_metrics.update_demand(
                    util_pct,
                    demand_borrow_pct,
                    lent_pct,
                    delta_running,
                    delta_borrowed,
                    delta_lent,
                    delta_queued,
                    load,
                    demand,
                );
                subcell_metrics.update_name_and_cpus(name, *nr_cpus as u32);
                if self.enable_rebalancing {
                    subcell_metrics.smoothed_util_pct = self.smoothed_subcell_util[cell][subcell];
                    subcell_metrics.smoothed_load = self.smoothed_subcell_load[cell][subcell];
                    subcell_metrics.smoothed_demand = self.smoothed_subcell_demand[cell][subcell];
                }
            }
        }

        Ok(())
    }

    /// Write applied_cpuset_seq to BSS, closing the rejection-skip window.
    fn update_applied_cpuset_seq(&mut self) {
        unsafe {
            let ptr = &mut self
                .skel
                .maps
                .bss_data
                .as_mut()
                .expect("BUG: bss_data missing after scheduler load")
                .applied_cpuset_seq as *mut u32;
            std::ptr::write_volatile(ptr, self.last_cpuset_seq);
        }
    }

    /// Check if any cell's cpuset was modified and recompute if so.
    fn check_cpuset_changes(&mut self) -> Result<()> {
        let current_seq = unsafe {
            let ptr = &self
                .skel
                .maps
                .bss_data
                .as_ref()
                .expect("BUG: bss_data missing after scheduler load")
                .cpuset_seq as *const u32;
            (ptr as *const AtomicU32)
                .as_ref()
                .expect("BUG: cpuset_seq pointer cast yielded null")
                .load(Ordering::Acquire)
        };

        if current_seq == self.last_cpuset_seq {
            return Ok(());
        }
        self.last_cpuset_seq = current_seq;

        match self.cell_source {
            CellSource::Parent(_) => self.recompute_parent_cells_after_cpuset_change(),
            CellSource::Configured(_) => self.recompute_configured_cells_after_cpuset_change(),
        }
    }

    fn recompute_parent_cells_after_cpuset_change(&mut self) -> Result<()> {
        if !self
            .cell_manager_mut()?
            .refresh_cpusets()
            .context("refreshing cell cpusets")?
        {
            // seq changed but no cpusets on our cells changed
            self.update_applied_cpuset_seq();
            return Ok(());
        }

        let cpu_assignments = self
            .compute_and_apply_cell_config(&[])
            .context("recomputing cell configuration after cpuset change")?;
        self.update_applied_cpuset_seq();
        info!(
            "Cpuset change detected, recomputed config: {}",
            self.cell_manager()?.format_cell_config(&cpu_assignments)
        );
        Ok(())
    }

    fn recompute_configured_cells_after_cpuset_change(&mut self) -> Result<()> {
        let cpu_assignments = self
            .compute_and_apply_configured_cell_config()
            .context("recomputing configured cell configuration after cpuset change")?;
        self.last_config_refresh = Instant::now();
        self.update_applied_cpuset_seq();
        info!(
            "Cpuset change detected, recomputed configured cells: {}",
            self.configured_cells()?
                .format_cell_config(&cpu_assignments)
        );
        Ok(())
    }

    /// Rescan the cgroup tree for --cell-config mode once the interval elapses.
    fn maybe_refresh_configured_cells(&mut self) -> Result<()> {
        if !matches!(self.cell_source, CellSource::Configured(_))
            || self.last_config_refresh.elapsed() < self.reconfiguration_interval
        {
            return Ok(());
        }

        let cpu_assignments = self
            .compute_and_apply_configured_cell_config()
            .context("recomputing configured cells")?;
        self.last_config_refresh = Instant::now();

        info!(
            "Refreshed configured cells: {}",
            self.configured_cells()?
                .format_cell_config(&cpu_assignments)
        );
        Ok(())
    }

    fn refresh_bpf_cells(&mut self) -> Result<()> {
        let applied_configuration = unsafe {
            let ptr = &self
                .skel
                .maps
                .bss_data
                .as_ref()
                .expect("BUG: bss_data missing after scheduler load")
                .applied_configuration_seq as *const u32;
            (ptr as *const std::sync::atomic::AtomicU32)
                .as_ref()
                .expect("BUG: applied_configuration_seq pointer cast yielded null")
                .load(std::sync::atomic::Ordering::Acquire)
        };
        if self
            .last_configuration_seq
            .is_some_and(|seq| applied_configuration == seq)
        {
            return Ok(());
        }
        // collect all cpus per cell.
        let mut cell_to_cpus: HashMap<u32, Cpumask> = HashMap::new();
        let cpu_ctxs =
            read_cpu_ctxs(&self.skel).context("reading per-CPU contexts for BPF cell refresh")?;
        for (i, cpu_ctx) in cpu_ctxs.iter().enumerate() {
            cell_to_cpus
                .entry(cpu_ctx.cell)
                .or_insert_with(Cpumask::new)
                .set_cpu(i)
                .expect("set cpu in existing mask");
        }

        // Create cells we don't have yet, drop cells that are no longer in use.
        // If we continue to drop cell metrics once a cell is removed, we'll need to make sure we
        // flush metrics for a cell before we remove it completely.
        //
        // IMPORTANT: We determine which cells exist based on CPU assignments (which are
        // synchronized by applied_configuration_seq), NOT by reading the in_use field
        // separately. This avoids a TOCTOU race where a cell's in_use is set before
        // CPUs are assigned.

        // Cell 0 (root cell) always exists even if it has no CPUs temporarily
        let cells_with_cpus: HashSet<u32> = cell_to_cpus.keys().copied().collect();
        let mut active_cells = cells_with_cpus.clone();
        active_cells.insert(0);

        for cell_idx in &active_cells {
            let cpus = cell_to_cpus
                .get(cell_idx)
                .cloned()
                .unwrap_or_else(Cpumask::new);
            self.cells.entry(*cell_idx).or_insert_with(Cell::new).cpus = cpus;
            self.metrics.cells.insert(*cell_idx, CellMetrics::default());
        }

        // Remove cells that no longer have CPUs assigned
        self.cells.retain(|&k, _| active_cells.contains(&k));
        self.metrics.cells.retain(|&k, _| active_cells.contains(&k));

        self.refresh_bpf_subcells(&active_cells)?;

        self.last_configuration_seq = Some(applied_configuration);

        Ok(())
    }
}

fn validate_subcell_assignments(
    cell: &CpuAssignment,
    subcells: &[CpuAssignment],
    share_sibling_cpus: bool,
) -> Result<()> {
    if subcells.is_empty() {
        bail!("Cell {} has no subcell assignments", cell.id);
    }
    if subcells.len() > MAX_SUBCELLS_PER_CELL {
        bail!(
            "Cell {} has too many subcells: {} > MAX_SUBCELLS_PER_CELL ({})",
            cell.id,
            subcells.len(),
            MAX_SUBCELLS_PER_CELL
        );
    }

    let mut seen = [false; MAX_SUBCELLS_PER_CELL];
    let mut primary_union = Cpumask::new();

    for subcell in subcells {
        if subcell.id >= bpf_intf::consts_MAX_SUBCELLS_PER_CELL {
            bail!(
                "Subcell ID {} for cell {} exceeds MAX_SUBCELLS_PER_CELL ({})",
                subcell.id,
                cell.id,
                bpf_intf::consts_MAX_SUBCELLS_PER_CELL
            );
        }

        let slot = subcell.id as usize;
        if seen[slot] {
            bail!("Duplicate subcell ID {} for cell {}", subcell.id, cell.id);
        }
        seen[slot] = true;

        if subcell.primary.weight() == 0 {
            bail!(
                "Subcell {} in cell {} has no primary CPUs",
                subcell.id,
                cell.id
            );
        }
        if subcell.primary.and(&cell.primary.not()).weight() != 0 {
            bail!(
                "Subcell {} primary CPUs are outside cell {}",
                subcell.id,
                cell.id
            );
        }
        if primary_union.and(&subcell.primary).weight() != 0 {
            bail!(
                "Subcell {} overlaps another subcell in cell {}",
                subcell.id,
                cell.id
            );
        }

        if share_sibling_cpus {
            let expected = cell.primary.and(&subcell.primary.not());
            let borrowable = subcell.borrowable.as_ref().ok_or_else(|| {
                anyhow::anyhow!(
                    "Subcell {} in cell {} is missing its borrowable mask",
                    subcell.id,
                    cell.id
                )
            })?;
            if borrowable != &expected {
                bail!(
                    "Subcell {} in cell {} has an invalid borrowable mask",
                    subcell.id,
                    cell.id
                );
            }
        } else if subcell.borrowable.is_some() {
            bail!(
                "Subcell {} in cell {} has borrowing disabled but a borrowable mask",
                subcell.id,
                cell.id
            );
        }

        primary_union = primary_union.or(&subcell.primary);
    }

    if primary_union != cell.primary {
        bail!("Subcell primary CPUs do not cover cell {}", cell.id);
    }

    Ok(())
}

fn write_cpumask_to_config(cpumask: &Cpumask, dest: &mut [u8]) {
    let raw_slice = cpumask.as_raw_slice();
    for (word_idx, word) in raw_slice.iter().enumerate() {
        let byte_start = word_idx * 8;
        let bytes = word.to_le_bytes();
        for (j, byte) in bytes.iter().enumerate() {
            let idx = byte_start + j;
            if idx < dest.len() {
                dest[idx] = *byte;
            }
        }
    }
}

fn write_subcell_matches(
    matches: &[Vec<SubcellMatch>],
    dest: &mut types::subcell_config,
) -> Result<()> {
    if matches.len() > bpf_intf::consts_MAX_SUBCELL_MATCH_ORS as usize {
        bail!(
            "Too many subcell match OR groups: {} > {}",
            matches.len(),
            bpf_intf::consts_MAX_SUBCELL_MATCH_ORS
        );
    }

    dest.nr_match_ors = matches.len() as u32;
    for (or_idx, ands) in matches.iter().enumerate() {
        if ands.len() > bpf_intf::consts_MAX_SUBCELL_MATCH_ANDS as usize {
            bail!(
                "Too many subcell match AND terms: {} > {}",
                ands.len(),
                bpf_intf::consts_MAX_SUBCELL_MATCH_ANDS
            );
        }

        dest.matches[or_idx].nr_matches = ands.len() as u32;
        for (and_idx, subcell_match) in ands.iter().enumerate() {
            match subcell_match {
                SubcellMatch::CommPrefix(prefix) => {
                    dest.matches[or_idx].matches[and_idx].kind =
                        bpf_intf::subcell_match_kind_SUBCELL_MATCH_COMM_PREFIX;
                    copy_cstr(&mut dest.matches[or_idx].matches[and_idx].value, prefix)?;
                }
            }
        }
    }

    Ok(())
}

fn copy_cstr(dest: &mut [i8], src: &str) -> Result<()> {
    if dest.is_empty() {
        return Ok(());
    }
    if src.as_bytes().contains(&0) {
        bail!("subcell match string contains NUL byte");
    }

    for value in dest.iter_mut() {
        *value = 0;
    }

    let len = src.len().min(dest.len() - 1);
    for (dst, src) in dest.iter_mut().zip(src.as_bytes().iter()).take(len) {
        *dst = *src as i8;
    }

    Ok(())
}

fn default_subcell_name(id: u32) -> String {
    if id == 0 {
        "rest".to_string()
    } else {
        format!("subcell{}", id)
    }
}

fn read_cpumask_from_bytes(src: &[u8]) -> Result<Cpumask> {
    let mut cpumask = Cpumask::new();

    for (byte_idx, byte) in src.iter().copied().enumerate() {
        let mut bits = byte;
        while bits != 0 {
            let bit = bits.trailing_zeros() as usize;
            bits &= !(1u8 << bit);
            cpumask.set_cpu((byte_idx * 8) + bit)?;
        }
    }

    Ok(cpumask)
}

fn read_bpf_cell(skel: &BpfSkel, cell_id: u32) -> Result<bpf_intf::cell> {
    let raw = skel
        .maps
        .cells
        .lookup(&cell_id.to_ne_bytes(), libbpf_rs::MapFlags::ANY)
        .with_context(|| format!("Failed to lookup cell {}", cell_id))?
        .ok_or_else(|| anyhow::anyhow!("Cell {} missing from BPF map", cell_id))?;

    if raw.len() != std::mem::size_of::<bpf_intf::cell>() {
        bail!(
            "Unexpected cell {} value size: {} != {}",
            cell_id,
            raw.len(),
            std::mem::size_of::<bpf_intf::cell>()
        );
    }

    Ok(unsafe { std::ptr::read_unaligned(raw.as_ptr() as *const bpf_intf::cell) })
}

fn read_cpu_ctxs(skel: &BpfSkel) -> Result<Vec<bpf_intf::cpu_ctx>> {
    let mut cpu_ctxs = vec![];
    let cpu_ctxs_vec = skel
        .maps
        .cpu_ctxs
        .lookup_percpu(&0u32.to_ne_bytes(), libbpf_rs::MapFlags::ANY)
        .context("Failed to lookup cpu_ctx")?
        .expect("BUG: cpu_ctxs lookup_percpu returned None for key 0");
    if cpu_ctxs_vec.len() < *NR_CPUS_POSSIBLE {
        bail!(
            "Percpu map returned {} entries but expected {}",
            cpu_ctxs_vec.len(),
            *NR_CPUS_POSSIBLE
        );
    }
    for cpu in 0..*NR_CPUS_POSSIBLE {
        cpu_ctxs.push(*unsafe {
            &*(cpu_ctxs_vec[cpu].as_slice().as_ptr() as *const bpf_intf::cpu_ctx)
        });
    }
    Ok(cpu_ctxs)
}

/// Read the per-CPU accounting slots of `active_subcells` from BPF.
fn read_subcell_accounts<I>(
    skel: &BpfSkel,
    active_subcells: I,
) -> Result<HashMap<(usize, usize), Vec<bpf_intf::subcell_account>>>
where
    I: IntoIterator<Item = (usize, usize)>,
{
    let mut accounts_by_subcell = HashMap::new();

    for (cell, subcell) in active_subcells {
        if cell >= MAX_CELLS || subcell >= MAX_SUBCELLS_PER_CELL {
            bail!(
                "Invalid subcell account key {}:{} (max {}:{})",
                cell,
                subcell,
                MAX_CELLS,
                MAX_SUBCELLS_PER_CELL
            );
        }

        let packed = ((cell * MAX_SUBCELLS_PER_CELL) + subcell) as u32;
        let per_cpu_values = skel
            .maps
            .subcell_accounts
            .lookup_percpu(&packed.to_ne_bytes(), libbpf_rs::MapFlags::ANY)
            .with_context(|| format!("Failed to lookup subcell_accounts key {}", packed))?
            .with_context(|| format!("subcell_accounts key {} is missing", packed))?;

        if per_cpu_values.len() < *NR_CPUS_POSSIBLE {
            bail!(
                "subcell_accounts returned {} entries but expected {}",
                per_cpu_values.len(),
                *NR_CPUS_POSSIBLE
            );
        }

        let mut values = Vec::with_capacity(*NR_CPUS_POSSIBLE);
        for cpu in 0..*NR_CPUS_POSSIBLE {
            let value = per_cpu_values[cpu].as_slice();
            if value.len() < std::mem::size_of::<bpf_intf::subcell_account>() {
                bail!(
                    "subcell_accounts key {} cpu {} value is too small: {}",
                    packed,
                    cpu,
                    value.len()
                );
            }
            values.push(unsafe {
                std::ptr::read_unaligned(value.as_ptr() as *const bpf_intf::subcell_account)
            });
        }

        accounts_by_subcell.insert((cell, subcell), values);
    }

    Ok(accounts_by_subcell)
}

fn run(opts: Opts) -> Result<()> {
    if opts.version {
        println!(
            "scx_mitosis {}",
            build_id::full_version(env!("CARGO_PKG_VERSION"))
        );
        return Ok(());
    }

    let env_filter = EnvFilter::try_from_default_env()
        .or_else(|_| match EnvFilter::try_new(&opts.log_level) {
            Ok(filter) => Ok(filter),
            Err(e) => {
                eprintln!(
                    "invalid log envvar: {}, using info, err is: {}",
                    opts.log_level, e
                );
                EnvFilter::try_new("info")
            }
        })
        .unwrap_or_else(|_| EnvFilter::new("info"));

    match tracing_subscriber::fmt()
        .with_env_filter(env_filter)
        .with_target(true)
        .with_thread_ids(true)
        .with_file(true)
        .with_line_number(true)
        .try_init()
    {
        Ok(()) => {}
        Err(e) => eprintln!("failed to init logger: {}", e),
    }

    if opts.verbose > 0 {
        warn!("Setting verbose via -v is deprecated and will be an error in future releases.");
    }

    debug!("opts={:?}", &opts);

    if let Some(run_id) = opts.run_id {
        info!("scx_mitosis run_id: {}", run_id);
    }

    let shutdown = Arc::new(AtomicBool::new(false));
    let shutdown_clone = shutdown.clone();
    ctrlc::set_handler(move || {
        shutdown_clone.store(true, Ordering::Relaxed);
    })
    .context("Error setting Ctrl-C handler")?;

    if let Some(intv) = opts.monitor.or(opts.monitor_sync) {
        let shutdown_clone = shutdown.clone();
        let jh = std::thread::spawn(move || {
            match stats::monitor(Duration::from_secs_f64(intv), shutdown_clone) {
                Ok(_) => {
                    debug!("stats monitor thread finished successfully")
                }
                Err(error_object) => {
                    warn!(
                        "stats monitor thread finished because of an error {}",
                        error_object
                    )
                }
            }
        });
        if opts.monitor.is_some() {
            let _ = jh.join();
            return Ok(());
        }
    }

    let mut open_object = MaybeUninit::uninit();
    loop {
        let mut sched =
            Scheduler::init(&opts, &mut open_object).context("initializing scheduler")?;
        if !sched
            .run(shutdown.clone())
            .context("running scheduler main loop")?
            .should_restart()
        {
            break;
        }
    }

    Ok(())
}

fn main() -> Result<()> {
    let parsed = undefok_flags::parse_args::<Opts>()?;
    // Emit undefok notices before logger setup so ignored rollout flags are always visible.
    for undefok in &parsed.ignored_undefok_flags {
        eprintln!("warning: ignoring undefok flag --{}", undefok.long);
    }
    run(parsed.opts)
}

#[cfg(test)]
mod tests {
    use super::{CpuAssignment, Cpumask, Opts, validate_subcell_assignments};
    use clap::Parser;

    fn cpumask(cpus: &[usize]) -> Cpumask {
        let mut mask = Cpumask::new();
        for &cpu in cpus {
            mask.set_cpu(cpu).unwrap();
        }
        mask
    }

    fn assignment(id: u32, primary: &[usize], borrowable: Option<&[usize]>) -> CpuAssignment {
        CpuAssignment {
            id,
            primary: cpumask(primary),
            borrowable: borrowable.map(cpumask),
        }
    }

    #[test]
    fn requires_cell_parent_cgroup_for_scheduler_mode() {
        assert!(Opts::try_parse_from(["scx_mitosis"]).is_err());
    }

    #[test]
    fn allows_cell_config_for_scheduler_mode() {
        assert!(Opts::try_parse_from(["scx_mitosis", "--cell-config", "/tmp/cells.json"]).is_ok());
    }

    #[test]
    fn rejects_both_cell_modes() {
        assert!(
            Opts::try_parse_from([
                "scx_mitosis",
                "--cell-parent-cgroup",
                "/workloads",
                "--cell-config",
                "/tmp/cells.json"
            ])
            .is_err()
        );
    }

    #[test]
    fn allows_monitor_without_cell_parent_cgroup() {
        assert!(Opts::try_parse_from(["scx_mitosis", "--monitor", "1"]).is_ok());
    }

    #[test]
    fn allows_version_without_cell_parent_cgroup() {
        assert!(Opts::try_parse_from(["scx_mitosis", "--version"]).is_ok());
    }

    #[test]
    fn accepts_complete_unpinned_subcell_partition() {
        let cell = assignment(3, &[0, 1, 2, 3], None);
        let subcells = [
            assignment(0, &[0, 1], Some(&[2, 3])),
            assignment(1, &[2, 3], Some(&[0, 1])),
        ];

        validate_subcell_assignments(&cell, &subcells, true).unwrap();
    }

    #[test]
    fn rejects_overlapping_subcell_primaries() {
        let cell = assignment(3, &[0, 1, 2, 3], None);
        let subcells = [
            assignment(0, &[0, 1, 2], None),
            assignment(1, &[2, 3], None),
        ];

        assert!(validate_subcell_assignments(&cell, &subcells, false).is_err());
    }

    #[test]
    fn rejects_incomplete_subcell_partition() {
        let cell = assignment(3, &[0, 1, 2, 3], None);
        let subcells = [assignment(0, &[0, 1, 2], None)];

        assert!(validate_subcell_assignments(&cell, &subcells, false).is_err());
    }
}