Log energy via socket

3 files changed, 114 insertions, 16 deletions

diff --git a/src/energy.rs b/src/energy.rs
index cb80047..c989685 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -9,6 +9,7 @@ use std::thread;
 use std::time::Duration;
 
 use crate::freq::FrequencyKHZ;
+use crate::socket;
 use crate::Pid;
 use dashmap::DashMap;
 
@@ -27,9 +28,9 @@ pub struct ProcessInfo {
     parent: Pid,
 }
 
-struct EnergyLog {
-    energy_of_exited_tasks: u64,
-    energy_of_running_tasks: u64,
+pub struct EnergyLog {
+    pub energy_of_exited_tasks: u64,
+    pub energy_of_running_tasks: u64,
 }
 
 pub struct EnergyService {
@@ -44,7 +45,7 @@ pub struct EnergyService {
     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>,
+    logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
 }
 
 impl EnergyService {
@@ -57,6 +58,7 @@ impl 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>>>,
+        logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
     ) -> Self {
         Self {
             estimator,
@@ -70,7 +72,7 @@ impl EnergyService {
             shared_policy_frequency_ranges,
             shared_cpu_current_frequencies,
             energy_logging_children_to_root_of_tree: HashMap::new(),
-            logged_energy: HashMap::new(),
+            logged_energy,
         }
     }
 
@@ -121,6 +123,8 @@ impl EnergyService {
                     {
                         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;
@@ -133,7 +137,7 @@ impl EnergyService {
                 Request::LogEnergyTask(pid) => {
                     self.energy_logging_children_to_root_of_tree
                         .insert(pid, pid);
-                    self.logged_energy.insert(
+                    self.logged_energy.write().unwrap().insert(
                         pid,
                         EnergyLog {
                             energy_of_exited_tasks: 0,
@@ -146,14 +150,15 @@ impl EnergyService {
     }
 
     fn update_measurements(&mut self) {
-        for (_, val) in self.logged_energy.iter_mut() {
-            val.energy_of_running_tasks = 0u64;
+        let mut logged_energy = self.logged_energy.write().unwrap();
+        for (_, value) in logged_energy.iter_mut() {
+            value.energy_of_running_tasks = 0u64;
         }
         for &pid in &self.active_processes {
             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
+                    logged_energy
                         .get_mut(root_pid)
                         .unwrap()
                         .energy_of_running_tasks += energy;
@@ -163,7 +168,7 @@ impl EnergyService {
                 {
                     self.energy_logging_children_to_root_of_tree
                         .insert(pid, root_pid);
-                    self.logged_energy
+                    logged_energy
                         .get_mut(&root_pid)
                         .unwrap()
                         .energy_of_running_tasks += energy;
@@ -172,12 +177,6 @@ impl EnergyService {
                 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) {
@@ -239,6 +238,8 @@ 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()));
+
     // Create and start the energy service
     let service = EnergyService::new(
         estimator,
@@ -249,9 +250,11 @@ pub fn start_energy_service(
         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))?;
 
     Ok(request_sender)
 }
diff --git a/src/main.rs b/src/main.rs
index c4bce3b..0caae55 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,6 +9,7 @@ mod energy;
 mod freq;
 mod model;
 mod scheduler;
+mod socket;
 
 #[rustfmt::skip]
 mod bpf;
diff --git a/src/socket.rs b/src/socket.rs
new file mode 100644
index 0000000..0503d1f
--- /dev/null
+++ b/src/socket.rs
@@ -0,0 +1,94 @@
+use std::{
+    collections::HashMap,
+    fs,
+    io::{self, Write},
+    os::unix::net::UnixListener,
+    path::Path,
+    sync::{Arc, RwLock},
+    thread,
+    time::Duration,
+};
+
+use crate::{energy::EnergyLog, Pid};
+
+pub struct LoggingSocketService {
+    path: String,
+    logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
+    update_interval: Duration,
+}
+
+impl LoggingSocketService {
+    pub fn new(
+        path: &str,
+        logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
+        update_interval: Duration,
+    ) -> 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));
+        }
+    }
+
+    pub fn run(self) {
+        thread::spawn(move || {
+            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 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}");
+                    line.clear();
+                }*/
+            }
+        });
+    }
+}
+
+pub fn start_logging_socket_service(
+    path: &str,
+    logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
+    update_interval: Duration,
+) -> io::Result<()> {
+    if Path::new(path).exists() {
+        fs::remove_file(path)?;
+    }
+    let socket = LoggingSocketService::new(path, logged_energy, update_interval);
+    socket.run();
+    Ok(())
+}