Added initial scheduler implementation

1 files changed, 106 insertions, 9 deletions

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