ktstr 0.23.0

//! Basic and mixed-workload scenario implementations.

use super::Ctx;
use super::ops::{CgroupDef, Op, Step, execute_steps};
use crate::assert::AssertResult;
use crate::workload::*;
use anyhow::Result;
use std::time::Duration;

fn host_cgroup_contention_steps(ctx: &Ctx) -> Vec<Step> {
    vec![
        Step::with_defs(
            vec![CgroupDef::named("cg_0"), CgroupDef::named("cg_1")],
            ctx.settled_hold(1.0),
        )
        .set_ops(vec![Op::spawn_host(
            WorkSpec::default().workers(ctx.topo.total_cpus()),
        )]),
    ]
}

/// Two managed cgroups with host-level contention from workers in the
/// test runner's own cgroup (`SpawnPlacement::RunnerCgroup`, typically
/// the root cgroup). Spawns `total_cpus` workers outside any managed cgroup
/// alongside two default cgroups.
pub fn custom_host_cgroup_contention(ctx: &Ctx) -> Result<AssertResult> {
    execute_steps(ctx, host_cgroup_contention_steps(ctx))
}

fn sched_mixed_steps(ctx: &Ctx) -> Vec<Step> {
    let configs = [
        (SchedPolicy::Normal, WorkType::SpinWait),
        (SchedPolicy::Batch, WorkType::SpinWait),
        (SchedPolicy::Idle, WorkType::SpinWait),
        (
            SchedPolicy::Fifo(1),
            WorkType::bursty(Duration::from_millis(500), Duration::from_millis(250)),
        ),
    ];

    let mut ops = vec![Op::add_cgroup("cg_0"), Op::add_cgroup("cg_1")];
    for name in ["cg_0", "cg_1"] {
        for &(policy, ref wtype) in &configs {
            ops.push(Op::spawn_workers(
                name,
                WorkSpec::default()
                    .workers(2)
                    .sched_policy(policy)
                    .work_type(wtype.clone()),
            ));
        }
    }

    vec![Step::new(ops, ctx.settled_hold(1.0))]
}

/// Two cgroups each running Normal, Batch, Idle, and FIFO(1) workers
/// concurrently. FIFO workers use bursty workloads to avoid monopolizing
/// CPUs.
pub fn custom_sched_mixed(ctx: &Ctx) -> Result<AssertResult> {
    execute_steps(ctx, sched_mixed_steps(ctx))
}

/// Light workload for scheduler-DEATH tests: one cgroup of 4 `YieldHeavy`
/// workers held for the scenario duration.
///
/// Deliberately light — few continuously-runnable tasks — so that when the
/// sched_ext scheduler dies mid-run, the kernel's scheduler-disable
/// (`scx_ops_disable`) bypass-drain has trivially few runnable tasks to
/// migrate. Both death paths funnel through the same drain: an `scx_bpf_error`
/// crash (host-injected via `bpf_map_write`, fires on any `ktstr_dispatch`
/// call) AND a watchdog stall (`--stall-after`, which the kernel's
/// per-rq sched_ext watchdog fires: `scx_watchdog_workfn`'s delayed
/// workqueue runs `check_rq_for_timeouts` (kernel/sched/ext.c), which
/// calls `scx_exit(..., SCX_EXIT_ERROR_STALL, ...)` once a runnable
/// task exceeds `watchdog_timeout`). The scheduler only stops
/// dispatching under `--stall-after` (`if (stall) return;`); it runs
/// no timer of its own. A heavy runnable
/// workload (e.g. [`custom_sched_mixed`]'s 12 SpinWait tasks) triggers the
/// Linux 6.14 per-node global-DSQ bypass-drain livelock, which strands the
/// whole guest until the host watchdog fires; that livelock is fixed upstream
/// only in sched_ext/for-7.1 (per-CPU bypass DSQs + an interruptible aborting
/// consume loop), so ktstr cannot patch it and instead keeps the death-test
/// workload light. The death still fires, and the death -> scx exit ->
/// auto-repro -> wprof-capture path is exercised identically, so no coverage
/// is lost.
pub fn custom_crash_light(ctx: &Ctx) -> Result<AssertResult> {
    let ops = vec![
        Op::add_cgroup("cg_0"),
        Op::spawn_workers(
            "cg_0",
            WorkSpec::default()
                .workers(4)
                .work_type(WorkType::YieldHeavy),
        ),
    ];
    execute_steps(ctx, vec![Step::new(ops, ctx.settled_hold(1.0))])
}

fn cgroup_pipe_io_steps(ctx: &Ctx) -> Vec<Step> {
    let mut ops = vec![Op::add_cgroup("cg_0"), Op::add_cgroup("cg_1")];
    for name in ["cg_0", "cg_1"] {
        ops.push(Op::spawn_workers(
            name,
            WorkSpec::default()
                .workers(2)
                .work_type(WorkType::pipe_io(1024)),
        ));
        ops.push(Op::spawn_workers(
            name,
            WorkSpec::default().workers(ctx.workers_per_cgroup),
        ));
    }

    vec![Step::new(ops, ctx.settled_hold(1.0))]
}

/// Two cgroups each with paired PipeIo workers and SpinWait workers.
/// Exercises cross-CPU wake placement from pipe I/O under CPU load.
pub fn custom_cgroup_pipe_io(ctx: &Ctx) -> Result<AssertResult> {
    execute_steps(ctx, cgroup_pipe_io_steps(ctx))
}

#[cfg(test)]
mod tests {
    use super::super::ops::{Setup, SpawnPlacement};
    use super::*;
    use crate::cgroup::CgroupManager;
    use crate::topology::TestTopology;
    use std::time::Duration;

    fn ctx_for_test<'a>(cgroups: &'a CgroupManager, topo: &'a TestTopology) -> Ctx<'a> {
        Ctx {
            cgroups,
            topo,
            duration: Duration::from_secs(1),
            workers_per_cgroup: 3,
            sched_pid: Some(1),
            settle: Duration::from_millis(100),
            work_type_override: None,
            assert: crate::assert::Assert::default_checks(),
            wait_for_map_write: false,
            current_step: std::sync::Arc::new(std::sync::atomic::AtomicU16::new(0)),
            entry_name: None,
            variant_hash: 0,
        }
    }

    fn def_names(step: &Step) -> Vec<String> {
        match &step.setup {
            Setup::Defs(defs) => defs.iter().map(|d| d.name.to_string()).collect(),
            Setup::Factory(_) => Vec::new(),
        }
    }

    #[test]
    fn host_cgroup_contention_builds_two_defs_and_host_spawn() {
        let cgroups = CgroupManager::new("/nonexistent");
        let topo = TestTopology::from_vm_topology(&crate::vmm::topology::Topology::new(1, 1, 4, 1));
        let ctx = ctx_for_test(&cgroups, &topo);

        let steps = host_cgroup_contention_steps(&ctx);
        assert_eq!(steps.len(), 1);
        assert_eq!(def_names(&steps[0]), ["cg_0", "cg_1"]);
        assert_eq!(steps[0].ops.len(), 1);
        match &steps[0].ops[0] {
            Op::Spawn {
                placement: SpawnPlacement::RunnerCgroup,
                work,
            } => {
                assert_eq!(work.num_workers, Some(topo.total_cpus()));
            }
            other => panic!("expected Op::Spawn(RunnerCgroup), got {other:?}"),
        }
    }

    #[test]
    fn sched_mixed_builds_two_add_cgroups_and_eight_spawns() {
        let cgroups = CgroupManager::new("/nonexistent");
        let topo = TestTopology::from_vm_topology(&crate::vmm::topology::Topology::new(1, 1, 4, 1));
        let ctx = ctx_for_test(&cgroups, &topo);

        let steps = sched_mixed_steps(&ctx);
        assert_eq!(steps.len(), 1);
        let ops = &steps[0].ops;
        let adds = ops
            .iter()
            .filter(|o| matches!(o, Op::AddCgroup { .. }))
            .count();
        let spawns = ops
            .iter()
            .filter(|o| {
                matches!(
                    o,
                    Op::Spawn {
                        placement: SpawnPlacement::Cgroup(_),
                        ..
                    }
                )
            })
            .count();
        assert_eq!(adds, 2, "two cgroups added");
        assert_eq!(spawns, 8, "4 policies × 2 cgroups = 8 spawns");
        for op in ops {
            if let Op::Spawn {
                placement: SpawnPlacement::Cgroup(_),
                work,
            } = op
            {
                assert_eq!(work.num_workers, Some(2));
            }
        }
    }

    #[test]
    fn cgroup_pipe_io_spawn_counts_follow_workers_per_cgroup() {
        let cgroups = CgroupManager::new("/nonexistent");
        let topo = TestTopology::from_vm_topology(&crate::vmm::topology::Topology::new(1, 1, 4, 1));
        let ctx = ctx_for_test(&cgroups, &topo);

        let steps = cgroup_pipe_io_steps(&ctx);
        assert_eq!(steps.len(), 1);
        let ops = &steps[0].ops;
        let spawns: Vec<_> = ops
            .iter()
            .filter_map(|o| match o {
                Op::Spawn {
                    placement: SpawnPlacement::Cgroup(cgroup),
                    work,
                } => Some((cgroup.to_string(), work.num_workers)),
                _ => None,
            })
            .collect();
        assert_eq!(spawns.len(), 4, "pipe_io spawn + spinwait spawn per cgroup");
        let spinwait_workers: Vec<_> = spawns
            .iter()
            .filter(|(_, n)| *n == Some(ctx.workers_per_cgroup))
            .collect();
        assert_eq!(spinwait_workers.len(), 2);
        let pipe_workers: Vec<_> = spawns.iter().filter(|(_, n)| *n == Some(2)).collect();
        assert_eq!(pipe_workers.len(), 2);
    }
}