scx_rustland: determine the amount of cores via /proc/stat

libbpf_rs::num_possible_cpus() may take into account multi-threads
multi-cores information, that are not used efficiently by the scheduler
at the moment.

For simplicity rely on /proc/stat to determine the amount of CPUs that
can be used by the scheduler and provide a proper abstraction to access
this information from the bpf Rust module.

Signed-off-by: Andrea Righi <[email protected]>

Author: Andrea Righi

scheds/rust/scx_rustland/src/bpf.rs

@@ -7,6 +7,8 @@ use crate::bpf_intf;
 use crate::bpf_skel::*;
 
 use std::ffi::CStr;
+use std::fs::File;
+use std::io::{self, BufRead};
 
 use anyhow::Context;
 use anyhow::Result;
@@ -200,6 +202,13 @@ impl<'a> BpfScheduler<'a> {
         let skel_builder = BpfSkelBuilder::default();
         let mut skel = skel_builder.open().context("Failed to open BPF program")?;
 
+        // Initialize online CPUs counter.
+        //
+        // We should probably refresh this counter during the normal execution to support cpu
+        // hotplugging, but for now let's keep it simple and set this only at initialization).
+        let nr_cpus_online = Self::count_cpus()?;
+        skel.rodata_mut().num_possible_cpus = nr_cpus_online;
+
         // Set scheduler options (defined in the BPF part).
         skel.bss_mut().usersched_pid = std::process::id();
         skel.rodata_mut().slice_ns = slice_us * 1000;
@@ -226,6 +235,30 @@ impl<'a> BpfScheduler<'a> {
         }
     }
 
+    // Return the amount of available CPUs in the system (according to /proc/stat).
+    fn count_cpus() -> io::Result<i32> {
+        let file = File::open("/proc/stat")?;
+        let reader = io::BufReader::new(file);
+        let mut cpu_count = -1;
+
+        for line in reader.lines() {
+            let line = line?;
+            if line.starts_with("cpu") {
+                cpu_count += 1;
+            } else {
+                break;
+            }
+        }
+
+        Ok(cpu_count)
+    }
+
+    // Override the default scheduler time slice (in us).
+    #[allow(dead_code)]
+    pub fn get_nr_cpus(&self) -> i32 {
+        self.skel.rodata().num_possible_cpus
+    }
+
     // Override the default scheduler time slice (in us).
     #[allow(dead_code)]
     pub fn set_effective_slice_us(&mut self, slice_us: u64) {

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

@@ -44,7 +44,7 @@ char _license[] SEC("license") = "GPL";
 #define MAX_CPUS 1024
 
 /* !0 for veristat, set during init */
-const volatile u32 num_possible_cpus = 8;
+const volatile s32 num_possible_cpus = 8;
 
 /*
  * Exit info (passed to the user-space counterpart).

scheds/rust/scx_rustland/src/main.rs

@@ -164,7 +164,6 @@ struct Scheduler<'a> {
     task_pool: TaskTree,   // tasks ordered by vruntime
     task_map: TaskInfoMap, // map pids to the corresponding task information
     min_vruntime: u64,     // Keep track of the minimum vruntime across all tasks
-    nr_cpus_online: i32,   // Amount of the available CPUs in the system
     slice_ns: u64,         // Default time slice (in ns)
 }
 
@@ -186,19 +185,12 @@ impl<'a> Scheduler<'a> {
         // Initialize global minimum vruntime.
         let min_vruntime: u64 = 0;
 
-        // Initialize online CPUs counter.
-        //
-        // We should probably refresh this counter during the normal execution to support cpu
-        // hotplugging, but for now let's keep it simple and set this only at initialization).
-        let nr_cpus_online = libbpf_rs::num_possible_cpus().unwrap() as i32;
-
         // Return scheduler object.
         Ok(Self {
             bpf,
             task_pool,
             task_map,
             min_vruntime,
-            nr_cpus_online,
             slice_ns,
         })
     }
@@ -207,7 +199,7 @@ impl<'a> Scheduler<'a> {
     fn get_idle_cpus(&self) -> Vec<i32> {
         let mut idle_cpus = Vec::new();
 
-        for cpu in 0..self.nr_cpus_online {
+        for cpu in 0..self.bpf.get_nr_cpus() {
             let pid = self.bpf.get_cpu_pid(cpu);
             if pid == 0 {
                 idle_cpus.push(cpu);
@@ -327,7 +319,7 @@ impl<'a> Scheduler<'a> {
         let nr_queued = *self.bpf.nr_queued_mut();
         let nr_scheduled = *self.bpf.nr_scheduled_mut();
         let nr_waiting = nr_queued + nr_scheduled;
-        let nr_cpus = self.nr_cpus_online as u64;
+        let nr_cpus = self.bpf.get_nr_cpus() as u64;
 
         // Scale time slice, but never scale below 1 ms.
         let scaling = nr_waiting / nr_cpus + 1;
@@ -475,7 +467,7 @@ impl<'a> Scheduler<'a> {
             Err(_) => -1,
         };
         info!("Running tasks:");
-        for cpu in 0..self.nr_cpus_online {
+        for cpu in 0..self.bpf.get_nr_cpus() {
             let pid = if cpu == sched_cpu {
                 "[self]".to_string()
             } else {