use crate::error::{NucleusError, Result};
use std::fs;
use std::path::Path;
#[derive(Debug, Clone)]
pub struct ResourceStats {
pub memory_usage: u64,
pub memory_limit: u64,
pub memory_swap_usage: u64,
pub cpu_usage_ns: u64,
pub pid_count: u64,
pub memory_percent: f64,
}
impl ResourceStats {
pub fn from_cgroup(cgroup_path: &str) -> Result<Self> {
let cgroup_path = Path::new(cgroup_path);
let memory_usage = Self::read_memory_current(cgroup_path)?;
let memory_limit = Self::read_memory_max(cgroup_path)?;
let memory_percent = if memory_limit > 0 {
(memory_usage as f64 / memory_limit as f64) * 100.0
} else {
0.0
};
let memory_swap_usage = Self::read_memory_swap(cgroup_path).unwrap_or(0);
let cpu_usage_ns = Self::read_cpu_usage(cgroup_path)?;
let pid_count = Self::read_pid_current(cgroup_path)?;
Ok(Self {
memory_usage,
memory_limit,
memory_swap_usage,
cpu_usage_ns,
pid_count,
memory_percent,
})
}
fn read_memory_current(cgroup_path: &Path) -> Result<u64> {
let path = cgroup_path.join("memory.current");
Self::read_u64_file(&path)
}
fn read_memory_max(cgroup_path: &Path) -> Result<u64> {
let path = cgroup_path.join("memory.max");
let content = fs::read_to_string(&path).map_err(|e| {
NucleusError::ResourceError(format!("Failed to read {:?}: {}", path, e))
})?;
if content.trim() == "max" {
Ok(0)
} else {
content.trim().parse().map_err(|e| {
NucleusError::ResourceError(format!("Failed to parse memory.max: {}", e))
})
}
}
fn read_memory_swap(cgroup_path: &Path) -> Result<u64> {
let path = cgroup_path.join("memory.swap.current");
Self::read_u64_file(&path)
}
fn read_cpu_usage(cgroup_path: &Path) -> Result<u64> {
let path = cgroup_path.join("cpu.stat");
let content = fs::read_to_string(&path).map_err(|e| {
NucleusError::ResourceError(format!("Failed to read {:?}: {}", path, e))
})?;
for line in content.lines() {
if let Some(value_str) = line.strip_prefix("usage_usec ") {
let usec: u64 = value_str.parse().map_err(|e| {
NucleusError::ResourceError(format!("Failed to parse CPU usage: {}", e))
})?;
return Ok(usec * 1000);
}
}
Err(NucleusError::ResourceError(format!(
"cpu.stat at {:?} does not contain 'usage_usec' key",
path
)))
}
fn read_pid_current(cgroup_path: &Path) -> Result<u64> {
let path = cgroup_path.join("pids.current");
Self::read_u64_file(&path)
}
fn read_u64_file(path: &Path) -> Result<u64> {
let content = fs::read_to_string(path).map_err(|e| {
NucleusError::ResourceError(format!("Failed to read {:?}: {}", path, e))
})?;
content
.trim()
.parse()
.map_err(|e| NucleusError::ResourceError(format!("Failed to parse {:?}: {}", path, e)))
}
pub fn format_memory(bytes: u64) -> String {
const KB: u64 = 1024;
const MB: u64 = KB * 1024;
const GB: u64 = MB * 1024;
if bytes >= GB {
format!("{:.2}G", bytes as f64 / GB as f64)
} else if bytes >= MB {
format!("{:.2}M", bytes as f64 / MB as f64)
} else if bytes >= KB {
format!("{:.2}K", bytes as f64 / KB as f64)
} else {
format!("{}B", bytes)
}
}
pub fn format_cpu_time(ns: u64) -> String {
let seconds = ns as f64 / 1_000_000_000.0;
format!("{:.2}s", seconds)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_format_memory() {
assert_eq!(ResourceStats::format_memory(512), "512B");
assert_eq!(ResourceStats::format_memory(1024), "1.00K");
assert_eq!(ResourceStats::format_memory(1024 * 1024), "1.00M");
assert_eq!(ResourceStats::format_memory(1024 * 1024 * 1024), "1.00G");
assert_eq!(ResourceStats::format_memory(512 * 1024 * 1024), "512.00M");
}
#[test]
fn test_format_cpu_time() {
assert_eq!(ResourceStats::format_cpu_time(0), "0.00s");
assert_eq!(ResourceStats::format_cpu_time(1_000_000_000), "1.00s");
assert_eq!(ResourceStats::format_cpu_time(5_500_000_000), "5.50s");
}
}