Skip to content

scx_p2dq: Refactor arena initialization #2009

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 28, 2025
Merged

scx_p2dq: Refactor arena initialization #2009

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 96 additions & 20 deletions scheds/rust/scx_chaos/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ mod bpf_skel;

use bpf_skel::BpfSkel;

use scx_p2dq::P2dqArenaProgs;
use scx_p2dq::SchedulerOpts as P2dqOpts;
use scx_userspace_arena::alloc::Allocator;
use scx_userspace_arena::alloc::HeapAllocator;
Expand All @@ -17,14 +16,16 @@ use scx_utils::scx_ops_load;
use scx_utils::scx_ops_open;
use scx_utils::uei_exited;
use scx_utils::uei_report;
use scx_utils::Core;
use scx_utils::Llc;
use scx_utils::Topology;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use clap::Parser;
use libbpf_rs::OpenObject;
use libbpf_rs::ProgramInput;
use libbpf_rs::ProgramOutput;
use log::debug;
use log::info;
use nix::unistd::Pid;
Expand Down Expand Up @@ -74,24 +75,6 @@ unsafe impl Allocator for ArenaAllocator {
}
}

impl P2dqArenaProgs for BpfSkel<'_> {
fn run_arena_init<'b>(&self, input: ProgramInput<'b>) -> Result<ProgramOutput<'b>> {
Ok(self.progs.p2dq_arena_init.test_run(input)?)
}

fn run_alloc_mask<'b>(&self, input: ProgramInput<'b>) -> Result<ProgramOutput<'b>> {
Ok(self.progs.p2dq_alloc_mask.test_run(input)?)
}

fn run_topology_node_init<'b>(&self, input: ProgramInput<'b>) -> Result<ProgramOutput<'b>> {
Ok(self.progs.p2dq_topology_node_init.test_run(input)?)
}

fn setup_ptr(&self) -> u64 {
self.maps.bss_data.setup_ptr
}
}

#[derive(Debug)]
pub enum Trait {
RandomDelays {
Expand Down Expand Up @@ -194,6 +177,96 @@ impl Scheduler {
}

impl Builder<'_> {
fn setup_arenas(&self, skel: &mut BpfSkel) -> Result<()> {
// Allocate the arena memory from the BPF side so userspace initializes it before starting
// the scheduler. Despite the function call's name this is neither a test nor a test run,
// it's the recommended way of executing SEC("syscall") probes.
let input = ProgramInput {
..Default::default()
};

let output = skel.progs.p2dq_arena_init.test_run(input)?;
if output.return_value != 0 {
bail!(
"Could not initialize arenas, p2dq_setup returned {}",
output.return_value as i32
);
}

Ok(())
}

fn setup_topology_node(&self, skel: &mut BpfSkel, mask: &[u64]) -> Result<()> {
// Copy the address of ptr to the kernel to populate it from BPF with the arena pointer.
let input = ProgramInput {
..Default::default()
};

let output = skel.progs.p2dq_alloc_mask.test_run(input)?;
if output.return_value != 0 {
bail!(
"Could not initialize arenas, setup_topology_node returned {}",
output.return_value as i32
);
}

let ptr =
unsafe { std::mem::transmute::<u64, &mut [u64; 10]>(skel.maps.bss_data.setup_ptr) };

let (valid_mask, _) = ptr.split_at_mut(mask.len());
valid_mask.clone_from_slice(mask);

let input = ProgramInput {
..Default::default()
};
let output = skel.progs.p2dq_topology_node_init.test_run(input)?;
if output.return_value != 0 {
bail!(
"p2dq_topology_node_init returned {}",
output.return_value as i32
);
}

Ok(())
}

fn setup_topology(&self, skel: &mut BpfSkel) -> Result<()> {
let topo = Topology::new().expect("Failed to build host topology");

self.setup_topology_node(skel, topo.span.as_raw_slice())?;

for (_, node) in topo.nodes {
self.setup_topology_node(skel, node.span.as_raw_slice())?;
}

for (_, llc) in topo.all_llcs {
self.setup_topology_node(
skel,
Arc::<Llc>::into_inner(llc)
.expect("missing llc")
.span
.as_raw_slice(),
)?;
}

for (_, core) in topo.all_cores {
self.setup_topology_node(
skel,
Arc::<Core>::into_inner(core)
.expect("missing core")
.span
.as_raw_slice(),
)?;
}
for (_, cpu) in topo.all_cpus {
let mut mask = [0; 9];
mask[cpu.id.checked_shr(64).unwrap_or(0)] |= 1 << (cpu.id % 64);
self.setup_topology_node(skel, &mask)?;
}

Ok(())
}

fn load_skel(&self) -> Result<Pin<Rc<SkelWithObject>>> {
let mut out: Rc<MaybeUninit<SkelWithObject>> = Rc::new_uninit();
let uninit_skel = Rc::get_mut(&mut out).expect("brand new rc should be unique");
Expand Down Expand Up @@ -306,6 +379,9 @@ impl Builder<'_> {
let mut skel = scx_ops_load!(open_skel, chaos, uei)?;
scx_p2dq::init_skel!(&mut skel);

self.setup_arenas(&mut skel)?;
self.setup_topology(&mut skel)?;

let out = unsafe {
// SAFETY: initialising field by field. open_object is already "initialised" (it's
// permanently MaybeUninit so any state is fine), hence the structure will be
Expand Down
113 changes: 0 additions & 113 deletions scheds/rust/scx_p2dq/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,8 @@
pub use scx_utils::CoreType;
use scx_utils::Topology;
pub use scx_utils::NR_CPU_IDS;
use scx_utils::{Core, Llc};

use anyhow::{bail, Result};
use clap::Parser;
use libbpf_rs::ProgramInput;
use libbpf_rs::ProgramOutput;

use std::sync::Arc;

lazy_static::lazy_static! {
pub static ref TOPO: Topology = Topology::new().unwrap();
Expand Down Expand Up @@ -145,110 +139,6 @@ pub fn dsq_slice_ns(dsq_index: u64, min_slice_us: u64, dsq_shift: u64) -> u64 {
result
}

/// Trait for interfacing with BPF arena programs
pub trait P2dqArenaProgs {
/// Run the arena initialization program and return the result
fn run_arena_init<'a>(&self, input: ProgramInput<'a>) -> Result<ProgramOutput<'a>>;

/// Run the allocation mask program and return the result
fn run_alloc_mask<'a>(&self, input: ProgramInput<'a>) -> Result<ProgramOutput<'a>>;

/// Run the topology node initialization program and return the result
fn run_topology_node_init<'a>(&self, input: ProgramInput<'a>) -> Result<ProgramOutput<'a>>;

/// Access to the setup pointer in BSS data
fn setup_ptr(&self) -> u64;
}

pub fn setup_arenas<T: P2dqArenaProgs>(skel: &T) -> Result<()> {
// Allocate the arena memory from the BPF side so userspace initializes it before starting
// the scheduler. Despite the function call's name this is neither a test nor a test run,
// it's the recommended way of executing SEC("syscall") probes.
let input = ProgramInput {
..Default::default()
};

let output = skel.run_arena_init(input)?;
if output.return_value != 0 {
bail!(
"Could not initialize arenas, p2dq_setup returned {}",
output.return_value as i32
);
}

Ok(())
}

fn setup_topology_node<T: P2dqArenaProgs>(skel: &T, mask: &[u64]) -> Result<()> {
// Copy the address of ptr to the kernel to populate it from BPF with the arena pointer.
let input = ProgramInput {
..Default::default()
};

let output = skel.run_alloc_mask(input)?;
if output.return_value != 0 {
bail!(
"Could not initialize arenas, setup_topology_node returned {}",
output.return_value as i32
);
}

let ptr = unsafe { std::mem::transmute::<u64, &mut [u64; 10]>(skel.setup_ptr()) };

let (valid_mask, _) = ptr.split_at_mut(mask.len());
valid_mask.clone_from_slice(mask);

let input = ProgramInput {
..Default::default()
};
let output = skel.run_topology_node_init(input)?;
if output.return_value != 0 {
bail!(
"p2dq_topology_node_init returned {}",
output.return_value as i32
);
}

Ok(())
}

pub fn setup_topology<T: P2dqArenaProgs>(skel: &T) -> Result<()> {
let topo = Topology::new().expect("Failed to build host topology");

setup_topology_node(skel, topo.span.as_raw_slice())?;

for (_, node) in topo.nodes {
setup_topology_node(skel, node.span.as_raw_slice())?;
}

for (_, llc) in topo.all_llcs {
setup_topology_node(
skel,
Arc::<Llc>::into_inner(llc)
.expect("missing llc")
.span
.as_raw_slice(),
)?;
}

for (_, core) in topo.all_cores {
setup_topology_node(
skel,
Arc::<Core>::into_inner(core)
.expect("missing core")
.span
.as_raw_slice(),
)?;
}
for (_, cpu) in topo.all_cpus {
let mut mask = [0; 9];
mask[cpu.id.checked_shr(64).unwrap_or(0)] |= 1 << (cpu.id % 64);
setup_topology_node(skel, &mask)?;
}

Ok(())
}

#[macro_export]
macro_rules! init_open_skel {
($skel: expr, $opts: expr, $verbose: expr) => {
Expand Down Expand Up @@ -348,9 +238,6 @@ macro_rules! init_skel {
$skel.maps.bss_data.cpu_llc_ids[cpu.id] = cpu.llc_id as u64;
$skel.maps.bss_data.cpu_node_ids[cpu.id] = cpu.node_id as u64;
}

$crate::setup_arenas($skel)?;
$crate::setup_topology($skel)?;
};
}

Expand Down
Loading