// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2026 Galih Tama <galpt@v.recipes>
//
// This software may be used and distributed according to the terms of the GNU
// General Public License version 2.

//! scx_mlfq, a Multilevel Feedback Queue scheduler for sched_ext.
//!
//! Per-CPU, virtual-time-ordered user DSQs (Q1/Q2/Q3 per CPU) over an EEVDF
//! virtual-time substrate. Tasks are classified into queues by an EMA
//! interactivity gauge that climbs during run segments and decays at
//! wakeup by the half-life step, with the short-sleep and I/O boost and
//! the band hysteresis. Demotion flows through the run-out gate. See
//! README.md for the design overview.

mod bpf_skel;
pub use bpf_skel::*;
pub mod bpf_intf;
pub use bpf_intf::*;

mod config;
mod stats;
mod topology;
mod webui;

use std::mem::size_of;
use std::mem::MaybeUninit;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use clap::CommandFactory;
use clap::Parser;
use clap_complete::generate;
use clap_complete::Shell;
use crossbeam::channel::RecvTimeoutError;
use libbpf_rs::MapCore;
use log::info;
use scx_stats::prelude::*;
use scx_utils::build_id;
use scx_utils::compat;
use scx_utils::libbpf_clap_opts::LibbpfOpts;
use scx_utils::pm;
use scx_utils::scx_ops_attach;
use scx_utils::scx_ops_load;
use scx_utils::scx_ops_open;
use scx_utils::try_set_rlimit_infinity;
use scx_utils::uei_exited;
use scx_utils::uei_report;
use scx_utils::UserExitInfo;

use config::Config;
use stats::Metrics;

const SCHEDULER_NAME: &str = "scx_mlfq";

/*
 * PM QoS idle-resume-latency cap in microseconds, applied to
 * /dev/cpu_dma_latency for the duration of the run. 10 us bans the deep
 * core and package C-states (18 us and 350 us exits on the target) while
 * keeping C1 (1 us), so wakeup latency is not dominated by deep-state
 * exits.
 */
const MLFQ_IDLE_RESUME_LATENCY_US: i32 = 10;

/* Compile-time bounds from src/bpf/intf.h, used by the web metrics. */
const MLFQ_MAX_CPUS: usize = crate::bpf_intf::mlfq_consts_MLFQ_MAX_CPUS as usize;

fn full_version() -> String {
    build_id::full_version(env!("CARGO_PKG_VERSION"))
}

#[derive(Debug, Parser)]
#[command(name = SCHEDULER_NAME, version, disable_version_flag = true)]
struct Opts {
    /// Enable periodic statistics monitoring at the given interval.
    #[clap(long)]
    stats: Option<f64>,

    /// Run in statistics monitoring mode. The scheduler is not launched.
    #[clap(long)]
    monitor: Option<f64>,

    /// Enable verbose libbpf/BPF debug logging.
    #[clap(short = 'd', long, action = clap::ArgAction::SetTrue)]
    debug: bool,

    /// Enable verbose output, including libbpf details.
    #[clap(short = 'v', long, action = clap::ArgAction::SetTrue)]
    verbose: bool,

    /// Size of the exit dump buffer in bytes. The kernel fills it with
    /// per-CPU and per-task state when the scheduler exits on an error.
    #[clap(long, default_value = "1048576")]
    exit_dump_len: u32,

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

    /// Show descriptions for statistics.
    #[clap(long)]
    help_stats: bool,

    /// Generate shell completions and exit.
    #[clap(long, value_name = "SHELL", hide = true)]
    completions: Option<Shell>,

    /// Disable the loopback web UI. The UI binds [::1]:50005 (falling
    /// back to 127.0.0.1:50005, then to the /tmp/scx_mlfq.sock unix
    /// socket when the loader sandbox blocks TCP) and is unauthenticated.
    #[clap(long = "no-webui", action = clap::ArgAction::SetTrue)]
    no_webui: bool,

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

/// The Scheduler facade owns the loaded skeleton, the struct_ops link, the
/// stats server and the web UI wiring. Drives the run loop until shutdown
/// or UEI exit.
struct Scheduler<'a> {
    skel: BpfSkel<'a>,
    struct_ops: Option<libbpf_rs::Link>,
    stats_server: StatsServer<(), Metrics>,
    /*
     * Web UI plumbing. The metrics sender (None when --no-webui, in
     * which case the webui thread is never spawned and the metrics are
     * never collected) and the once-per-attach per-CPU static seed
     * (freq, LLC, SMT) the web metrics merge the dynamic BPF state into.
     */
    webui_tx: Option<crossbeam::channel::Sender<stats::WebMetrics>>,
    webui_join: Option<std::thread::JoinHandle<()>>,
    cpu_static: Vec<stats::PerCpuMetrics>,
    /*
     * Per-CPU current-frequency cache for the web UI, refreshed from
     * sysfs at most once per second in the run loop.
     */
    cur_freq_khz: Vec<u64>,
    freq_read_at: Option<std::time::Instant>,
    started_at: std::time::Instant,
    /*
     * PM QoS idle-latency constraint on /dev/cpu_dma_latency, held for
     * the run. Closing the file restores the previous constraint.
     */
    #[expect(dead_code)]
    pm_qos_fd: Option<std::fs::File>,
}

impl<'a> Scheduler<'a> {
    fn init(
        opts: &'a Opts,
        open_object: &'a mut MaybeUninit<libbpf_rs::OpenObject>,
        shutdown: Arc<AtomicBool>,
    ) -> Result<Self> {
        try_set_rlimit_infinity();

        let mut skel_builder = BpfSkelBuilder::default();
        skel_builder.obj_builder.debug(opts.debug || opts.verbose);

        let open_opts = opts.libbpf.clone().into_bpf_open_opts();
        let mut skel = scx_ops_open!(skel_builder, open_object, mlfq_ops, open_opts)?;

        // Write the validated constants into rodata before load. The rodata
        // section becomes read-only once the object is loaded.
        let config = Config::default();
        config.validate()?;
        config.apply(&mut skel)?;
        info!("Config: {}", config.describe());

        // Hybrid-capacity, cache-domain and NUMA placement data also goes
        // into rodata pre-load.
        let topology_plan = topology::init_topology(&mut skel)?;

        /*
         * Ops flags. Honor exiting tasks, receive the SCX_ENQ_LAST enqueue
         * for the last runnable task on a CPU, never migrate
         * migration-disabled tasks, and allow queued-wakeup selection of
         * idle CPUs (the idle-CPU fast path depends on the latter two).
         */
        let mut flags = *compat::SCX_OPS_ENQ_EXITING
            | *compat::SCX_OPS_ENQ_LAST
            | *compat::SCX_OPS_ENQ_MIGRATION_DISABLED
            | *compat::SCX_OPS_ALLOW_QUEUED_WAKEUP;
        if *compat::SCX_OPS_KEEP_BUILTIN_IDLE != 0 {
            flags |= *compat::SCX_OPS_KEEP_BUILTIN_IDLE;
            skel.maps
                .rodata_data
                .as_mut()
                .expect("rodata missing, the BPF object has no .rodata section")
                .mlfq_idle_tracking = 1;
        } else {
            skel.struct_ops.mlfq_ops_mut().update_idle = std::ptr::null_mut();
        }
        skel.struct_ops.mlfq_ops_mut().flags = flags;

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

        let mut skel = scx_ops_load!(skel, mlfq_ops, uei)?;

        // The membership bitmaps are written after load (the maps are only
        // available on the loaded object).
        if let Err(e) = topology::write_primary_bitmap(&mut skel, &topology_plan.capacity) {
            log::warn!(
                "failed to write the primary bitmap, falling back to all-primary placement: {e:#}"
            );
        }
        if let Err(e) = topology::write_llc_bitmaps(&mut skel, &topology_plan.llcs) {
            log::warn!("failed to write the LLC bitmaps, disabling LLC-aware placement: {e:#}");
        }
        if let Err(e) = topology::write_llc_cpu_lists(&mut skel, &topology_plan.llcs) {
            log::warn!(
                "failed to write the per-LLC CPU lists, disabling LLC-aware stealing: {e:#}"
            );
        }

        let struct_ops = scx_ops_attach!(skel, mlfq_ops)?;

        /*
         * PM QoS. Hold a global idle-resume-latency constraint on
         * /dev/cpu_dma_latency for the duration of the run, so the
         * cpuidle governor keeps the CPUs in the shallowest idle states
         * that fit the cap and wakeup latency is not dominated by the
         * deep C-state exits.
         */
        let pm_qos_fd = if pm::cpu_idle_resume_latency_supported() {
            match pm::update_global_idle_resume_latency(MLFQ_IDLE_RESUME_LATENCY_US) {
                Ok(f) => {
                    info!(
                        "PM QoS idle resume latency held at {}us",
                        MLFQ_IDLE_RESUME_LATENCY_US
                    );
                    Some(f)
                }
                Err(e) => {
                    log::warn!("failed to set the PM QoS idle resume latency: {e:#}");
                    None
                }
            }
        } else {
            log::warn!(
                "PM QoS idle resume latency is not supported; the constraint is not applied"
            );
            None
        };

        let stats_server = StatsServer::new(stats::server_data()).launch()?;

        /*
         * The web UI is a small bounded metrics channel with capacity 16.
         * try_send drops a frame when the buffer is full, so the run
         * loop never blocks and the buffer never grows. It feeds a
         * detached server thread that exits on the shared shutdown
         * flag.
         */
        let (webui_tx, webui_join): (
            Option<crossbeam::channel::Sender<stats::WebMetrics>>,
            Option<std::thread::JoinHandle<()>>,
        ) = if opts.no_webui {
            (None, None)
        } else {
            let (tx, rx) = crossbeam::channel::bounded::<stats::WebMetrics>(16);
            let shutdown = shutdown.clone();
            let jh = std::thread::spawn(move || {
                webui::start(rx, shutdown);
            });
            (Some(tx), Some(jh))
        };
        let cpu_static = if opts.no_webui {
            Vec::new()
        } else {
            topology::web_cpu_static()
        };

        Ok(Self {
            skel,
            struct_ops: Some(struct_ops),
            stats_server,
            webui_tx,
            webui_join,
            cpu_static,
            cur_freq_khz: Vec::new(),
            freq_read_at: None,
            started_at: std::time::Instant::now(),
            pm_qos_fd,
        })
    }

    fn get_metrics(&self) -> Metrics {
        let s = self.read_mlfq_stats();
        Metrics {
            on_cpu: s[1],
            total_runtime: s[0],
            uptime_ns: self.started_at.elapsed().as_nanos() as u64,
            q1_placements: s[2],
            q2_placements: s[3],
            q3_placements: s[4],
            promotions: s[5],
            demotions: s[6],
            aging_boosts: s[7],
            short_sleep_boosts: s[8],
            preemption_kicks: s[9],
            cpuperf_boosts: s[10],
            steals: s[11],
            steals_same_llc: s[12],
            steals_cross_llc: s[13],
            keep_running: s[14],
            enq_no_tctx: s[15],
            enq_bad_weight: s[16],
            enq_no_deadline: s[17],
            enq_fastpath: s[18],
            enq_regular: s[19],
            enq_pinned_idle: s[20],
            enq_pinned_busy: s[21],
            enq_pinned_global: s[22],
            rt_takeovers: s[23],
            rt_evacuations: s[24],
            rt_redirects: s[25],
            rt_reenqs: s[26],
            op_lat: self.read_op_lat(),
            wakeup_total: self.read_wakeup_total(),
        }
    }

    /// Web-metrics snapshot. The raw scheduler counters plus the per-CPU
    /// state and the runnable gauges, pushed to the web UI every run-loop
    /// iteration.
    fn get_web_metrics(&mut self) -> stats::WebMetrics {
        let bss_data = self
            .skel
            .maps
            .bss_data
            .as_ref()
            .expect("bss_data missing, the BPF object has no .bss section");

        // Refresh the per-CPU current frequencies at most once per second.
        let now = std::time::Instant::now();
        if self
            .freq_read_at
            .is_none_or(|t| now.duration_since(t).as_secs() >= 1)
        {
            let nr = (bss_data.nr_cpu_ids as usize).min(MLFQ_MAX_CPUS);
            self.cur_freq_khz = (0..nr)
                .map(|cpu| topology::current_freq_khz(cpu as u32))
                .collect();
            self.freq_read_at = Some(now);
        }

        let nr_cpus = (bss_data.nr_cpu_ids as usize).min(MLFQ_MAX_CPUS);
        let mut statics = vec![None; nr_cpus];
        for s in &self.cpu_static {
            if (s.id as usize) < nr_cpus {
                statics[s.id as usize] = Some(s.clone());
            }
        }
        let mut per_cpu = Vec::with_capacity(nr_cpus);
        for (cpu, static_entry) in statics.iter().enumerate() {
            let mut entry = static_entry.clone().unwrap_or_default();
            entry.id = cpu as u32;
            entry.cur_freq_khz = self.cur_freq_khz.get(cpu).copied().unwrap_or(0);

            let state = self.read_cpu_state(cpu);
            entry.running_queue = state.running_queue;
            entry.running_pid = state.running_pid;

            let rt = self.read_rtdl_state(cpu);
            entry.rt_occupied = rt.flags & crate::bpf_intf::MLFQ_RTDL_OCCUPIED != 0;
            per_cpu.push(entry);
        }

        stats::WebMetrics {
            stats: self.get_metrics(),
            per_cpu,
            queue_runnable: self.read_queue_runnable(),
            llc_runnable: self.read_llc_runnable(),
        }
    }

    /// Read one CPU's dynamic state from the per-CPU array map. A failed
    /// lookup or an unexpected value size yields an all-zero state.
    fn read_cpu_state(&self, cpu: usize) -> mlfq_cpu_state {
        let key = (cpu as u32).to_ne_bytes();
        match self
            .skel
            .maps
            .cpu_state_stor
            .lookup(&key, libbpf_rs::MapFlags::ANY)
        {
            Ok(Some(bytes)) if bytes.len() >= size_of::<mlfq_cpu_state>() => {
                // SAFETY: mlfq_cpu_state is the repr(C) bindgen mirror of
                // the map's value type. Reading the value bytes as the
                // struct is a plain reinterpretation of integer fields.
                unsafe { std::ptr::read_unaligned(bytes.as_ptr().cast::<mlfq_cpu_state>()) }
            }
            _ => mlfq_cpu_state {
                running_queue: 0,
                running_pid: 0,
                steal_scan_off: 0,
                perf_level: 0,
                busy_win_ns: 0,
                busy_win_start: 0,
                running_deadline: 0,
                run_start_at: 0,
            },
        }
    }

    /// Read one CPU's realtime-occupancy state from the per-CPU array
    /// map. A failed lookup yields an all-zero state (not occupied).
    fn read_rtdl_state(&self, cpu: usize) -> mlfq_rtdl_state {
        let key = (cpu as u32).to_ne_bytes();
        match self
            .skel
            .maps
            .rtdl_state_stor
            .lookup(&key, libbpf_rs::MapFlags::ANY)
        {
            Ok(Some(bytes)) if bytes.len() >= size_of::<mlfq_rtdl_state>() => unsafe {
                std::ptr::read_unaligned(bytes.as_ptr().cast::<mlfq_rtdl_state>())
            },
            _ => mlfq_rtdl_state {
                flags: 0,
                pad: 0,
                last_drain_at: 0,
            },
        }
    }

    /// Tracked runnable tasks per queue (index 0 unused, 1..3 = Q1..Q3).
    fn read_queue_runnable(&self) -> Vec<u64> {
        let bss_data = self
            .skel
            .maps
            .bss_data
            .as_ref()
            .expect("bss_data missing, the BPF object has no .bss section");
        bss_data
            .mlfq_queue_runnable
            .iter()
            .map(|v| *v as u64)
            .collect()
    }

    /// Tracked runnable tasks per LLC domain (MLFQ_MAX_LLCS entries).
    fn read_llc_runnable(&self) -> Vec<u64> {
        let bss_data = self
            .skel
            .maps
            .bss_data
            .as_ref()
            .expect("bss_data missing, the BPF object has no .bss section");
        bss_data
            .mlfq_llc_runnable
            .iter()
            .map(|v| *v as u64)
            .collect()
    }

    fn exited(&self) -> bool {
        uei_exited!(&self.skel, uei)
    }

    /// Sum the per-CPU mlfq_stats slots into a flat u64 array in struct
    /// field order (see struct mlfq_stats in src/bpf/intf.h). Each CPU
    /// bumps its own slot non-atomically on its own cache line, so the
    /// sum is the system-wide counter.
    fn read_mlfq_stats(&self) -> [u64; 27] {
        let cpu_values = self
            .skel
            .maps
            .mlfq_stats
            .lookup_percpu(&0u32.to_ne_bytes(), libbpf_rs::MapFlags::ANY)
            .ok()
            .flatten()
            .unwrap_or_default();
        let mut total = [0u64; 27];
        for chunk in cpu_values {
            for (i, slot) in chunk.chunks_exact(8).take(27).enumerate() {
                total[i] = total[i].wrapping_add(u64::from_ne_bytes(
                    slot.try_into().expect("an 8-byte mlfq_stats slot"),
                ));
            }
        }
        total
    }

    /// Sum the per-CPU op-latency histogram into a flat per-op vector
    /// (MLFQ_OP_LAT_OPS x MLFQ_OP_LAT_BUCKETS entries, op-major).
    fn read_op_lat(&self) -> Vec<u64> {
        let nr_ops = crate::bpf_intf::mlfq_op_lat_slots_MLFQ_OP_LAT_OPS as usize;
        let buckets = crate::bpf_intf::mlfq_op_lat_consts_MLFQ_OP_LAT_BUCKETS as usize;
        let mut out = vec![0u64; nr_ops * buckets];

        for op in 0..nr_ops {
            let cpu_values = self
                .skel
                .maps
                .mlfq_op_lat
                .lookup_percpu(&(op as u32).to_ne_bytes(), libbpf_rs::MapFlags::ANY)
                .ok()
                .flatten()
                .unwrap_or_default();
            for chunk in cpu_values {
                // The per-CPU value leads with the sample gate (sampled_cnt)
                // before the bucket vector; skip it so the histogram reads land
                // on the bucket array.
                for (b, slot) in chunk.chunks_exact(8).skip(1).take(buckets).enumerate() {
                    out[op * buckets + b] = out[op * buckets + b].wrapping_add(u64::from_ne_bytes(
                        slot.try_into().expect("an 8-byte histogram slot"),
                    ));
                }
            }
        }
        out
    }

    /// Sum the per-CPU lifetime wakeup totals from the mlfq_wakeup_stats
    /// map. Each CPU's total is bumped atomically on the wakeup path.
    fn read_wakeup_total(&self) -> u64 {
        let cpu_values = self
            .skel
            .maps
            .mlfq_wakeup_stats
            .lookup_percpu(&0u32.to_ne_bytes(), libbpf_rs::MapFlags::ANY)
            .ok()
            .flatten()
            .unwrap_or_default();
        cpu_values
            .iter()
            .map(|chunk| {
                u64::from_ne_bytes(
                    chunk[0..8]
                        .try_into()
                        .expect("the total field is the first 8 bytes"),
                )
            })
            .fold(0u64, u64::wrapping_add)
    }

    fn run(&mut self, shutdown: Arc<AtomicBool>) -> Result<UserExitInfo> {
        let (res_ch, req_ch) = self.stats_server.channels();

        while !shutdown.load(Ordering::Relaxed) && !self.exited() {
            match req_ch.recv_timeout(Duration::from_millis(100)) {
                Ok(()) => {
                    // Push a web snapshot on the stats request too, so
                    // the UI cadence is the run-loop cadence, not the
                    // browser's 1 s poll.
                    let web = self.get_web_metrics();
                    if let Some(ref tx) = self.webui_tx {
                        let _ = tx.try_send(web);
                    }
                    res_ch.send(self.get_metrics())?
                }
                Err(RecvTimeoutError::Timeout) => {
                    let web = self.get_web_metrics();
                    if let Some(ref tx) = self.webui_tx {
                        let _ = tx.try_send(web);
                    }
                }
                Err(e) => Err(e)?,
            }
        }

        let m = self.get_metrics();
        log::info!(
            "mlfq exit counters: Q1={} Q2={} Q3={} fastpath={} regular={} \
             pin_idle={} pin_busy={} pin_global={} drop_tctx={} drop_weight={} \
             drop_deadline={} promotions={} demotions={} aging_boosts={} \
             short_sleep_boosts={} cpuperf_boosts={} preempt_kicks={} \
             runtime={} on_cpu={} steals={} steals_same_llc={} \
             steals_cross_llc={} keep_running={} rt_takeovers={} \
             rt_evacuations={} rt_redirects={} rt_reenqs={} wakeups={}",
            m.q1_placements,
            m.q2_placements,
            m.q3_placements,
            m.enq_fastpath,
            m.enq_regular,
            m.enq_pinned_idle,
            m.enq_pinned_busy,
            m.enq_pinned_global,
            m.enq_no_tctx,
            m.enq_bad_weight,
            m.enq_no_deadline,
            m.promotions,
            m.demotions,
            m.aging_boosts,
            m.short_sleep_boosts,
            m.cpuperf_boosts,
            m.preemption_kicks,
            m.total_runtime,
            m.on_cpu,
            m.steals,
            m.steals_same_llc,
            m.steals_cross_llc,
            m.keep_running,
            m.rt_takeovers,
            m.rt_evacuations,
            m.rt_redirects,
            m.rt_reenqs,
            m.wakeup_total
        );
        let _ = self.struct_ops.take();
        uei_report!(&self.skel, uei)
    }
}

impl Drop for Scheduler<'_> {
    fn drop(&mut self) {
        /*
         * Dropping pm_qos_fd closes the /dev/cpu_dma_latency fd, which
         * makes the kernel drop the PM QoS request and restore the
         * previous idle-latency constraint. The field drop below does
         * this on every exit path.
         */
        /*
         * Join the web UI thread so its unblock-write flag (see
         * webui.rs) is visible before the caller's restore decision. The
         * thread exits within its poll interval of the shutdown
         * flag, so the join is bounded.
         */
        if let Some(jh) = self.webui_join.take() {
            let _ = jh.join();
        }
        info!("Unregister {SCHEDULER_NAME} scheduler");
    }
}

fn main() -> Result<()> {
    let opts = Opts::parse();

    if let Some(shell) = opts.completions {
        generate(
            shell,
            &mut Opts::command(),
            SCHEDULER_NAME,
            &mut std::io::stdout(),
        );
        return Ok(());
    }

    let monitor_only = opts.monitor.is_some();

    // Print the version before any other work, so `scx_mlfq -V` works
    // without attaching anything.
    if opts.version {
        println!("{} {}", SCHEDULER_NAME, full_version());
        return Ok(());
    }

    if opts.help_stats {
        stats::server_data().describe_meta(&mut std::io::stdout(), None)?;
        return Ok(());
    }

    if !monitor_only {
        simplelog::TermLogger::init(
            if opts.debug {
                simplelog::LevelFilter::Debug
            } else {
                simplelog::LevelFilter::Info
            },
            simplelog::Config::default(),
            simplelog::TerminalMode::Stderr,
            simplelog::ColorChoice::Auto,
        )?;

        // The SMT annotation is appended when the host exposes it.
        let smt_suffix = match topology::smt_enabled() {
            Some(true) => " SMT on",
            Some(false) => " SMT off",
            None => "",
        };
        info!("{} {}{}", SCHEDULER_NAME, full_version(), smt_suffix);
        info!(
            "scheduler options: {}",
            std::env::args().skip(1).collect::<Vec<_>>().join(" ")
        );
    }

    let shutdown = Arc::new(AtomicBool::new(false));
    let shutdown_clone = shutdown.clone();

    ctrlc::set_handler(move || {
        shutdown_clone.store(true, Ordering::Relaxed);
    })?;

    if let Some(intv) = opts.monitor.or(opts.stats) {
        let monitor_shutdown = shutdown.clone();
        let jh = std::thread::spawn(move || {
            if let Err(err) = stats::monitor(Duration::from_secs_f64(intv), monitor_shutdown) {
                log::warn!("stats monitor thread finished with error: {err}");
            }
        });

        if monitor_only {
            let _ = jh.join();
            return Ok(());
        }
    }

    let mut open_object = MaybeUninit::<libbpf_rs::OpenObject>::uninit();
    loop {
        let mut sched = Scheduler::init(&opts, &mut open_object, shutdown.clone())?;
        if !sched.run(shutdown.clone())?.should_restart() {
            break;
        }
        // Give the kernel time to finish the previous detachment before
        // re-initializing the BPF object for the next incarnation.
        std::thread::sleep(Duration::from_millis(100));
    }

    /*
     * If this run wrote the runtime loader-sandbox unblock (see
     * webui.rs), restore it now. The drop-in is per-boot state under
     * /run, so it must be undone once, at the final exit after the
     * restart loop, not on every internal restart.
     */
    webui::restore_loader_sandbox();

    info!("Scheduler exited");

    Ok(())
}

#[cfg(test)]
mod math_test {
    use std::process::Command;

    fn host_cc() -> String {
        if let Ok(cc) = std::env::var("CC") {
            return cc;
        }
        // CI ships clang-19. Prefer it, then fall back to the system compiler.
        for cand in ["clang", "cc", "gcc"] {
            if Command::new(cand).arg("--version").status().is_ok() {
                return cand.to_string();
            }
        }
        "cc".to_string()
    }

    #[test]
    fn mlfq_pure_math_native_harness() {
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        let src = format!("{manifest_dir}/src/bpf/mlfq_math_test.c");
        let intf_dir = format!("{manifest_dir}/src/bpf");
        let out_dir = std::env::temp_dir().join("scx_mlfq_math_test");
        std::fs::create_dir_all(&out_dir).unwrap();
        let exe = out_dir.join("mlfq_math_test");

        let status = Command::new(host_cc())
            .args(["-O2", "-Wall", "-Werror", "-std=c11"])
            .args(["-I", &intf_dir])
            .arg(&src)
            .args(["-o", exe.to_str().unwrap()])
            .status()
            .expect("failed to compile mlfq_math_test.c");

        assert!(status.success(), "native harness failed to compile");

        let output = Command::new(&exe)
            .output()
            .expect("failed to run the native harness");

        let stdout = String::from_utf8_lossy(&output.stdout);
        print!("{stdout}");

        assert!(
            output.status.success(),
            "native harness failed:\n{stdout}{}",
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(
            stdout.contains("All tests passed"),
            "native harness did not report success"
        );
    }
}
