Gather basic memory and CPU stats for services and processes

src/resources/mod.rs

@@ -6,3 +6,5 @@ pub mod network;
 pub mod system;
 /// This module contains functions that are responsible to stream information via websockets.
 pub mod stream;
+/// This module contains functions related to service and process monitoring.
+pub mod operations;

src/resources/operations.rs

@@ -0,0 +1,161 @@
+use serde::{Deserialize, Serialize};
+use crate::squire;
+use sysinfo::{Pid, System};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Usage {
+    name: String,
+    pid: u32,
+    cpu: String,
+    memory: String,
+}
+
+pub fn process_monitor(system: &System, process_names: &Vec<String>) -> Vec<Usage> {
+    let mut usages: Vec<Usage> = Vec::new();
+    for (pid, process) in system.processes() {
+        let process_name = process.name().to_str().unwrap().to_string();
+        if process_names.iter().any(|given_name| process_name.contains(given_name)) {
+            let cpu_usage = process.cpu_usage();
+            let memory_usage = process.memory();
+            let cpu = format!("{}", cpu_usage);
+            let memory = squire::util::size_converter(memory_usage);
+            let pid_32 = pid.as_u32();
+            usages.push(Usage {
+                name: process_name,
+                pid: pid_32,
+                cpu,
+                memory,
+            });
+        }
+    }
+    usages
+}
+
+pub fn service_monitor(system: &System, service_names: &Vec<String>) -> Vec<Usage> {
+    let mut usages: Vec<Usage> = Vec::new();
+    for service_name in service_names {
+        match service_monitor_fn(system, &service_name) {
+            Ok(usage) => usages.push(usage),
+            Err(err) => {
+                log::debug!("{}", err);
+                usages.push(Usage {
+                    name: service_name.to_string(),
+                    pid: 0000,
+                    memory: "N/A".to_string(),
+                    cpu: "N/A".to_string(),
+                });
+            }
+        };
+    }
+    usages
+}
+
+fn service_monitor_fn(system: &System, service_name: &String) -> Result<Usage, String> {
+    let pid = match get_service_pid(service_name) {
+        Some(pid) => pid,
+        None => return Err(format!("Failed to get PID for service: {}", service_name)),
+    };
+    let sys_pid: Pid = Pid::from(pid as usize);
+    if let Some(process) = system.process(sys_pid) {
+        let cpu_usage = process.cpu_usage();
+        let memory_usage = process.memory();
+        let cpu = format!("{}", cpu_usage);
+        let memory = squire::util::size_converter(memory_usage);
+        let pid_32 = sys_pid.as_u32();
+        Ok(Usage {
+            name: service_name.to_string(),
+            pid: pid_32,
+            cpu,
+            memory,
+        })
+    } else {
+        Err(format!("Process with PID {} not found", pid))
+    }
+}
+
+/// Function to get PID of a service (OS-agnostic)
+///
+/// # See Also
+///
+/// Service names are case-sensitive, so use the following command to get the right name.
+///
+/// * macOS: `launchctl list | grep {{ service_name }}`
+/// * Linux: `systemctl show {{ service_name }} --property=MainPID`
+/// * Windows: `sc query {{ service_name }}`
+fn get_service_pid(service_name: &str) -> Option<i32> {
+    let operating_system = std::env::consts::OS;
+    match operating_system {
+        "macos" => get_service_pid_macos(service_name, "/bin/launchctl"),
+        "linux" => get_service_pid_linux(service_name, "/usr/bin/systemctl"),
+        "windows" => get_service_pid_windows(service_name, "C:\\Windows\\System32\\sc.exe"),
+        _ => {
+            log::error!("Unsupported operating system: {}", operating_system);
+            None
+        }
+    }
+}
+
+// Linux: Use systemctl to get the service PID
+fn get_service_pid_linux(service_name: &str, lib_path: &str) -> Option<i32> {
+    let result = squire::util::run_command(
+        lib_path,
+        &["show", service_name, "--property=MainPID"],
+        true,
+    );
+    let output = match result {
+        Ok(output) => output,
+        Err(_) => return None,
+    };
+    if let Some(line) = output.lines().find(|line| line.starts_with("MainPID=")) {
+        if let Some(pid_str) = line.split('=').nth(1) {
+            if let Ok(pid) = pid_str.trim().parse::<i32>() {
+                return Some(pid);
+            }
+        }
+    }
+    None
+}
+
+// macOS: Use launchctl to get the service PID
+fn get_service_pid_macos(service_name: &str, lib_path: &str) -> Option<i32> {
+    let result = squire::util::run_command(
+        lib_path,
+        &["list"],
+        true,
+    );
+    let output = match result {
+        Ok(output) => output,
+        Err(_) => return None,
+    };
+    for line in output.lines() {
+        if line.contains(service_name) {
+            // Split the line and extract the PID (usually first column)
+            let parts: Vec<&str> = line.split_whitespace().collect();
+            if let Ok(pid) = parts[0].parse::<i32>() {
+                return Some(pid);
+            }
+        }
+    }
+    None
+}
+
+// Windows: Use sc query or PowerShell to get the service PID
+fn get_service_pid_windows(service_name: &str, lib_path: &str) -> Option<i32> {
+    let result = squire::util::run_command(
+        lib_path,
+        &["query", service_name],
+        true,
+    );
+    let output = match result {
+        Ok(output) => output,
+        Err(_) => return None,
+    };
+    if let Some(line) = output.lines().find(|line| line.contains("PID")) {
+        if let Some(pid_str) = line.split(':').nth(1) {
+            if let Ok(pid) = pid_str.trim().parse::<i32>() {
+                return Some(pid);
+            }
+        }
+    }
+    None
+}

src/resources/stream.rs

@@ -62,6 +62,22 @@ fn get_docker_stats() -> Result<Vec<serde_json::Value>, Box<dyn std::error::Erro
     Ok(stats)
 }
 
+fn get_service_stats(
+    system: &System,
+    config: &squire::settings::Config
+) -> Vec<serde_json::Value> {
+    let usages = resources::operations::service_monitor(&system, &config.services);
+    usages.into_iter().map(|usage| serde_json::to_value(usage).unwrap()).collect()
+}
+
+fn get_process_stats(
+    system: &System,
+    config: &squire::settings::Config
+) -> Vec<serde_json::Value> {
+    let usages = resources::operations::process_monitor(&system, &config.processes);
+    usages.into_iter().map(|usage| serde_json::to_value(usage).unwrap()).collect()
+}
+
 /// Function to get CPU usage percentage.
 ///
 /// # Returns
@@ -85,10 +101,7 @@ fn get_cpu_percent() -> Vec<String> {
 /// # Returns
 ///
 /// A `HashMap` containing the system metrics with CPU load average, memory and swap usage.
-fn get_system_metrics() -> HashMap<String, serde_json::Value> {
-    let mut system = System::new_all();
-    system.refresh_all();
-
+fn get_system_metrics(system: &System) -> HashMap<String, serde_json::Value> {
     // https://docs.rs/sysinfo/0.31.4/sysinfo/struct.System.html#method.load_average
     // Currently this doesn't work on Windows
     let load_avg = System::load_average();
@@ -129,11 +142,21 @@ fn get_system_metrics() -> HashMap<String, serde_json::Value> {
 /// # Returns
 ///
 /// A `HashMap` containing the system information with basic system information and memory/storage information.
-pub fn system_resources() -> HashMap<String, serde_json::Value> {
-    let mut system_metrics = get_system_metrics();
+pub fn system_resources(config: &squire::settings::Config) -> HashMap<String, serde_json::Value> {
+    let mut system = System::new_all();
+    system.refresh_all();
+    let mut system_metrics = get_system_metrics(&system);
     let cpu_percent = get_cpu_percent();
     let docker_stats = get_docker_stats().unwrap();
     system_metrics.insert("cpu_usage".to_string(), serde_json::json!(cpu_percent));
     system_metrics.insert("docker_stats".to_string(), serde_json::json!(docker_stats));
+    if !config.services.is_empty() {
+        let service_stats = get_service_stats(&system, &config);
+        system_metrics.insert("service_stats".to_string(), serde_json::json!(service_stats));
+    }
+    if !config.processes.is_empty() {
+        let process_stats = get_process_stats(&system, &config);
+        system_metrics.insert("process_stats".to_string(), serde_json::json!(process_stats));
+    }
     system_metrics
 }

src/routes/websocket.rs

@@ -13,11 +13,15 @@ use std::time::Duration;
 /// # Arguments
 ///
 /// * `request` - A reference to the Actix web `HttpRequest` object.
-async fn send_system_resources(request: HttpRequest, mut session: actix_ws::Session) {
+async fn send_system_resources(
+    request: HttpRequest,
+    mut session: actix_ws::Session,
+    config: web::Data<Arc<squire::settings::Config>>,
+) {
     let host = request.connection_info().host().to_string();
     let disk_stats = resources::stream::get_disk_stats();
     loop {
-        let mut system_resources = resources::stream::system_resources();
+        let mut system_resources = resources::stream::system_resources(&config);
         system_resources.insert("disk_info".to_string(), disk_stats.clone());
         let serialized = serde_json::to_string(&system_resources).unwrap();
         match session.text(serialized).await {
@@ -123,7 +127,7 @@ async fn echo(
         .aggregate_continuations();
     rt::spawn(async move {
         log::warn!("Connection established");
-        let send_task = send_system_resources(request.clone(), session.clone());
+        let send_task = send_system_resources(request.clone(), session.clone(), config.clone());
         let receive_task = receive_messages(session.clone(), stream);
         let session_task = session_handler(session.clone(), config.session_duration);
         future::join3(send_task, receive_task, session_task).await;

src/squire/settings.rs

@@ -25,6 +25,10 @@ pub struct Config {
     pub max_connections: usize,
     /// List of websites (supports regex) to add to CORS configuration.
     pub websites: Vec<String>,
+    /// List of services to monitor.
+    pub services: Vec<String>,
+    /// List of processes to monitor.
+    pub processes: Vec<String>,
 }
 
 /// Returns the default value for debug flag.
@@ -70,5 +74,5 @@ pub fn default_workers() -> usize {
 /// Returns the default maximum number of concurrent connections (3)
 pub fn default_max_connections() -> usize { 3 }
 
-/// Returns an empty list as the default website (CORS configuration)
-pub fn default_websites() -> Vec<String> { Vec::new() }
+/// Returns an empty vec
+pub fn default_vec() -> Vec<String> { Vec::new() }

src/squire/startup.rs

@@ -223,7 +223,9 @@ fn load_env_vars() -> settings::Config {
     let session_duration = parse_i64("session_duration").unwrap_or(settings::default_session_duration());
     let workers = parse_usize("workers").unwrap_or(settings::default_workers());
     let max_connections = parse_usize("max_connections").unwrap_or(settings::default_max_connections());
-    let websites = parse_vec("websites").unwrap_or(settings::default_websites());
+    let websites = parse_vec("websites").unwrap_or(settings::default_vec());
+    let services = parse_vec("services").unwrap_or(settings::default_vec());
+    let processes = parse_vec("processes").unwrap_or(settings::default_vec());
     settings::Config {
         username,
         password,
@@ -235,6 +237,8 @@ fn load_env_vars() -> settings::Config {
         workers,
         max_connections,
         websites,
+        services,
+        processes
     }
 }
 
@@ -356,6 +360,7 @@ pub fn get_config(metadata: &constant::MetaData) -> std::sync::Arc<settings::Con
     let env_file_path = std::env::current_dir()
         .unwrap_or_default()
         .join(env_file);
+    // https://github.com/dotenv-rs/dotenv/issues/94
     let _ = dotenv::from_path(env_file_path.as_path());
     std::sync::Arc::new(validate_vars())
 }