Fix scheduler lockup

4 files changed, 131 insertions, 48 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8306e86..22af256 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -68,6 +68,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6b964d184e89d9b6b67dd2715bc8e74cf3107fb2b529990c90cf517326150bf4"
 
 [[package]]
+name = "autocfg"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26"
+
+[[package]]
 name = "bindgen"
 version = "0.71.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -341,6 +347,20 @@ dependencies = [
 ]
 
 [[package]]
+name = "dashmap"
+version = "6.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf"
+dependencies = [
+ "cfg-if",
+ "crossbeam-utils",
+ "hashbrown",
+ "lock_api",
+ "once_cell",
+ "parking_lot_core",
+]
+
+[[package]]
 name = "deranged"
 version = "0.3.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -408,6 +428,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2"
 
 [[package]]
+name = "hashbrown"
+version = "0.14.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
+
+[[package]]
 name = "heck"
 version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -539,6 +565,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
 
 [[package]]
+name = "lock_api"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17"
+dependencies = [
+ "autocfg",
+ "scopeguard",
+]
+
+[[package]]
 name = "log"
 version = "0.4.26"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -609,6 +645,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e"
 
 [[package]]
+name = "parking_lot_core"
+version = "0.9.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_syscall",
+ "smallvec",
+ "windows-targets",
+]
+
+[[package]]
 name = "paste"
 version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -652,6 +701,7 @@ dependencies = [
  "anyhow",
  "clap",
  "ctrlc",
+ "dashmap",
  "iocuddle",
  "libbpf-rs 0.24.8",
  "libc",
@@ -850,6 +900,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
+[[package]]
 name = "scx_rustland_core"
 version = "2.2.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -957,6 +1013,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
 [[package]]
+name = "smallvec"
+version = "1.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd"
+
+[[package]]
 name = "sscanf"
 version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 8448899..a1ebd01 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,7 @@ iocuddle = "0.1.1"
 clap = { version = "4.5" , features = ["derive"] }
 perf-event = "0.4.8"
 procfs = { version = "0.17.0", default-features = false }
+dashmap = "6.1.0"
 
 [build-dependencies]
 scx_rustland_core = "2.2.8"
diff --git a/src/main.rs b/src/main.rs
index 95c31a8..6f5d2ae 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,7 @@ mod bpf;
 use bpf::*;
 
 use clap::{Arg, ArgAction, Command};
+use dashmap::DashMap;
 use e_core_selector::{ECoreSelector, RoundRobinSelector};
 use mock::perf::PerfEstimator;
 use scx_utils::UserExitInfo;
@@ -19,12 +20,14 @@ use libbpf_rs::OpenObject;
 use mock::{KernelDriver, KernelModule, MockModule};
 use task_state::TaskState;
 
+use std::any::Any;
 use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
 use std::mem::MaybeUninit;
 use std::ops::Range;
 use std::sync::mpsc::{Receiver, Sender, SyncSender};
+use std::sync::{Arc, OnceLock};
 use std::time::Duration;
 use std::{process, thread};
 
@@ -45,8 +48,9 @@ struct Scheduler<'a> {
     p_cores: Range<i32>,
     e_cores: Range<i32>,
     e_core_selector: Box<dyn ECoreSelector>,
-    reciever: Receiver<(Pid, Response)>,
+    // reciever: Receiver<(Pid, Response)>,
     sender: SyncSender<(Pid, Request)>,
+    hashmap: Arc<DashMap<Pid, (u64, i32)>>,
 }
 
 enum Request {
@@ -73,6 +77,9 @@ impl<'a> Scheduler<'a> {
             false, // debug (false = debug mode off)
         )?;
         dbg!("registering rust user space scheduler");
+        let map: DashMap<i32, (u64, i32)> = DashMap::with_capacity(100000);
+        let map = Arc::new(map);
+        let thread_map = map.clone();
 
         let read_cores = |core_type: &str| {
             let e_cores_file = File::open(format!("/sys/devices/cpu_{core_type}/cpus"))?;
@@ -87,8 +94,8 @@ impl<'a> Scheduler<'a> {
             };
             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 e_cores = read_cores("atom").unwrap_or(0..4);
+        let p_cores = read_cores("core").unwrap_or(4..8);
 
         let selector = Box::new(RoundRobinSelector::new(&e_cores));
 
@@ -109,9 +116,13 @@ impl<'a> Scheduler<'a> {
                 i += 1;
                 // if i % 1000 == 0 {
                 // println!("reading energy");
-                for &pid in tasks.iter().nth((i / 100) % 1000) {
-                    let _ = response_send
-                        .try_send((pid, Response::Energy(module.read_consumption(pid as u64))));
+                for pid in tasks.iter()
+                /*.nth((i / 100) % tasks.len().max(100))*/
+                {
+                    let energy = module.read_consumption(*pid as u64);
+                    if let Some(mut entry) = thread_map.try_get_mut(pid).try_unwrap() {
+                        entry.0 = energy;
+                    }
                 }
                 // }
 
@@ -125,10 +136,11 @@ impl<'a> Scheduler<'a> {
                             .unwrap_or_default();
                             module.start_trace(pid as u64);
                             tasks.insert(pid);
+                            thread_map.insert(pid, (0, parent));
 
-                            response_send
-                                .try_send((pid, Response::Parent(parent)))
-                                .unwrap();
+                            // response_send
+                            //     .try_send((pid, Response::Parent(parent)))
+                            //     .unwrap();
                         }
                         Request::ReadPerf => {
                             response_send
@@ -145,7 +157,7 @@ impl<'a> Scheduler<'a> {
                         Request::Heartbeat => {}
                     }
                 }
-                std::thread::sleep(Duration::from_micros(50000));
+                std::thread::sleep(Duration::from_micros(5000));
             }
         });
 
@@ -161,7 +173,8 @@ impl<'a> Scheduler<'a> {
             e_cores,
             e_core_selector: selector,
             sender: request_send,
-            reciever: response_recieve,
+            // reciever: response_recieve,
+            hashmap: map,
         })
     }
 
@@ -182,34 +195,38 @@ impl<'a> Scheduler<'a> {
         //
         // Although the FIFO scheduler doesn't use these fields, they can provide valuable data for
         // implementing more sophisticated scheduling policies.
+        // let mut contended = Vec::with_capacity(100);
 
         while let Ok(Some(task)) = self.bpf.dequeue_task() {
             // The scheduler itself has to be scheduled regardless of its energy usage
             if task.pid == self.own_pid {
-                self.task_queue.push_back(task);
+                self.task_queue.push_front(task);
+                self.dispatch_next_task();
                 continue;
             }
 
-            if let Some(task_state) = self.managed_tasks.get_mut(&(task.pid)) {
-                // self.sender.try_send((task.pid, Request::ReadPerf)).unwrap();
-                if task_state.budget == 0 {
-                    println!("budget zero");
-                    self.no_budget_task_queue.push_back(task);
-                } else {
+            match self.hashmap.try_get(&task.pid) {
+                dashmap::try_result::TryResult::Present(value) => {
+                    let (energy, parent) = *value;
+
+                    // dbg!(energy);
+                    // TODO: fix
+                    if energy == 0 {
+                        // println!("budget zero");
+                        self.no_budget_task_queue.push_back(task);
+                    } else {
+                        self.task_queue.push_back(task);
+                    }
+                }
+                dashmap::try_result::TryResult::Absent => {
+                    self.hashmap.insert(task.pid, (0, 0));
+                    self.sender.try_send((task.pid, Request::NewTask)).unwrap();
+                    self.task_queue.push_back(task);
+                }
+                dashmap::try_result::TryResult::Locked => {
+                    // println!("locked");
                     self.task_queue.push_back(task);
                 }
-            } else {
-                self.sender.try_send((task.pid, Request::NewTask)).unwrap();
-
-                self.managed_tasks.insert(
-                    task.pid,
-                    TaskState {
-                        previous_energy_usage: 0,
-                        budget: self.maximum_budget,
-                        parent: 0,
-                    },
-                );
-                self.task_queue.push_back(task);
             }
         }
     }
@@ -275,7 +292,7 @@ impl<'a> Scheduler<'a> {
     fn dispatch_tasks(&mut self) {
         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();
+            // self.reset_budgets_and_garbage_collect();
 
             // Consume all tasks before dispatching any.
             self.consume_all_tasks();
@@ -291,22 +308,22 @@ impl<'a> Scheduler<'a> {
                 self.bpf.notify_complete(0);
 
                 // self.sender.try_send((0, Request::Heartbeat)).unwrap();
-                for (pid, response) in self.reciever.try_iter() {
-                    let Some(task_state) = self.managed_tasks.get_mut(&(pid)) else {
-                        continue;
-                    };
-                    match response {
-                        Response::Parent(parent) => {
-                            task_state.parent = parent;
-                        }
-                        Response::Energy(energy) => {
-                            let used_energy = energy - task_state.previous_energy_usage.min(energy);
-                            task_state.previous_energy_usage = energy;
+                // for (pid, response) in self.reciever.try_iter() {
+                //     let Some(task_state) = self.managed_tasks.get_mut(&(pid)) else {
+                //         continue;
+                //     };
+                //     match response {
+                //         Response::Parent(parent) => {
+                //             task_state.parent = parent;
+                //         }
+                //         Response::Energy(energy) => {
+                //             let used_energy = energy - task_state.previous_energy_usage.min(energy);
+                //             task_state.previous_energy_usage = energy;
 
-                            task_state.budget -= used_energy.min(task_state.budget);
-                        }
-                    }
-                }
+                //             task_state.budget -= used_energy.min(task_state.budget);
+                //         }
+                //     }
+                // }
 
                 break;
             }
diff --git a/src/mock/perf.rs b/src/mock/perf.rs
index d35510d..3dde836 100644
--- a/src/mock/perf.rs
+++ b/src/mock/perf.rs
@@ -21,7 +21,10 @@ static EVENT_TYPES: &[(f32, Event)] = &[
 
 impl super::KernelModule for PerfEstimator {
     fn start_trace(&mut self, pid: u64) {
-        let mut group = Group::new().expect("failed to create Group");
+        let Ok(mut group) = Group::new() else {
+            println!("failed to create Group");
+            return;
+        };
         let counters = EVENT_TYPES
             .iter()
             .map(|(_, kind)| {
@@ -51,7 +54,7 @@ impl super::KernelModule for PerfEstimator {
         for ((factor, _ty), count) in EVENT_TYPES.iter().zip(counts.iter()) {
             sum += *factor as u64 * count.1;
         }
-        dbg!(sum);
+        // dbg!(sum);
         sum
     }
 }