3 files changed, 40 insertions, 28 deletions

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