Update logging via socket

5 files changed, 81 insertions, 148 deletions

diff --git a/src/energy.rs b/src/energy.rs
index 1b52c13..02aaf62 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -22,37 +22,30 @@ pub use trackers::{KernelDriver, PerfEstimator};
 pub enum Request {
     NewTask(Pid, i32),
     RemoveTask(Pid),
-    LogEnergyTask(Pid),
 }
 
 #[derive(Clone)]
 pub struct ProcessInfo {
-    energy: f64,
-    tree_energy: f64,
-    last_update: std::time::Instant,
-    parent: Pid,
-    cpu: i32,
-}
-
-pub struct EnergyLog {
-    pub energy_of_exited_tasks: f64,
-    pub energy_of_running_tasks: f64,
+    pub energy: f64,
+    pub tree_energy: f64,
+    pub last_update: std::time::Instant,
+    pub parent: Pid,
+    pub cpu: i32,
 }
 
 pub struct EnergyService {
     estimator: Box<dyn Estimator>,
     budget_policy: Option<Box<dyn BudgetPolicy>>,
-    // TODO: remove
+    // contains the same data as the keys of process_info but having this reduces contention and
+    // avoids unnecessary clone
     active_processes: BTreeSet<Pid>,
-    process_info: HashMap<Pid, ProcessInfo>,
+    process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
     shared_budgets: Arc<DashMap<Pid, u64>>,
     request_receiver: mpsc::Receiver<Request>,
     update_interval: Duration,
     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: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
     rapl_offset: f64,
     old_rapl: f64,
     bias: f64,
@@ -63,27 +56,25 @@ impl EnergyService {
     pub fn new(
         estimator: Box<dyn Estimator>,
         budget_policy: Box<dyn BudgetPolicy>,
+        process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
         shared_budgets: Arc<DashMap<Pid, u64>>,
         request_receiver: mpsc::Receiver<Request>,
         update_interval: Duration,
         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>>>,
-        logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
     ) -> Self {
         Self {
             estimator,
             budget_policy: Some(budget_policy),
             active_processes: BTreeSet::new(),
-            process_info: HashMap::new(),
+            process_info,
             shared_budgets,
             request_receiver,
             update_interval,
             shared_cpu_frequency_ranges,
             shared_policy_frequency_ranges,
             shared_cpu_current_frequencies,
-            energy_logging_children_to_root_of_tree: HashMap::new(),
-            logged_energy,
             rapl_offset: rapl::read_package_energy().unwrap(),
             old_rapl: 0.,
             bias: 1.,
@@ -118,17 +109,16 @@ impl EnergyService {
     fn handle_request(&mut self, request: Request) {
         match request {
             Request::NewTask(pid, cpu) => {
-                if self.process_info.contains_key(&pid) {
+                if self.process_info.read().unwrap().contains_key(&pid) {
                     return;
                 }
                 self.estimator.start_trace(pid as u64);
-                self.active_processes.insert(pid);
                 let parent = (|| {
                     let process = procfs::process::Process::new(pid)?;
                     process.stat().map(|stat| stat.ppid)
                 })()
                 .unwrap_or_default();
-                self.process_info.insert(
+                self.process_info.write().unwrap().insert(
                     pid,
                     ProcessInfo {
                         energy: 0.,
@@ -138,9 +128,10 @@ impl EnergyService {
                         cpu,
                     },
                 );
+                self.active_processes.insert(pid);
                 // Initialize with default budget
                 self.shared_budgets.insert(pid, u64::MAX);
-                if !self.process_info.contains_key(&parent) && parent != 0 {
+                if !self.process_info.read().unwrap().contains_key(&parent) && parent != 0 {
                     self.handle_request(Request::NewTask(parent, cpu));
                 }
             }
@@ -150,46 +141,28 @@ impl EnergyService {
                 }
 
                 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
-                            .write()
-                            .unwrap()
-                            .get_mut(&root_pid)
-                            .unwrap()
-                            .energy_of_exited_tasks += info.energy;
-                    }
-                }
-                self.active_processes.remove(&pid);
-                self.process_info.remove(&pid);
+                self.process_info.write().unwrap().remove(&pid);
+                self.process_info.write().unwrap().remove(&pid);
                 self.shared_budgets.remove(&pid);
             }
-            Request::LogEnergyTask(pid) => {
-                self.energy_logging_children_to_root_of_tree
-                    .insert(pid, pid);
-                self.logged_energy.write().unwrap().insert(
-                    pid,
-                    EnergyLog {
-                        energy_of_exited_tasks: 0.,
-                        energy_of_running_tasks: 0.,
-                    },
-                );
-            }
         }
     }
 
     fn update_measurements(&mut self) {
-        let mut logged_energy = self.logged_energy.write().unwrap();
-        for (_, value) in logged_energy.iter_mut() {
-            value.energy_of_running_tasks = 0f64;
-        }
-        let old_energy = self.process_info.get(&1).unwrap().tree_energy;
-        for &pid in &self.active_processes {
-            if let Some(info) = self.process_info.get_mut(&pid) {
-                if let Some(energy) = self.estimator.read_consumption(pid as u64, info.cpu) {
+        let old_energy = self
+            .process_info
+            .read()
+            .unwrap()
+            .get(&1)
+            .unwrap()
+            .tree_energy;
+        for pid in &self.active_processes {
+            let mut process_info = self.process_info.write().unwrap();
+            if let Some(info) = process_info.get_mut(&pid) {
+                if let Some(energy) = self.estimator.read_consumption(*pid as u64, info.cpu) {
                     info.energy += energy * self.bias;
                     let mut parent = info.parent;
-                    while let Some(info) = self.process_info.get_mut(&parent) {
+                    while let Some(info) = process_info.get_mut(&parent) {
                         info.tree_energy += energy * self.bias;
                         info.last_update = std::time::Instant::now();
                         parent = info.parent;
@@ -197,33 +170,23 @@ impl EnergyService {
                 }
             }
         }
-        if let Some(init) = self.process_info.get_mut(&1) {
+        if let Some(init) = self.process_info.write().unwrap().get_mut(&1) {
             let rapl = rapl::read_package_energy().unwrap() - self.rapl_offset;
             let rapl_diff = rapl - self.old_rapl;
-            let est_diff = (init.tree_energy - old_energy);
+            let est_diff = init.tree_energy - old_energy;
             if est_diff < 0.1 {
                 self.offset = (self.offset + (rapl_diff - est_diff)) * 0.5;
             }
             self.old_rapl = rapl;
-            init.tree_energy = (init.tree_energy + self.offset);
+            init.tree_energy = init.tree_energy + self.offset;
             let offset_bias = (rapl / init.tree_energy).clamp(0.1, 2.);
             let diff_bias = (rapl_diff / est_diff).clamp(0.1, 2.);
             let current_bias = (offset_bias + diff_bias) * 0.5;
-            //dbg!(current_bias, self.bias);
             self.bias = (self.bias * ((1. / 3.) * current_bias + (2. / 3.))).clamp(0.1, 20.);
-            //let v = logged_energy.get(&1).unwrap();
             println!(
                 "Energy estimation: {:.1} rapl: {:.1}, est diff: {:.1} rapl diff: {:.1}",
-                init.tree_energy,
-                rapl,
-                est_diff,
-                rapl_diff,
-                //(rapl_diff - (init.tree_energy - old_energy)) / rapl_diff * 100.,
+                init.tree_energy, rapl, est_diff, rapl_diff,
             );
-            // if init.tree_energy > 1000. || rapl > 1000. {
-            //     self.rapl_offset += rapl;
-            //     init.tree_energy = 0.;
-            // }
         }
     }
 
@@ -251,13 +214,19 @@ impl EnergyService {
     }
 
     pub fn process_energy(&self, pid: Pid) -> Option<f64> {
-        self.process_info.get(&pid).map(|info| info.energy)
+        self.process_info
+            .read()
+            .unwrap()
+            .get(&pid)
+            .map(|info| info.energy)
     }
 
     pub fn all_process_energies(&self) -> HashMap<Pid, f64> {
         self.process_info
+            .read()
+            .unwrap()
             .iter()
-            .map(|(&pid, info)| (pid, info.energy))
+            .map(|(&key, info)| (key, info.energy))
             .collect()
     }
 }
@@ -286,23 +255,23 @@ pub fn start_energy_service(
     // shouldn't be a problem because we are privileged
     // if PackageEnergy::check_paranoid().unwrap_or(3) > 0 {}
 
-    let logged_energy = Arc::new(RwLock::new(HashMap::new()));
+    let process_info = Arc::new(RwLock::new(HashMap::new()));
 
     // Create and start the energy service
     let service = EnergyService::new(
         estimator,
         budget_policy,
+        process_info.clone(),
         shared_budgets,
         request_receiver,
         Duration::from_millis(50), // 50ms update interval
         shared_cpu_frequency_ranges,
         shared_policy_frequency_ranges,
         shared_cpu_current_frequencies,
-        logged_energy.clone(),
     );
 
     service.run();
-    //socket::start_logging_socket_service("/tmp/pm-sched", logged_energy, Duration::from_secs(1))?;
+    socket::start_logging_socket_service("/tmp/pm-sched", process_info.clone())?;
 
     Ok(request_sender)
 }
diff --git a/src/energy/budget.rs b/src/energy/budget.rs
index 7002dd3..6154187 100644
--- a/src/energy/budget.rs
+++ b/src/energy/budget.rs
@@ -42,8 +42,8 @@ impl BudgetPolicy for SimpleCappingPolicy {
             }
         } else {
             // Under power cap, assign maximum budget to all
-            for &pid in energy_service.active_processes() {
-                budgets.insert(pid, u64::MAX);
+            for pid in energy_service.active_processes() {
+                budgets.insert(*pid, u64::MAX);
             }
         }
 
diff --git a/src/main.rs b/src/main.rs
index 6242917..c56a41c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -47,23 +47,11 @@ 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 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();
@@ -74,8 +62,7 @@ fn main() -> Result<()> {
         return Ok(());
     }
     loop {
-        let mut sched =
-            Scheduler::init(&mut open_object, use_mocking, power_cap, &pids_to_monitor)?;
+        let mut sched = Scheduler::init(&mut open_object, use_mocking, power_cap)?;
         if !sched.run()?.should_restart() {
             break;
         }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index fc4d448..359b9da 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -45,7 +45,6 @@ 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,
@@ -76,12 +75,6 @@ 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();
@@ -115,9 +108,6 @@ impl<'a> Scheduler<'a> {
         frequency_sender
             .try_send(FrequencyRequest::GetPossibleCPUFrequencyRange)
             .unwrap();
-        frequency_sender
-            .try_send(FrequencyRequest::GetCurrentFrequencies)
-            .unwrap();
 
         Ok(Self {
             bpf,
diff --git a/src/socket.rs b/src/socket.rs
index 0503d1f..0b42abe 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     fs,
-    io::{self, Write},
+    io::{self, BufRead, BufReader, Write},
     os::unix::net::UnixListener,
     path::Path,
     sync::{Arc, RwLock},
@@ -9,44 +9,20 @@ use std::{
     time::Duration,
 };
 
-use crate::{energy::EnergyLog, Pid};
+use dashmap::DashMap;
+
+use crate::{energy::ProcessInfo, Pid};
 
 pub struct LoggingSocketService {
     path: String,
-    logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
-    update_interval: Duration,
+    process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
 }
 
 impl LoggingSocketService {
-    pub fn new(
-        path: &str,
-        logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
-        update_interval: Duration,
-    ) -> Self {
+    pub fn new(path: &str, process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>) -> Self {
         LoggingSocketService {
             path: path.to_string(),
-            logged_energy,
-            update_interval,
-        }
-    }
-
-    fn generate_report(&self) -> String {
-        loop {
-            if let Ok(logged_energy) = self.logged_energy.try_read() {
-                let mut result = logged_energy
-                    .iter()
-                    .map(|(key, value)| {
-                        format!(
-                            "{key}: {}J",
-                            value.energy_of_running_tasks + value.energy_of_exited_tasks
-                        )
-                    })
-                    .collect::<Vec<_>>()
-                    .join(", ");
-                result += "\n";
-                break result;
-            }
-            thread::sleep(Duration::from_millis(10));
+            process_info,
         }
     }
 
@@ -55,26 +31,38 @@ impl LoggingSocketService {
             let listener = UnixListener::bind(self.path.clone()).unwrap();
             for stream in listener.incoming() {
                 let mut socket = stream.unwrap();
-                loop {
-                    if socket
-                        .write_all(format!("{}", self.generate_report()).as_bytes())
-                        .is_err()
-                    {
-                        break;
-                    }
-                    thread::sleep(self.update_interval);
-                }
-                //TODO: if we ever want to read too
-                /*let mut reader = BufReader::new(socket.try_clone().unwrap());
+                let mut reader = BufReader::new(socket.try_clone().unwrap());
                 let mut line = String::new();
                 while let Ok(bytes_read) = reader.read_line(&mut line) {
                     if bytes_read == 0 {
                         break;
                     }
-                    socket.write_all(format!("got: {line}").as_bytes());
-                    println!("got: {line}");
+                    if let Ok(pid) = line.trim().parse::<Pid>() {
+                        if let Some(info) = self.process_info.read().unwrap().get(&pid).clone() {
+                            socket
+                                .write_all(
+                                    format!(
+                                        "pid: {pid} process: {}J process tree: {}J\n",
+                                        info.energy, info.tree_energy
+                                    )
+                                    .as_bytes(),
+                                )
+                                .unwrap();
+                        } else {
+                            socket
+                                .write_all(format!("Unknown pid: {pid}\n").as_bytes())
+                                .unwrap();
+                        }
+                    } else {
+                        socket
+                            .write_all(
+                                format!("Failed to get energy for pid: {}\n", line.trim_end())
+                                    .as_bytes(),
+                            )
+                            .unwrap();
+                    }
                     line.clear();
-                }*/
+                }
             }
         });
     }
@@ -82,13 +70,12 @@ impl LoggingSocketService {
 
 pub fn start_logging_socket_service(
     path: &str,
-    logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
-    update_interval: Duration,
+    process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
 ) -> io::Result<()> {
     if Path::new(path).exists() {
         fs::remove_file(path)?;
     }
-    let socket = LoggingSocketService::new(path, logged_energy, update_interval);
+    let socket = LoggingSocketService::new(path, process_info);
     socket.run();
     Ok(())
 }