Fix locking

4 files changed, 70 insertions, 179 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 6710683..8306e86 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -275,12 +275,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "core-foundation-sys"
-version = "0.8.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
-
-[[package]]
 name = "crossbeam"
 version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -594,15 +588,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "ntapi"
-version = "0.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4"
-dependencies = [
- "winapi",
-]
-
-[[package]]
 name = "num-conv"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -672,10 +657,10 @@ dependencies = [
  "libc",
  "perf-event",
  "plain",
+ "procfs",
  "rand",
  "scx_rustland_core",
  "scx_utils",
- "sysinfo",
 ]
 
 [[package]]
@@ -713,6 +698,28 @@ dependencies = [
 ]
 
 [[package]]
+name = "procfs"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc5b72d8145275d844d4b5f6d4e1eef00c8cd889edb6035c21675d1bb1f45c9f"
+dependencies = [
+ "bitflags",
+ "hex",
+ "procfs-core",
+ "rustix",
+]
+
+[[package]]
+name = "procfs-core"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "239df02d8349b06fc07398a3a1697b06418223b1c7725085e801e7c0fc6a12ec"
+dependencies = [
+ "bitflags",
+ "hex",
+]
+
+[[package]]
 name = "quote"
 version = "1.0.38"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -759,26 +766,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "rayon"
-version = "1.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa"
-dependencies = [
- "either",
- "rayon-core",
-]
-
-[[package]]
-name = "rayon-core"
-version = "1.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2"
-dependencies = [
- "crossbeam-deque",
- "crossbeam-utils",
-]
-
-[[package]]
 name = "redox_syscall"
 version = "0.5.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1020,20 +1007,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "sysinfo"
-version = "0.33.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4fc858248ea01b66f19d8e8a6d55f41deaf91e9d495246fd01368d99935c6c01"
-dependencies = [
- "core-foundation-sys",
- "libc",
- "memchr",
- "ntapi",
- "rayon",
- "windows",
-]
-
-[[package]]
 name = "tap"
 version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1197,22 +1170,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "winapi"
-version = "0.3.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
-dependencies = [
- "winapi-i686-pc-windows-gnu",
- "winapi-x86_64-pc-windows-gnu",
-]
-
-[[package]]
-name = "winapi-i686-pc-windows-gnu"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
-
-[[package]]
 name = "winapi-util"
 version = "0.1.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1222,65 +1179,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "winapi-x86_64-pc-windows-gnu"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
-
-[[package]]
-name = "windows"
-version = "0.57.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "12342cb4d8e3b046f3d80effd474a7a02447231330ef77d71daa6fbc40681143"
-dependencies = [
- "windows-core",
- "windows-targets",
-]
-
-[[package]]
-name = "windows-core"
-version = "0.57.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d2ed2439a290666cd67ecce2b0ffaad89c2a56b976b736e6ece670297897832d"
-dependencies = [
- "windows-implement",
- "windows-interface",
- "windows-result",
- "windows-targets",
-]
-
-[[package]]
-name = "windows-implement"
-version = "0.57.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn",
-]
-
-[[package]]
-name = "windows-interface"
-version = "0.57.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn",
-]
-
-[[package]]
-name = "windows-result"
-version = "0.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5e383302e8ec8515204254685643de10811af0ed97ea37210dc26fb0032647f8"
-dependencies = [
- "windows-targets",
-]
-
-[[package]]
 name = "windows-sys"
 version = "0.59.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 9b36ff8..8448899 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,7 +15,7 @@ rand = "0.9.0"
 iocuddle = "0.1.1"
 clap = { version = "4.5" , features = ["derive"] }
 perf-event = "0.4.8"
-sysinfo = "0.33.1"
+procfs = { version = "0.17.0", default-features = false }
 
 [build-dependencies]
 scx_rustland_core = "2.2.8"
diff --git a/src/main.rs b/src/main.rs
index 5a5051d..7bce3c7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,7 +25,7 @@ use std::io::{BufRead, BufReader};
 use std::mem::MaybeUninit;
 use std::ops::Range;
 use std::process;
-use sysinfo::System;
+use std::sync::mpsc::{Receiver, Sender};
 
 use anyhow::Result;
 
@@ -43,7 +43,8 @@ struct Scheduler<'a> {
     p_cores: Range<i32>,
     e_cores: Range<i32>,
     e_core_selector: Box<dyn ECoreSelector>,
-    system: System,
+    reciever: Receiver<(i32, i32)>,
+    sender: Sender<i32>,
 }
 
 impl<'a> Scheduler<'a> {
@@ -60,39 +61,43 @@ impl<'a> Scheduler<'a> {
         )?;
         dbg!("registering rust user space scheduler");
         let module: Box<dyn KernelModule> = if use_mocking {
-            Box::new(PerfEstimator::default())
+            Box::new(MockModule::default())
         } else {
             Box::new(KernelDriver::default())
         };
 
-        // let e_cores_file = File::open("/sys/devices/cpu_atom/cpus")?;
-        // let e_cores_reader = BufReader::new(e_cores_file);
-        // let e_cores = if let Some(Ok(line)) = e_cores_reader.lines().next() {
-        //     let cores: Vec<&str> = line.split('-').collect();
-        //     cores[0].parse::<i32>().unwrap()..cores[1].parse::<i32>().unwrap()
-        // } else {
-        //     panic!(
-        //         "Was not able to differentiate between core types. Does the system have e cores?"
-        //     )
-        // };
-
-        // let p_cores_file = File::open("/sys/devices/cpu_core/cpus")?;
-        // let p_cores_reader = BufReader::new(p_cores_file);
-        // let p_cores = if let Some(Ok(line)) = p_cores_reader.lines().next() {
-        //     let cores: Vec<&str> = line.split('-').collect();
-        //     cores[0].parse::<i32>().unwrap()..cores[1].parse::<i32>().unwrap()
-        // } else {
-        //     panic!(
-        //         "Was not able to differentiate between core types. Does the system have e cores?"
-        //     )
-        // };
-        let e_cores = 0..2;
-        let p_cores = 2..4;
+        let read_cores = |core_type: &str| {
+            let e_cores_file = File::open(format!("/sys/devices/cpu_{core_type}/cpus"))?;
+            let e_cores_reader = BufReader::new(e_cores_file);
+            let e_cores = if let Some(Ok(line)) = e_cores_reader.lines().next() {
+                let cores: Vec<&str> = line.split('-').collect();
+                cores[0].parse::<i32>().unwrap()..cores[1].parse::<i32>().unwrap()
+            } else {
+                panic!(
+                    "Was not able to differentiate between core types. Does the system have e cores?"
+                )
+            };
+            Ok::<_, std::io::Error>(e_cores)
+        };
+        let e_cores = read_cores("atom").unwrap_or(0..0);
+        let p_cores = read_cores("core").unwrap_or(0..4);
 
         let selector = Box::new(RoundRobinSelector::new(&e_cores));
 
-        let mut system = System::new_all();
-        system.refresh_processes(sysinfo::ProcessesToUpdate::All, true);
+        let (pid_send, pid_recieve) = std::sync::mpsc::channel();
+        let (parent_send, parent_recieve) = std::sync::mpsc::channel();
+
+        std::thread::spawn(move || loop {
+            if let Ok(pid) = pid_recieve.recv() {
+                let parent = (|| {
+                    let process = procfs::process::Process::new(pid)?;
+                    process.stat().map(|stat| stat.ppid)
+                })()
+                .unwrap_or_default();
+
+                parent_send.send((pid, parent)).unwrap();
+            }
+        });
 
         Ok(Self {
             bpf,
@@ -106,7 +111,8 @@ impl<'a> Scheduler<'a> {
             p_cores,
             e_cores,
             e_core_selector: selector,
-            system,
+            sender: pid_send,
+            reciever: parent_recieve,
         })
     }
 
@@ -147,12 +153,15 @@ impl<'a> Scheduler<'a> {
                 }
             } else {
                 self.module.start_trace(task.pid as u64);
+
+                self.sender.send(task.pid).unwrap();
+
                 self.managed_tasks.insert(
                     task.pid as u32,
                     TaskState {
                         previous_energy_usage: 0,
                         budget: self.maximum_budget,
-                        children: vec![],
+                        parent: 0,
                     },
                 );
                 self.task_queue.push_back(task);
@@ -219,12 +228,9 @@ impl<'a> Scheduler<'a> {
     }
 
     fn dispatch_tasks(&mut self) {
-        panic!();
         loop {
             //TODO: we should probably not do this every time, but instead wait a predefined amount of time between invocations
             self.reset_budgets_and_garbage_collect();
-            //TODO: we should probably not do this every time, but instead wait a predefined amount of time between invocations
-            self.refresh_children();
 
             // Consume all tasks before dispatching any.
             self.consume_all_tasks();
@@ -232,6 +238,10 @@ impl<'a> Scheduler<'a> {
             // Dispatch one task from the queue.
             self.dispatch_next_task();
 
+            for (pid, parent) in self.reciever.try_iter() {
+                self.managed_tasks.get_mut(&(pid as u32)).unwrap().parent = parent as u32;
+            }
+
             // If no task is ready to run (or in case of error), stop dispatching tasks and notify
             // the BPF component that all tasks have been scheduled / dispatched, with no remaining
             // pending tasks.
@@ -244,7 +254,7 @@ impl<'a> Scheduler<'a> {
     }
 
     fn run(&mut self) -> Result<UserExitInfo> {
-        while dbg!(!self.bpf.exited()) {
+        while !self.bpf.exited() {
             self.dispatch_tasks();
         }
         self.bpf.shutdown_and_report()
@@ -269,23 +279,6 @@ impl<'a> Scheduler<'a> {
             was_scheduled
         });
     }
-
-    //TODO: this is probably not very efficient
-    fn refresh_children(&mut self) {
-        self.system
-            .refresh_processes(sysinfo::ProcessesToUpdate::All, true);
-        for entry in self.managed_tasks.values_mut() {
-            entry.children.clear();
-        }
-
-        for (pid, process) in self.system.processes() {
-            if let Some(parent_pid) = process.parent() {
-                if let Some(entry) = self.managed_tasks.get_mut(&parent_pid.as_u32()) {
-                    entry.children.push(pid.as_u32());
-                }
-            }
-        }
-    }
 }
 
 fn main() -> Result<()> {
diff --git a/src/task_state.rs b/src/task_state.rs
index e56b109..6043cf4 100644
--- a/src/task_state.rs
+++ b/src/task_state.rs
@@ -1,5 +1,5 @@
 pub struct TaskState {
     pub previous_energy_usage: u64,
     pub budget: u64,
-    pub children: Vec<u32>,
-}
\ No newline at end of file
+    pub parent: u32,
+}