use chrono::{DateTime, Utc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use sysinfo::System;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceSnapshot {
pub timestamp: DateTime<Utc>,
pub cpu_percent: f32,
pub memory_used_mb: u64,
pub memory_total_mb: u64,
pub active_agents: usize,
pub pending_tasks: usize,
pub total_token_usage: u64,
pub disk_used_gb: f64,
pub load_avg_1m: f32,
}
#[derive(Debug, Clone, Copy)]
pub struct OverloadThreshold {
pub cpu_percent: f32,
pub memory_percent: f32,
pub load_avg: f32,
}
impl Default for OverloadThreshold {
fn default() -> Self {
Self {
cpu_percent: 90.0,
memory_percent: 90.0,
load_avg: 8.0,
}
}
}
pub struct ResourceMonitor {
interval_secs: u64,
history_max: usize,
history: RwLock<VecDeque<ResourceSnapshot>>,
total_token_usage: AtomicU64,
active_agents: AtomicUsize,
pending_tasks: AtomicUsize,
overload_threshold: OverloadThreshold,
sys: parking_lot::Mutex<System>,
}
impl Default for ResourceMonitor {
fn default() -> Self {
Self::new(60, 60)
}
}
impl ResourceMonitor {
pub fn new(interval_secs: u64, history_max: usize) -> Self {
Self {
interval_secs,
history_max,
history: RwLock::new(VecDeque::with_capacity(history_max)),
total_token_usage: AtomicU64::new(0),
active_agents: AtomicUsize::new(0),
pending_tasks: AtomicUsize::new(0),
overload_threshold: OverloadThreshold::default(),
sys: parking_lot::Mutex::new(System::new_all()),
}
}
pub fn snapshot(&self) -> ResourceSnapshot {
let mut sys = self.sys.lock();
sys.refresh_all();
let cpu_percent =
sys.cpus().iter().map(|c| c.cpu_usage()).sum::<f32>() / sys.cpus().len().max(1) as f32;
let total_memory = sys.total_memory();
let used_memory = sys.used_memory();
let memory_total_mb = total_memory / (1024 * 1024);
let memory_used_mb = used_memory / (1024 * 1024);
let load_avg_1m = System::load_average().one as f32;
let disk_used_gb = estimate_disk_usage();
ResourceSnapshot {
timestamp: Utc::now(),
cpu_percent,
memory_used_mb,
memory_total_mb,
active_agents: self.active_agents.load(Ordering::Relaxed),
pending_tasks: self.pending_tasks.load(Ordering::Relaxed),
total_token_usage: self.total_token_usage.load(Ordering::Relaxed),
disk_used_gb,
load_avg_1m,
}
}
pub fn record_snapshot(&self) {
let snap = self.snapshot();
let mut history = self.history.write();
if history.len() >= self.history_max {
history.pop_front();
}
history.push_back(snap);
}
pub fn start_sampling(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
let monitor = Arc::clone(self);
let interval = self.interval_secs;
tokio::spawn(async move {
let mut ticker = tokio::time::interval(std::time::Duration::from_secs(interval));
loop {
ticker.tick().await;
monitor.record_snapshot();
}
})
}
pub fn history(&self, last_n: usize) -> Vec<ResourceSnapshot> {
let guard = self.history.read();
let n = last_n.min(guard.len());
guard.iter().rev().take(n).cloned().collect()
}
pub fn is_overloaded(&self) -> bool {
let snap = self.snapshot();
let memory_percent = if snap.memory_total_mb > 0 {
(snap.memory_used_mb as f32 / snap.memory_total_mb as f32) * 100.0
} else {
0.0
};
snap.cpu_percent >= self.overload_threshold.cpu_percent
|| memory_percent >= self.overload_threshold.memory_percent
|| snap.load_avg_1m >= self.overload_threshold.load_avg
}
pub fn set_active_agents(&self, count: usize) {
self.active_agents.store(count, Ordering::Relaxed);
}
pub fn set_pending_tasks(&self, count: usize) {
self.pending_tasks.store(count, Ordering::Relaxed);
}
pub fn add_token_usage(&self, tokens: u64) {
self.total_token_usage.fetch_add(tokens, Ordering::Relaxed);
}
pub fn overload_threshold(&self) -> OverloadThreshold {
self.overload_threshold
}
}
fn estimate_disk_usage() -> f64 {
let cwd = std::env::current_dir().unwrap_or_default();
walk_dir_size(&cwd) as f64 / (1024.0 * 1024.0 * 1024.0)
}
fn walk_dir_size(path: &std::path::Path) -> u64 {
let mut total = 0u64;
if let Ok(entries) = std::fs::read_dir(path) {
for entry in entries.flatten() {
let meta = entry.metadata();
if let Ok(m) = meta {
if m.is_file() {
total += m.len();
} else if m.is_dir() {
total += walk_dir_size(&entry.path());
}
}
}
}
total
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_snapshot_structure() {
let monitor = ResourceMonitor::default();
let snap = monitor.snapshot();
assert!(snap.timestamp <= Utc::now());
assert!(snap.cpu_percent >= 0.0);
assert!(snap.memory_used_mb >= 0);
assert!(snap.memory_total_mb >= 0);
assert!(snap.active_agents >= 0);
assert!(snap.pending_tasks >= 0);
assert!(snap.total_token_usage >= 0);
assert!(snap.disk_used_gb >= 0.0);
assert!(snap.load_avg_1m >= 0.0);
}
#[test]
fn test_is_overloaded_default_threshold() {
let monitor = ResourceMonitor::default();
let _ = monitor.is_overloaded();
}
#[test]
fn test_is_overloaded_high_thresholds_not_overloaded() {
let monitor = ResourceMonitor::default();
let result = monitor.is_overloaded();
let _ = result;
}
#[test]
fn test_history_management() {
let monitor = ResourceMonitor::new(1, 5);
assert!(monitor.history(10).is_empty());
for _ in 0..3 {
monitor.record_snapshot();
}
let history = monitor.history(10);
assert_eq!(history.len(), 3);
}
#[test]
fn test_history_eviction() {
let monitor = ResourceMonitor::new(1, 3);
for _ in 0..5 {
monitor.record_snapshot();
}
let history = monitor.history(10);
assert_eq!(history.len(), 3);
}
#[test]
fn test_set_active_agents() {
let monitor = ResourceMonitor::default();
monitor.set_active_agents(5);
let snap = monitor.snapshot();
assert_eq!(snap.active_agents, 5);
}
#[test]
fn test_set_pending_tasks() {
let monitor = ResourceMonitor::default();
monitor.set_pending_tasks(3);
let snap = monitor.snapshot();
assert_eq!(snap.pending_tasks, 3);
}
#[test]
fn test_add_token_usage() {
let monitor = ResourceMonitor::default();
monitor.add_token_usage(100);
monitor.add_token_usage(200);
let snap = monitor.snapshot();
assert_eq!(snap.total_token_usage, 300);
}
#[test]
fn test_overload_threshold_default() {
let threshold = OverloadThreshold::default();
assert_eq!(threshold.cpu_percent, 90.0);
assert_eq!(threshold.memory_percent, 90.0);
assert_eq!(threshold.load_avg, 8.0);
}
#[test]
fn test_overload_threshold_custom() {
let threshold = OverloadThreshold {
cpu_percent: 75.0,
memory_percent: 80.0,
load_avg: 4.0,
};
assert_eq!(threshold.cpu_percent, 75.0);
assert_eq!(threshold.memory_percent, 80.0);
assert_eq!(threshold.load_avg, 4.0);
}
#[test]
fn test_history_last_n() {
let monitor = ResourceMonitor::new(1, 10);
let empty = monitor.history(5);
assert!(empty.is_empty());
let many = monitor.history(100);
assert!(many.is_empty());
}
#[test]
fn test_load_average_struct() {
let la = System::load_average();
assert!(la.one >= 0.0);
assert!(la.five >= 0.0);
assert!(la.fifteen >= 0.0);
}
}