Added initial scheduler implementation

4 files changed, 150 insertions, 12 deletions

diff --git a/src/e_core_selector.rs b/src/e_core_selector.rs
new file mode 100644
index 0000000..76d8233
--- /dev/null
+++ b/src/e_core_selector.rs
@@ -0,0 +1,30 @@
+use std::ops::Range;
+
+pub trait ECoreSelector {
+    fn next_core(&mut self) -> i32;
+}
+
+//TODO: better core selection e.g. try to leave task on the same core
+pub struct RoundRobinSelector {
+    offset: u32,
+    num_cores: u32,
+    last_used: u32
+}
+
+impl RoundRobinSelector {
+    //TODO: use range instead
+    pub fn new(e_cores: Range<u32>) -> RoundRobinSelector {
+        Self {
+            offset: e_cores.start,
+            num_cores: e_cores.len() as u32,
+            last_used: 0,
+        }
+    }
+}
+
+impl ECoreSelector for RoundRobinSelector {
+    fn next_core(&mut self) -> i32 {
+        self.last_used += 1;
+        (self.last_used % self.num_cores) as i32
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 23b972d..9dcd496 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,18 +2,27 @@ mod bpf_skel;
 pub use bpf_skel::*;
 pub mod bpf_intf;
 
+mod mock;
+mod task_state;
+mod e_core_selector;
+
 #[rustfmt::skip]
 mod bpf;
 use bpf::*;
 
-mod mock;
-
+use e_core_selector::{ECoreSelector, RoundRobinSelector};
 use scx_utils::UserExitInfo;
 
 use libbpf_rs::OpenObject;
+use task_state::TaskState;
+use mock::{MockModule, KernelModule, KernelDriver};
 
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use std::mem::MaybeUninit;
+use std::time::SystemTime;
+use std::{env, u64};
+use std::process;
+use std::ops::Range;
 
 use anyhow::Result;
 
@@ -21,10 +30,22 @@ const SLICE_US: u64 = 5000;
 
 struct Scheduler<'a> {
     bpf: BpfScheduler<'a>,
+    module: Box<dyn KernelModule>,
     task_queue: VecDeque<QueuedTask>,
+    no_budget_task_queue: VecDeque<QueuedTask>,
+    managed_tasks: HashMap<u32, TaskState>,
+    maximum_budget: u64,
+    own_pid: u32,
+    p_cores: Range<i32>,
+    e_cores: Range<i32>,
+    e_core_selector: Box<dyn ECoreSelector>
 }
 
 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> {
         let bpf = BpfScheduler::init(
             open_object,
@@ -33,9 +54,19 @@ impl<'a> Scheduler<'a> {
             false, // debug (false = debug mode off)
         )?;
         dbg!("registering rust user space scheduler");
+        let module = Box::new(KernelDriver::default());
+        let selector = Box::new(RoundRobinSelector::new(8..16)); //TODO: get real value
         Ok(Self {
             bpf,
+            module,
             task_queue: VecDeque::new(),
+            no_budget_task_queue: VecDeque::new(),
+            managed_tasks: HashMap::new(),
+            maximum_budget: u64::MAX,
+            own_pid: process::id(),
+            p_cores: 0..8, //TODO: real value
+            e_cores: 8..16, //TODO: real value
+            e_core_selector: selector,
         })
     }
 
@@ -56,8 +87,37 @@ 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() {
-            self.task_queue.push_back(task);
+            // The scheduler itself has to be scheduled regardless of its energy usage
+            if task.pid as u32 == self.own_pid {
+                self.task_queue.push_back(task);
+                continue;
+            }
+
+            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);
+                } else {
+                    task_state.budget -= used_energy;
+                    self.task_queue.push_back(task);
+                }
+            } 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);
+            }
         }
     }
 
@@ -96,12 +156,32 @@ impl<'a> Scheduler<'a> {
 
             // Notify the BPF component of the number of pending tasks and immediately give a
             // chance to run to the dispatched task.
-            self.bpf.notify_complete(self.task_queue.len() as u64);
+            //TODO: should the pending tasks include the ones that don't have any budget left?
+            self.bpf.notify_complete(self.task_queue.len() as u64 + self.no_budget_task_queue.len() as u64);
+        } else if let Some(task) = self.no_budget_task_queue.pop_front() {
+            //TODO: do something with tasks that have no budget left
+            let mut dispatched_task = DispatchedTask::new(&task);
+
+            let cpu = self.e_core_selector.next_core();
+            if cpu >= 0 {
+                dispatched_task.cpu = cpu;
+            } else {
+                dispatched_task.flags |= RL_CPU_ANY as u64;
+            }
+
+            dispatched_task.slice_ns = SLICE_US;
+            self.bpf.dispatch_task(&dispatched_task).unwrap();
+            self.bpf.notify_complete(self.task_queue.len() as u64 + self.no_budget_task_queue.len() as u64);
         }
     }
 
     fn dispatch_tasks(&mut self) {
         loop {
+            //TODO: sometimes reset budget of tasks
+
+            //TODO: sometimes remove inactive tasks
+            //self.remove_inactive_tasks();
+
             // Consume all tasks before dispatching any.
             self.consume_all_tasks();
 
@@ -124,17 +204,34 @@ 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
+    }
 }
 
+//TODO: get children 
 fn main() -> Result<()> {
     // Initialize and load the FIFO scheduler.
     let mut open_object = MaybeUninit::uninit();
-    loop {
-        let mut sched = Scheduler::init(&mut open_object)?;
-        if !sched.run()?.should_restart() {
-            break;
+    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;
+            }
         }
     }
 
+
     Ok(())
 }
diff --git a/src/mock.rs b/src/mock.rs
index 4ee9ed4..3b3e3cf 100644
--- a/src/mock.rs
+++ b/src/mock.rs
@@ -2,13 +2,13 @@ use iocuddle::*;
 use libc::rand;
 use rand::Rng;
 
-trait KernelModule {
+pub trait KernelModule {
     fn start_trace(&mut self, _pid: u64) {}
     fn stop_trace(&mut self, _pid: u64) {}
     fn read_consumption(&mut self, pid: u64) -> u64;
 }
 
-struct MockModule;
+pub struct MockModule;
 
 impl KernelModule for MockModule {
     fn start_trace(&mut self, pid: u64) {
@@ -25,7 +25,7 @@ impl KernelModule for MockModule {
 }
 
 #[derive(Debug)]
-struct KernelDriver {
+pub struct KernelDriver {
     file: std::fs::File,
 }
 
diff --git a/src/task_state.rs b/src/task_state.rs
new file mode 100644
index 0000000..4f1bbb8
--- /dev/null
+++ b/src/task_state.rs
@@ -0,0 +1,11 @@
+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