Allow setting energy limit via socket

7 files changed, 80 insertions, 29 deletions

diff --git a/energy-monitor/profile.json.gz b/energy-monitor/profile.json.gz
new file mode 100644
index 0000000..a6b0871
--- /dev/null
+++ b/energy-monitor/profile.json.gz
diff --git a/energy-monitor/src/energy.rs b/energy-monitor/src/energy.rs
index 808b990..a0a6ce1 100644
--- a/energy-monitor/src/energy.rs
+++ b/energy-monitor/src/energy.rs
@@ -45,7 +45,7 @@ pub fn request_all_processes(
     stream: &mut std::os::unix::net::UnixStream,
 ) -> io::Result<HashMap<Pid, (f64, f64)>> {
     // Write -1 to get all processes
-    stream.write_all(b"-1\n")?;
+    stream.write_all(b"list\n")?;
 
     let mut reader = BufReader::new(stream);
     let mut data = Vec::new();
@@ -69,3 +69,26 @@ pub fn request_all_processes(
 
     Ok(result)
 }
+pub fn request_rapl(stream: &mut std::os::unix::net::UnixStream) -> io::Result<f64> {
+    // Write -1 to get all processes
+    stream.write_all(b"rapl\n")?;
+
+    let mut reader = BufReader::new(stream);
+
+    let mut out = String::new();
+    reader.read_line(&mut out)?;
+
+    let number = out.parse().unwrap_or_default();
+    Ok(number)
+}
+pub fn set_power_limit(stream: &mut std::os::unix::net::UnixStream, limit: u32) -> io::Result<()> {
+    // Write -1 to get all processes
+    stream.write_all(format!("limit {limit}\n").as_bytes())?;
+
+    let mut reader = BufReader::new(stream);
+    let mut data = Vec::new();
+
+    reader.read_until(b'#', &mut data)?;
+
+    Ok(())
+}
diff --git a/energy-monitor/src/main.rs b/energy-monitor/src/main.rs
index e608d72..7c65e2d 100644
--- a/energy-monitor/src/main.rs
+++ b/energy-monitor/src/main.rs
@@ -164,10 +164,12 @@ impl App {
 
     fn increase_power_limit(&mut self) {
         self.power_limit += POWER_LIMIT_STEP;
+        energy::set_power_limit(&mut self.process_data.socket, self.power_limit as u32).unwrap();
     }
 
     fn decrease_power_limit(&mut self) {
         self.power_limit = (self.power_limit - POWER_LIMIT_STEP).max(0.0);
+        energy::set_power_limit(&mut self.process_data.socket, self.power_limit as u32).unwrap();
     }
 }
 
diff --git a/energy-monitor/src/process.rs b/energy-monitor/src/process.rs
index 363f478..64b278a 100644
--- a/energy-monitor/src/process.rs
+++ b/energy-monitor/src/process.rs
@@ -1,4 +1,5 @@
 use std::{collections::HashMap, os::unix::net::UnixStream, time::Instant};
+use std::{fs, io};
 
 use procfs::process::Process;
 
@@ -116,3 +117,16 @@ fn get_process_name(pid: Pid) -> Option<String> {
         Err(_) => None,
     }
 }
+
+/// Read current package energy counter value in joules
+pub fn read_package_energy() -> io::Result<f64> {
+    let energy_path = "/sys/class/powercap/intel-rapl/intel-rapl:0/energy_uj";
+    let energy_str = fs::read_to_string(energy_path)?;
+    let energy_uj = energy_str
+        .trim()
+        .parse::<u64>()
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
+    // Convert from microjoules to joules
+    Ok(energy_uj as f64 / 1_000_000.0)
+}
diff --git a/src/energy.rs b/src/energy.rs
index 3f6931d..21283b8 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -384,7 +384,7 @@ pub fn start_energy_service(
     );
 
     service.run();
-    socket::start_logging_socket_service("/tmp/pm-sched", process_info.clone())?;
+    let budget = socket::start_logging_socket_service("/tmp/pm-sched", process_info.clone())?;
 
     Ok(request_sender)
 }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 4d4e987..c34e151 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -84,8 +84,8 @@ impl<'a> Scheduler<'a> {
         // We assume that the CPU IDs for each core type are assigned contiguously.
         e_core_ids.sort();
         p_core_ids.sort();
-        let e_cores = *e_core_ids.first().unwrap_or(&0)..(*e_core_ids.last().unwrap_or(&0) + 1);
-        let p_cores = *p_core_ids.first().unwrap_or(&0)..(*p_core_ids.last().unwrap_or(&0) + 1);
+        let e_cores = *e_core_ids.first().unwrap_or(&0)..(*e_core_ids.last().unwrap_or(&-1) + 1);
+        let p_cores = *p_core_ids.first().unwrap_or(&0)..(*p_core_ids.last().unwrap_or(&-1) + 1);
         let all_cores = 0..((e_cores.len() + p_cores.len()) as u32);
 
         let selector = Box::new(RoundRobinSelector::new(&e_cores));
diff --git a/src/socket.rs b/src/socket.rs
index e2ecdbf..2bb0dbb 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -4,7 +4,7 @@ use std::{
     io::{self, BufRead, BufReader, Write},
     os::unix::net::UnixListener,
     path::Path,
-    sync::{Arc, RwLock},
+    sync::{atomic::AtomicU32, Arc, RwLock},
     thread,
 };
 
@@ -12,14 +12,20 @@ use crate::{energy::ProcessInfo, Pid};
 
 pub struct LoggingSocketService {
     path: String,
+    power_limit: Arc<AtomicU32>,
     process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
 }
 
 impl LoggingSocketService {
-    pub fn new(path: &str, process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>) -> Self {
+    pub fn new(
+        path: &str,
+        process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
+        power_limit: Arc<AtomicU32>,
+    ) -> Self {
         LoggingSocketService {
             path: path.to_string(),
             process_info,
+            power_limit,
         }
     }
 
@@ -34,35 +40,40 @@ impl LoggingSocketService {
                     if bytes_read == 0 {
                         break;
                     }
-                    if let Ok(pid) = line.trim().parse::<Pid>() {
-                        socket
-                            .write_all(self.handle_request(pid).as_bytes())
-                            .unwrap();
-                    } else {
-                        socket
-                            .write_all(
-                                format!("Failed to get energy for pid: {}\n", line.trim_end())
-                                    .as_bytes(),
-                            )
-                            .unwrap();
-                    }
+                    let (command, args) = line.split_once(' ').unwrap_or((line.trim(), ""));
+                    let output = match dbg!(command) {
+                        "pid" => self.get_process(args.parse().unwrap_or_default()),
+                        "list" => self.list_processes(),
+                        "limit" => self.set_power_limit(args),
+                        _ => "Unrecognized command#".into(),
+                    };
+                    socket.write_all(output.as_bytes()).unwrap();
                     line.clear();
                 }
             }
         });
     }
-
-    fn handle_request(&self, pid: i32) -> String {
+    fn list_processes(&self) -> String {
         let mut output = String::new();
         use std::fmt::Write;
-        if pid == -1 {
-            for (pid, info) in self.process_info.read().unwrap().iter() {
-                writeln!(&mut output, "{pid},{},{}", info.energy, info.tree_energy).unwrap();
-            }
-            writeln!(&mut output, "#",).unwrap();
-            return output;
+        for (pid, info) in self.process_info.read().unwrap().iter() {
+            writeln!(&mut output, "{pid},{},{}", info.energy, info.tree_energy).unwrap();
+        }
+        writeln!(&mut output, "#",).unwrap();
+        output
+    }
+
+    fn set_power_limit(&self, args: &str) -> String {
+        if let Ok(power_limit) = args.parse() {
+            self.power_limit
+                .store(power_limit, std::sync::atomic::Ordering::Relaxed);
+            "#".into()
+        } else {
+            format!("failed to parse {args} as energy limit#")
         }
+    }
 
+    fn get_process(&self, pid: i32) -> String {
         if let Some(info) = self.process_info.read().unwrap().get(&pid) {
             format!(
                 "pid: {pid} process: {}J process tree: {}J\n",
@@ -77,11 +88,12 @@ impl LoggingSocketService {
 pub fn start_logging_socket_service(
     path: &str,
     process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
-) -> io::Result<()> {
+) -> io::Result<Arc<AtomicU32>> {
     if Path::new(path).exists() {
         fs::remove_file(path)?;
     }
-    let socket = LoggingSocketService::new(path, process_info);
+    let arc = Arc::new(AtomicU32::new(100));
+    let socket = LoggingSocketService::new(path, process_info, arc.clone());
     socket.run();
-    Ok(())
+    Ok(arc)
 }