1 files changed, 35 insertions, 48 deletions

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