Cleanup, argument parsing, garbage collection

3 files changed, 51 insertions, 40 deletions

diff --git a/src/main.rs b/src/main.rs
index 9dcd496..36de0bc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,6 +10,7 @@ mod e_core_selector;
 mod bpf;
 use bpf::*;
 
+use clap::{Arg, ArgAction, Command};
 use e_core_selector::{ECoreSelector, RoundRobinSelector};
 use scx_utils::UserExitInfo;
 
@@ -19,8 +20,7 @@ use mock::{MockModule, KernelModule, KernelDriver};
 
 use std::collections::{HashMap, VecDeque};
 use std::mem::MaybeUninit;
-use std::time::SystemTime;
-use std::{env, u64};
+use std::u64;
 use std::process;
 use std::ops::Range;
 
@@ -35,6 +35,7 @@ struct Scheduler<'a> {
     no_budget_task_queue: VecDeque<QueuedTask>,
     managed_tasks: HashMap<u32, TaskState>,
     maximum_budget: u64,
+    power_cap: u64,
     own_pid: u32,
     p_cores: Range<i32>,
     e_cores: Range<i32>,
@@ -42,11 +43,7 @@ struct Scheduler<'a> {
 }
 
 impl<'a> Scheduler<'a> {
-    fn init_mock(open_object: &'a mut MaybeUninit<OpenObject>) -> Result<Self> {
-        todo!();
-    }
-
-    fn init(open_object: &'a mut MaybeUninit<OpenObject>) -> Result<Self> {
+    fn init(open_object: &'a mut MaybeUninit<OpenObject>, use_mocking: bool, power_cap: u64) -> Result<Self> {
         let bpf = BpfScheduler::init(
             open_object,
             0,     // exit_dump_len (buffer size of exit info, 0 = default)
@@ -54,7 +51,11 @@ impl<'a> Scheduler<'a> {
             false, // debug (false = debug mode off)
         )?;
         dbg!("registering rust user space scheduler");
-        let module = Box::new(KernelDriver::default());
+        let module: Box<dyn KernelModule> = if use_mocking {
+            Box::new(MockModule)
+        } else { 
+            Box::new(KernelDriver::default())
+        };
         let selector = Box::new(RoundRobinSelector::new(8..16)); //TODO: get real value
         Ok(Self {
             bpf,
@@ -63,6 +64,7 @@ impl<'a> Scheduler<'a> {
             no_budget_task_queue: VecDeque::new(),
             managed_tasks: HashMap::new(),
             maximum_budget: u64::MAX,
+            power_cap,
             own_pid: process::id(),
             p_cores: 0..8, //TODO: real value
             e_cores: 8..16, //TODO: real value
@@ -87,7 +89,6 @@ 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 current_time = SystemTime::now();
 
         while let Ok(Some(task)) = self.bpf.dequeue_task() {
             // The scheduler itself has to be scheduled regardless of its energy usage
@@ -99,7 +100,6 @@ impl<'a> Scheduler<'a> {
             if let Some(task_state) = self.managed_tasks.get_mut(&(task.pid as u32)) {
                 let energy = self.module.read_consumption(task.pid as u64);
                 let used_energy = energy - task_state.previous_energy_usage;
-                task_state.scheduled_last = current_time;
                 if task_state.budget < used_energy {
                     task_state.budget = 0;
                     self.no_budget_task_queue.push_back(task);
@@ -110,11 +110,9 @@ impl<'a> Scheduler<'a> {
             } else {
                 self.module.start_trace(task.pid as u64);
                 self.managed_tasks.insert(task.pid as u32, TaskState {
-                    pid: task.pid as u64,
                     previous_energy_usage: 0,
                     budget: self.maximum_budget,
                     children: vec![],
-                    scheduled_last: current_time,
                 });
                 self.task_queue.push_back(task);
             }
@@ -177,10 +175,8 @@ impl<'a> Scheduler<'a> {
 
     fn dispatch_tasks(&mut self) {
         loop {
-            //TODO: sometimes reset budget of tasks
-
-            //TODO: sometimes remove inactive tasks
-            //self.remove_inactive_tasks();
+            //TODO: we should probably not do this every time but after a predefined amount of time
+            self.reset_budgets_and_garbage_collect();
 
             // Consume all tasks before dispatching any.
             self.consume_all_tasks();
@@ -205,30 +201,52 @@ impl<'a> Scheduler<'a> {
         self.bpf.shutdown_and_report()
     }
 
-    fn remove_inactive_tasks(&mut self) {
-        //TODO: remove tasks which have not been scheduled for a certain period of time
+    fn get_new_budget(&self) -> u64 {
+        //TODO
+        u64::MAX
+    }
+
+    fn reset_budgets_and_garbage_collect(&mut self) {
+        let old_budget = self.maximum_budget;
+        self.maximum_budget = self.get_new_budget();
+
+        self.managed_tasks.retain(|key, value| {
+            let was_scheduled = value.budget == old_budget;
+            if was_scheduled {
+                value.budget = self.maximum_budget;
+            } else {
+                self.module.stop_trace(*key as u64);
+            }
+            was_scheduled
+        });
     }
 }
 
 //TODO: get children 
 fn main() -> Result<()> {
+    let matches = Command::new("Energy User Space Scheduler")
+        .arg(Arg::new("mock")
+                .short('m')
+                .long("mock")
+                .help("Use this flag if to activate kernel module mocking")
+                .action(ArgAction::SetTrue)
+                .required(false))
+        .arg(Arg::new("power_cap")
+                .long("energy_cap")
+                .help("Set a power cap for the processor")
+                .required(false)
+                .value_name("power in watts"))
+        .get_matches();
+
+    let power_cap = *matches.get_one::<u64>("power_cap").unwrap_or(&u64::MAX);
+    let use_mocking = matches.get_flag("mock");
+
     // Initialize and load the FIFO scheduler.
     let mut open_object = MaybeUninit::uninit();
-    let args: Vec<String> = env::args().collect();
-    //TODO: get power cap from parameter or config file and set maximum budget accordingly
-    if args[1] != "mock" {
-        loop {
-            let mut sched = Scheduler::init(&mut open_object)?;
-            if !sched.run()?.should_restart() {
-                break;
-            }
-        }
-    } else {
-        loop {
-            let mut sched = Scheduler::init_mock(&mut open_object)?;
-            if !sched.run()?.should_restart() {
-                break;
-            }
+    loop {
+        let mut sched = Scheduler::init(&mut open_object, use_mocking, power_cap)?;
+        if !sched.run()?.should_restart() {
+            break;
         }
     }
 
diff --git a/src/mock.rs b/src/mock.rs
index 17d20f4..81bc32c 100644
--- a/src/mock.rs
+++ b/src/mock.rs
@@ -1,5 +1,4 @@
 use iocuddle::*;
-use libc::rand;
 use rand::Rng;
 
 pub trait KernelModule {
diff --git a/src/task_state.rs b/src/task_state.rs
index 4f1bbb8..3d1db4c 100644
--- a/src/task_state.rs
+++ b/src/task_state.rs
@@ -1,11 +1,5 @@
-use std::time::SystemTime;
-
-use crate::bpf::QueuedTask;
-
 pub struct TaskState {
-    pub pid: u64,
     pub previous_energy_usage: u64,
     pub budget: u64,
     pub children: Vec<u64>,
-    pub scheduled_last: SystemTime,
 }
\ No newline at end of file