Fix scheduler lockup

2 files changed, 68 insertions, 48 deletions

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
     }
 }