Add energy logging (stdout for now)

3 files changed, 79 insertions, 3 deletions

diff --git a/src/energy.rs b/src/energy.rs
index a26a2b3..cb80047 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -18,6 +18,7 @@ pub use trackers::{Estimator, KernelDriver, PerfEstimator};
 pub enum Request {
     NewTask(Pid),
     RemoveTask(Pid),
+    LogEnergyTask(Pid),
 }
 
 pub struct ProcessInfo {
@@ -26,6 +27,11 @@ pub struct ProcessInfo {
     parent: Pid,
 }
 
+struct EnergyLog {
+    energy_of_exited_tasks: u64,
+    energy_of_running_tasks: u64,
+}
+
 pub struct EnergyService {
     estimator: Box<dyn Estimator>,
     budget_policy: Option<Box<dyn BudgetPolicy>>,
@@ -37,6 +43,8 @@ pub struct EnergyService {
     shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
     shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
     shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
+    energy_logging_children_to_root_of_tree: HashMap<Pid, Pid>,
+    logged_energy: HashMap<Pid, EnergyLog>,
 }
 
 impl EnergyService {
@@ -61,6 +69,8 @@ impl EnergyService {
             shared_cpu_frequency_ranges,
             shared_policy_frequency_ranges,
             shared_cpu_current_frequencies,
+            energy_logging_children_to_root_of_tree: HashMap::new(),
+            logged_energy: HashMap::new(),
         }
     }
 
@@ -106,22 +116,68 @@ impl EnergyService {
                 }
                 Request::RemoveTask(pid) => {
                     self.estimator.stop_trace(pid as u64);
+                    if let Some(root_pid) =
+                        self.energy_logging_children_to_root_of_tree.remove(&pid)
+                    {
+                        if let Some(info) = self.process_info.remove(&pid) {
+                            self.logged_energy
+                                .get_mut(&root_pid)
+                                .unwrap()
+                                .energy_of_exited_tasks += info.energy;
+                        }
+                    }
                     self.active_processes.remove(&pid);
                     self.process_info.remove(&pid);
                     self.shared_budgets.remove(&pid);
                 }
+                Request::LogEnergyTask(pid) => {
+                    self.energy_logging_children_to_root_of_tree
+                        .insert(pid, pid);
+                    self.logged_energy.insert(
+                        pid,
+                        EnergyLog {
+                            energy_of_exited_tasks: 0,
+                            energy_of_running_tasks: 0,
+                        },
+                    );
+                }
             }
         }
     }
 
     fn update_measurements(&mut self) {
+        for (_, val) in self.logged_energy.iter_mut() {
+            val.energy_of_running_tasks = 0u64;
+        }
         for &pid in &self.active_processes {
-            let energy = self.estimator.read_consumption(pid as u64);
             if let Some(info) = self.process_info.get_mut(&pid) {
+                let energy = self.estimator.read_consumption(pid as u64);
+                if let Some(root_pid) = self.energy_logging_children_to_root_of_tree.get(&pid) {
+                    self.logged_energy
+                        .get_mut(root_pid)
+                        .unwrap()
+                        .energy_of_running_tasks += energy;
+                } else if let Some(&root_pid) = self
+                    .energy_logging_children_to_root_of_tree
+                    .get(&info.parent)
+                {
+                    self.energy_logging_children_to_root_of_tree
+                        .insert(pid, root_pid);
+                    self.logged_energy
+                        .get_mut(&root_pid)
+                        .unwrap()
+                        .energy_of_running_tasks += energy;
+                }
                 info.energy = energy;
                 info.last_update = std::time::Instant::now();
             }
         }
+        for (pid, log) in self.logged_energy.iter() {
+            println!(
+                "pid: {pid} energy: {}",
+                log.energy_of_exited_tasks + log.energy_of_running_tasks
+            );
+        }
     }
 
     fn update_budgets(&mut self) {
diff --git a/src/main.rs b/src/main.rs
index 4f7d7a4..c4bce3b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -46,6 +46,13 @@ fn main() -> Result<()> {
                 .required(false)
                 .value_name("power in watts"),
         )
+        .arg(
+            Arg::new("log_energy")
+                .long("log_energy")
+                .help("Provide one or more PIDs to log their process tress energy")
+                .num_args(1..)
+                .required(false),
+        )
         .get_matches();
 
     let model = model::load_model();
@@ -64,17 +71,23 @@ fn main() -> Result<()> {
     let power_cap = *matches.get_one::<u64>("power_cap").unwrap_or(&u64::MAX);
     let use_mocking = matches.get_flag("mock");
     let benchmark = matches.get_flag("benchmark");
+    let pids_to_monitor: Vec<Pid> = matches
+        .get_many::<String>("log_energy")
+        .unwrap_or_default()
+        .map(|v| v.parse::<Pid>().unwrap())
+        .collect();
 
     // Initialize and load the scheduler.
     let mut open_object = MaybeUninit::uninit();
     let log_path = "/tmp/logs.csv";
     if benchmark {
         let mut sched = BenchmarkScheduler::init(&mut open_object, log_path)?;
-        sched.run();
+        sched.run()?;
         return Ok(());
     }
     loop {
-        let mut sched = Scheduler::init(&mut open_object, use_mocking, power_cap)?;
+        let mut sched =
+            Scheduler::init(&mut open_object, use_mocking, power_cap, &pids_to_monitor)?;
         if !sched.run()?.should_restart() {
             break;
         }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 39db5c6..a9cad61 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -44,6 +44,7 @@ impl<'a> Scheduler<'a> {
         open_object: &'a mut MaybeUninit<OpenObject>,
         use_mocking: bool,
         power_cap: u64,
+        pids_to_monitor: &[Pid],
     ) -> Result<Self> {
         let bpf = BpfScheduler::init(
             open_object,
@@ -74,6 +75,12 @@ impl<'a> Scheduler<'a> {
             shared_cpu_current_frequencies.clone(),
         )?;
 
+        for pid in pids_to_monitor {
+            energy_sender
+                .try_send(EnergyRequest::LogEnergyTask(*pid))
+                .unwrap();
+        }
+
         let topology = Topology::new().unwrap();
         let mut e_core_ids = Vec::new();
         let mut p_core_ids = Vec::new();