lib, scx_p2dq: factor out arena interface programs

lib/arena.bpf.c

@@ -0,0 +1,83 @@
+/*
+ * SPDX-License-Identifier: GPL-2.0
+ * Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
+ */
+#include <scx/common.bpf.h>
+#include <lib/sdt_task.h>
+
+#include <lib/arena.h>
+#include <lib/percpu.h>
+#include <lib/cpumask.h>
+#include <lib/topology.h>
+
+/*
+ * "System-call" based API for arenas.
+ */
+
+struct task_ctx;
+u64 arena_topo_setup_ptr;
+
+SEC("syscall")
+int arena_init(struct arena_init_args *args)
+{
+	int ret;
+
+	ret = scx_static_init(args->static_pages);
+	if (ret)
+		return ret;
+
+	/* How many types to store all CPU IDs? */
+	ret = scx_bitmap_init(div_round_up(nr_cpu_ids, 8));
+	if (ret)
+		return ret;
+
+	ret = scx_percpu_storage_init();
+	if (ret)
+		return ret;
+
+	ret = scx_task_init(args->task_ctx_size);
+	if (ret)
+		return ret;
+
+	return 0;
+}
+
+SEC("syscall")
+int arena_alloc_mask(void)
+{
+	scx_bitmap_t bitmap;
+
+	bitmap = scx_bitmap_alloc();
+	if (!bitmap)
+		return -ENOMEM;
+
+	arena_topo_setup_ptr = (u64)&bitmap->bits;
+
+	return 0;
+}
+
+SEC("syscall")
+int arena_topology_node_init(void)
+{
+	scx_bitmap_t bitmap = (scx_bitmap_t)container_of(arena_topo_setup_ptr, struct scx_bitmap, bits);
+	int ret;
+
+	ret = topo_init(bitmap);
+	if (ret)
+		return ret;
+
+	arena_topo_setup_ptr = 0;
+
+	return 0;
+}
+
+SEC("syscall")
+int arena_topology_print(void)
+{
+	scx_arena_subprog_init();
+
+	topo_print();
+
+	return 0;
+}
+

scheds/include/lib/arena.h

@@ -0,0 +1,16 @@
+#pragma once
+
+struct arena_init_args {
+	u64 static_pages;
+	u64 task_ctx_size;
+};
+
+int arena_init(struct arena_init_args *args);
+int arena_alloc_mask(void);
+
+struct arena_topology_node_init_args {
+	u64 setup_ptr;
+};
+
+int arena_topology_node_init();
+int arena_topology_print(void);

scheds/rust/scx_chaos/build.rs

@@ -36,6 +36,7 @@ fn main() {
         .add_source("src/bpf/lib/sdt_alloc.bpf.c")
         .add_source("src/bpf/lib/bitmap.bpf.c")
         .add_source("src/bpf/lib/topology.bpf.c")
+        .add_source("src/bpf/lib/arena.bpf.c")
         .compile_link_gen()
         .unwrap();
 }

scheds/rust/scx_chaos/src/lib.rs

@@ -20,6 +20,10 @@ use scx_utils::Core;
 use scx_utils::Llc;
 use scx_utils::Topology;
 
+use scx_p2dq::bpf_intf::consts_STATIC_ALLOC_PAGES_GRANULARITY;
+use scx_p2dq::types;
+use std::ffi::c_ulong;
+
 use anyhow::bail;
 use anyhow::Context;
 use anyhow::Result;
@@ -181,11 +185,22 @@ impl Builder<'_> {
         // 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 mut args = types::arena_init_args {
+            static_pages: consts_STATIC_ALLOC_PAGES_GRANULARITY as c_ulong,
+            task_ctx_size: std::mem::size_of::<types::task_p2dq>() as c_ulong,
+        };
+
         let input = ProgramInput {
+            context_in: Some(unsafe {
+                std::slice::from_raw_parts_mut(
+                    &mut args as *mut _ as *mut u8,
+                    std::mem::size_of_val(&args),
+                )
+            }),
             ..Default::default()
         };
 
-        let output = skel.progs.p2dq_arena_init.test_run(input)?;
+        let output = skel.progs.arena_init.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "Could not initialize arenas, p2dq_setup returned {}",
@@ -202,24 +217,25 @@ impl Builder<'_> {
             ..Default::default()
         };
 
-        let output = skel.progs.p2dq_alloc_mask.test_run(input)?;
+        let output = skel.progs.arena_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 ptr = unsafe {
+            std::mem::transmute::<u64, &mut [u64; 10]>(skel.maps.bss_data.arena_topo_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)?;
+        let output = skel.progs.arena_topology_node_init.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "p2dq_topology_node_init returned {}",

scheds/rust/scx_p2dq/build.rs

@@ -12,6 +12,7 @@ fn main() {
         .add_source("src/bpf/lib/sdt_alloc.bpf.c")
         .add_source("src/bpf/lib/bitmap.bpf.c")
         .add_source("src/bpf/lib/topology.bpf.c")
+        .add_source("src/bpf/lib/arena.bpf.c")
         .compile_link_gen()
         .unwrap();
 }

scheds/rust/scx_p2dq/src/bpf/main.bpf.c

@@ -87,8 +87,6 @@ extern const volatile u32 nr_cpu_ids;
 const u64 lb_timer_intvl_ns = 250LLU * NSEC_PER_MSEC;
 const u64 lb_backoff_ns = 5LLU * NSEC_PER_MSEC;
 
-u64 setup_ptr;
-
 u64 cpu_llc_ids[MAX_CPUS];
 u64 cpu_node_ids[MAX_CPUS];
 u64 big_core_ids[MAX_CPUS];
@@ -1567,69 +1565,6 @@ s32 static start_timers(void)
 	return 0;
 }
 
-SEC("syscall")
-int p2dq_topo_print(void)
-{
-	scx_arena_subprog_init();
-	topo_print();
-	return 0;
-}
-
-SEC("syscall")
-int p2dq_arena_init(void)
-{
-	int ret;
-
-	ret = scx_static_init(STATIC_ALLOC_PAGES_GRANULARITY);
-	if (ret)
-		return ret;
-
-	/* How many types to store all CPU IDs? */
-	ret = scx_bitmap_init(div_round_up(nr_cpu_ids, 8));
-	if (ret)
-		return ret;
-
-	ret = scx_percpu_storage_init();
-	if (ret)
-		return ret;
-
-	ret = scx_task_init(sizeof(task_ctx));
-	if (ret)
-		return ret;
-
-	return 0;
-}
-
-SEC("syscall")
-int p2dq_alloc_mask(void)
-{
-	scx_bitmap_t bitmap;
-
-	bitmap = scx_bitmap_alloc();
-	if (!bitmap)
-		return -ENOMEM;
-
-	setup_ptr = (u64)&bitmap->bits;
-
-	return 0;
-}
-
-SEC("syscall")
-int p2dq_topology_node_init(void)
-{
-	scx_bitmap_t bitmap = (scx_bitmap_t)container_of(setup_ptr, struct scx_bitmap, bits);
-	int ret;
-
-	ret = topo_init(bitmap);
-	if (ret)
-		return ret;
-
-	setup_ptr = 0;
-
-	return 0;
-}
-
-
 static __always_inline s32 p2dq_init_impl()
 {
 	int i, ret;

scheds/rust/scx_p2dq/src/lib.rs

@@ -2,6 +2,11 @@
 
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2.
+pub mod bpf_intf;
+
+mod bpf_skel;
+pub use bpf_skel::*;
+
 pub use scx_utils::CoreType;
 use scx_utils::Topology;
 pub use scx_utils::NR_CPU_IDS;

scheds/rust/scx_p2dq/src/main.rs

@@ -39,6 +39,8 @@ use scx_utils::UserExitInfo;
 use scx_utils::NR_CPU_IDS;
 use scx_utils::{Core, Llc};
 
+use std::ffi::c_ulong;
+
 use bpf_intf::stat_idx_P2DQ_NR_STATS;
 use bpf_intf::stat_idx_P2DQ_STAT_DIRECT;
 use bpf_intf::stat_idx_P2DQ_STAT_DISPATCH_PICK2;
@@ -164,11 +166,22 @@ impl<'a> Scheduler<'a> {
         // 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 mut args = types::arena_init_args {
+            static_pages: bpf_intf::consts_STATIC_ALLOC_PAGES_GRANULARITY as c_ulong,
+            task_ctx_size: std::mem::size_of::<types::task_p2dq>() as c_ulong,
+        };
+
         let input = ProgramInput {
+            context_in: Some(unsafe {
+                std::slice::from_raw_parts_mut(
+                    &mut args as *mut _ as *mut u8,
+                    std::mem::size_of_val(&args),
+                )
+            }),
             ..Default::default()
         };
 
-        let output = self.skel.progs.p2dq_arena_init.test_run(input)?;
+        let output = self.skel.progs.arena_init.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "Could not initialize arenas, p2dq_setup returned {}",
@@ -185,7 +198,7 @@ impl<'a> Scheduler<'a> {
             ..Default::default()
         };
 
-        let output = self.skel.progs.p2dq_alloc_mask.test_run(input)?;
+        let output = self.skel.progs.arena_alloc_mask.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "Could not initialize arenas, setup_topology_node returned {}",
@@ -194,7 +207,7 @@ impl<'a> Scheduler<'a> {
         }
 
         let ptr = unsafe {
-            std::mem::transmute::<u64, &mut [u64; 10]>(self.skel.maps.bss_data.setup_ptr)
+            std::mem::transmute::<u64, &mut [u64; 10]>(self.skel.maps.bss_data.arena_topo_setup_ptr)
         };
 
         let (valid_mask, _) = ptr.split_at_mut(mask.len());
@@ -203,7 +216,7 @@ impl<'a> Scheduler<'a> {
         let input = ProgramInput {
             ..Default::default()
         };
-        let output = self.skel.progs.p2dq_topology_node_init.test_run(input)?;
+        let output = self.skel.progs.arena_topology_node_init.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "p2dq_topology_node_init returned {}",
@@ -269,7 +282,7 @@ impl<'a> Scheduler<'a> {
             ..Default::default()
         };
 
-        let output = self.skel.progs.p2dq_topo_print.test_run(input)?;
+        let output = self.skel.progs.arena_topology_print.test_run(input)?;
         if output.return_value != 0 {
             bail!(
                 "Could not initialize arenas, topo_print returned {}",