Get children

2 files changed, 30 insertions, 5 deletions

diff --git a/src/main.rs b/src/main.rs
index 36f22eb..6cbc1fa 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,7 @@ use libbpf_rs::OpenObject;
 use task_state::TaskState;
 use mock::{MockModule, KernelModule, KernelDriver};
 
+use sysinfo::System;
 use std::collections::{HashMap, VecDeque};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
@@ -40,7 +41,8 @@ struct Scheduler<'a> {
     own_pid: u32,
     p_cores: Range<i32>,
     e_cores: Range<i32>,
-    e_core_selector: Box<dyn ECoreSelector>
+    e_core_selector: Box<dyn ECoreSelector>,
+    system: System,
 }
 
 impl<'a> Scheduler<'a> {
@@ -77,6 +79,10 @@ impl<'a> Scheduler<'a> {
         };
 
         let selector = Box::new(RoundRobinSelector::new(&e_cores));
+
+        let mut system = System::new_all();
+        system.refresh_processes(sysinfo::ProcessesToUpdate::All, true);
+
         Ok(Self {
             bpf,
             module,
@@ -89,6 +95,7 @@ impl<'a> Scheduler<'a> {
             p_cores,
             e_cores,
             e_core_selector: selector,
+            system
         })
     }
 
@@ -195,9 +202,11 @@ impl<'a> Scheduler<'a> {
 
     fn dispatch_tasks(&mut self) {
         loop {
-            //TODO: we should probably not do this every time, and instead wait a predefined amount of time
+            //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();
-
+            //TODO: we should probably not do this every time, but instead wait a predefined amount of time between invocations
+            self.refresh_children();
+            
             // Consume all tasks before dispatching any.
             self.consume_all_tasks();
 
@@ -240,9 +249,25 @@ impl<'a> Scheduler<'a> {
             was_scheduled
         });
     }
+
+    //TODO: this is probably not very efficient
+    fn refresh_children(&mut self) {
+        self.system.refresh_processes(sysinfo::ProcessesToUpdate::All, true);
+        for entry in self.managed_tasks.values_mut() {
+            entry.children.clear();
+        }
+
+        for (pid, process) in self.system.processes() {
+            if let Some(parent_pid) = process.parent() {
+                if let Some(entry) = self.managed_tasks.get_mut(&parent_pid.as_u32()) {
+                    entry.children.push(pid.as_u32());
+                }
+            }
+        }
+    }
+
 }
 
-//TODO: get children 
 fn main() -> Result<()> {
     let matches = Command::new("Energy User Space Scheduler")
         .arg(Arg::new("mock")
diff --git a/src/task_state.rs b/src/task_state.rs
index 3d1db4c..e56b109 100644
--- a/src/task_state.rs
+++ b/src/task_state.rs
@@ -1,5 +1,5 @@
 pub struct TaskState {
     pub previous_energy_usage: u64,
     pub budget: u64,
-    pub children: Vec<u64>,
+    pub children: Vec<u32>,
 }
\ No newline at end of file