Merge pull request #218 from sched-ext/rusty_hotplug

Gracefully handle hotplug in scx_rusty

rust/scx_utils/src/user_exit_info.rs

@@ -22,6 +22,7 @@ pub enum ScxExitKind {
     None = bindings::scx_exit_kind_SCX_EXIT_NONE as isize,
     Done = bindings::scx_exit_kind_SCX_EXIT_DONE as isize,
     Unreg = bindings::scx_exit_kind_SCX_EXIT_UNREG as isize,
+    UnregBPF = bindings::scx_exit_kind_SCX_EXIT_UNREG_BPF as isize,
     SysRq = bindings::scx_exit_kind_SCX_EXIT_SYSRQ as isize,
     Error = bindings::scx_exit_kind_SCX_EXIT_ERROR as isize,
     ErrorBPF = bindings::scx_exit_kind_SCX_EXIT_ERROR_BPF as isize,
@@ -40,9 +41,14 @@ macro_rules! uei_read {
         scx_utils::paste! {
             let bpf_uei = $skel.data().$uei;
             let bpf_dump = scx_utils::UEI_DUMP_PTR_MUTEX.lock().unwrap().ptr;
+            let exit_code_ptr = match scx_utils::compat::struct_has_field("scx_exit_info", "exit_code") {
+                Ok(true) => &bpf_uei.exit_code as *const _,
+                _ => std::ptr::null(),
+            };
 
             scx_utils::UserExitInfo::new(
                 &bpf_uei.kind as *const _,
+                exit_code_ptr,
                 bpf_uei.reason.as_ptr() as *const _,
                 bpf_uei.msg.as_ptr() as *const _,
                 bpf_dump,
@@ -103,6 +109,7 @@ pub struct UserExitInfo {
     /// The C enum scx_exit_kind value. Test against ScxExitKind. None-zero
     /// value indicates that the BPF scheduler has exited.
     kind: i32,
+    exit_code: i64,
     reason: Option<String>,
     msg: Option<String>,
     dump: Option<String>,
@@ -116,11 +123,17 @@ impl UserExitInfo {
     /// type which then calls this method with the individual fields.
     pub fn new(
         kind_ptr: *const i32,
+        exit_code_ptr: *const i64,
         reason_ptr: *const c_char,
         msg_ptr: *const c_char,
         dump_ptr: *const c_char,
     ) -> Self {
         let kind = unsafe { std::ptr::read_volatile(kind_ptr) };
+        let exit_code = if exit_code_ptr.is_null() {
+            0
+        } else {
+            unsafe { std::ptr::read_volatile(exit_code_ptr) }
+        };
 
         let (reason, msg) = (
             Some(
@@ -153,6 +166,7 @@ impl UserExitInfo {
 
         Self {
             kind,
+            exit_code,
             reason,
             msg,
             dump,
@@ -180,11 +194,22 @@ impl UserExitInfo {
             _ => "<UNKNOWN>".into(),
         };
 
-        if self.kind <= ScxExitKind::Unreg as i32 {
+        if self.kind <= ScxExitKind::UnregBPF as i32 {
             eprintln!("{}", why);
             Ok(())
         } else {
             bail!("{}", why)
         }
     }
+
+    /// Return the exit code that the scheduler gracefully exited with. This
+    /// only applies when the BPF scheduler exits with scx_bpf_exit(), i.e. kind
+    /// ScxExitKind::UnregBPF.
+    pub fn exit_code(&self) -> Option<i64> {
+        if self.kind == ScxExitKind::UnregBPF as i32 {
+            Some(self.exit_code)
+        } else {
+            None
+        }
+    }
 }

scheds/include/scx/common.bpf.h

@@ -30,7 +30,7 @@ static inline void ___vmlinux_h_sanity_check___(void)
 
 void scx_bpf_error_bstr(char *fmt, unsigned long long *data, u32 data_len) __ksym;
 void scx_bpf_exit_bstr(s64 exit_code, char *fmt,
-		       unsigned long long *data, u32 data__sz) __ksym;
+		       unsigned long long *data, u32 data__sz) __ksym __weak;
 
 static inline __attribute__((format(printf, 1, 2)))
 void ___scx_bpf_exit_format_checker(const char *fmt, ...) {}

scheds/rust/scx_rusty/src/bpf/intf.h

@@ -45,6 +45,10 @@ enum consts {
 	MAX_DOM_ACTIVE_PIDS	= 1024,
 };
 
+enum rusty_exit_codes {
+	RUSTY_EXIT_HOTPLUG,
+};
+
 /* Statistics */
 enum stat_idx {
 	/* The following fields add up to all dispatched tasks */

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

@@ -1428,6 +1428,22 @@ static s32 initialize_cpu(s32 cpu)
 	return 0;
 }
 
+void BPF_STRUCT_OPS(rusty_cpu_online, s32 cpu)
+{
+	if (bpf_ksym_exists(scx_bpf_exit_bstr))
+		scx_bpf_exit(RUSTY_EXIT_HOTPLUG, "CPU %d went online", cpu);
+	else
+		scx_bpf_error("CPU %d went online", cpu);
+}
+
+void BPF_STRUCT_OPS(rusty_cpu_offline, s32 cpu)
+{
+	if (bpf_ksym_exists(scx_bpf_exit_bstr))
+		scx_bpf_exit(RUSTY_EXIT_HOTPLUG, "CPU %d went offline", cpu);
+	else
+		scx_bpf_error("CPU %d went offline", cpu);
+}
+
 s32 BPF_STRUCT_OPS_SLEEPABLE(rusty_init)
 {
 	struct bpf_cpumask *cpumask;
@@ -1497,6 +1513,8 @@ SCX_OPS_DEFINE(rusty,
 	       .set_cpumask		= (void *)rusty_set_cpumask,
 	       .init_task		= (void *)rusty_init_task,
 	       .exit_task		= (void *)rusty_exit_task,
+	       .cpu_online		= (void *)rusty_cpu_online,
+	       .cpu_offline		= (void *)rusty_cpu_offline,
 	       .init			= (void *)rusty_init,
 	       .exit			= (void *)rusty_exit,
 	       .name			= "rusty");

scheds/rust/scx_rusty/src/main.rs

@@ -39,9 +39,10 @@ use scx_utils::init_libbpf_logging;
 use scx_utils::scx_ops_attach;
 use scx_utils::scx_ops_load;
 use scx_utils::uei_exited;
-use scx_utils::uei_report;
+use scx_utils::uei_read;
 use scx_utils::Cpumask;
 use scx_utils::Topology;
+use scx_utils::UserExitInfo;
 
 const MAX_DOMS: usize = bpf_intf::consts_MAX_DOMS as usize;
 const MAX_CPUS: usize = bpf_intf::consts_MAX_CPUS as usize;
@@ -527,7 +528,7 @@ impl<'a> Scheduler<'a> {
         Ok(())
     }
 
-    fn run(&mut self, shutdown: Arc<AtomicBool>) -> Result<()> {
+    fn run(&mut self, shutdown: Arc<AtomicBool>) -> Result<UserExitInfo> {
         let now = Instant::now();
         let mut next_tune_at = now + self.tune_interval;
         let mut next_sched_at = now + self.sched_interval;
@@ -559,7 +560,9 @@ impl<'a> Scheduler<'a> {
         }
 
         self.struct_ops.take();
-        uei_report!(&self.skel, uei)
+        let uei = uei_read!(&self.skel, uei);
+        uei.report()?;
+        Ok(uei)
     }
 }
 
@@ -591,8 +594,6 @@ fn main() -> Result<()> {
         simplelog::ColorChoice::Auto,
     )?;
 
-    let mut sched = Scheduler::init(&opts)?;
-
     let shutdown = Arc::new(AtomicBool::new(false));
     let shutdown_clone = shutdown.clone();
     ctrlc::set_handler(move || {
@@ -600,5 +601,16 @@ fn main() -> Result<()> {
     })
     .context("Error setting Ctrl-C handler")?;
 
-    sched.run(shutdown)
+    while !shutdown.load(Ordering::Relaxed) {
+        let mut sched = Scheduler::init(&opts)?;
+
+        let uei = sched.run(shutdown.clone())?;
+        if let Some(exit_code) = uei.exit_code() {
+            if exit_code == bpf_intf::rusty_exit_codes_RUSTY_EXIT_HOTPLUG as i64 {
+                continue;
+            }
+        }
+        break;
+    }
+    Ok(())
 }