Add energy logging (stdout for now)

1 files changed, 15 insertions, 2 deletions

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;
         }