Skip to main content

sandbox_rs/monitoring/
monitor.rs

1//! Process monitoring via /proc
2//!
3//! Provides real-time monitoring of process resources using /proc filesystem.
4//! Tracks memory usage, CPU time, thread count, and process state.
5
6use std::fs;
7use std::path::Path;
8use std::time::{Duration, Instant};
9
10use nix::sys::signal::{Signal, kill};
11use nix::unistd::Pid;
12
13use sandbox_core::{Result, SandboxError};
14
15/// Process state enumeration
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum ProcessState {
18    /// Process is running
19    Running,
20    /// Process is sleeping
21    Sleeping,
22    /// Process is zombie
23    Zombie,
24    /// Process state is unknown
25    Unknown,
26}
27
28impl ProcessState {
29    /// Parse state from /proc stat first character
30    pub fn from_char(c: char) -> Self {
31        match c {
32            'R' => ProcessState::Running,
33            'S' => ProcessState::Sleeping,
34            'Z' => ProcessState::Zombie,
35            _ => ProcessState::Unknown,
36        }
37    }
38}
39
40/// Process statistics snapshot
41#[derive(Debug, Clone)]
42pub struct ProcessStats {
43    /// Process ID
44    pub pid: i32,
45    /// Virtual memory size in bytes
46    pub vsize: u64,
47    /// Resident set size in bytes (physical memory)
48    pub rss: u64,
49    /// RSS in MB (for convenience)
50    pub memory_usage_mb: u64,
51    /// CPU time in milliseconds
52    pub cpu_time_ms: u64,
53    /// Number of threads
54    pub num_threads: u32,
55    /// Current process state
56    pub state: ProcessState,
57    /// Timestamp of this snapshot
58    pub timestamp: Instant,
59}
60
61impl ProcessStats {
62    /// Create stats from /proc data
63    fn from_proc(pid: i32, timestamp: Instant) -> Result<Self> {
64        let stat_path = format!("/proc/{}/stat", pid);
65        let status_path = format!("/proc/{}/status", pid);
66
67        let stat_content = fs::read_to_string(&stat_path).map_err(|e| {
68            SandboxError::ProcessMonitoring(format!("Failed to read {}: {}", stat_path, e))
69        })?;
70
71        let parts: Vec<&str> = stat_content.split_whitespace().collect();
72        if parts.len() < 24 {
73            return Err(SandboxError::ProcessMonitoring(
74                "Invalid /proc/stat format".to_string(),
75            ));
76        }
77
78        let state = ProcessState::from_char(parts[2].chars().next().unwrap_or('?'));
79        let utime: u64 = parts[13]
80            .parse()
81            .map_err(|_| SandboxError::ProcessMonitoring("Invalid utime".to_string()))?;
82        let stime: u64 = parts[14]
83            .parse()
84            .map_err(|_| SandboxError::ProcessMonitoring("Invalid stime".to_string()))?;
85        let num_threads: u32 = parts[19]
86            .parse()
87            .map_err(|_| SandboxError::ProcessMonitoring("Invalid num_threads".to_string()))?;
88        let vsize: u64 = parts[22]
89            .parse()
90            .map_err(|_| SandboxError::ProcessMonitoring("Invalid vsize".to_string()))?;
91        let rss: u64 = parts[23]
92            .parse()
93            .map_err(|_| SandboxError::ProcessMonitoring("Invalid rss".to_string()))?;
94
95        let _status_content = fs::read_to_string(&status_path).unwrap_or_default();
96
97        let clk_tck = unsafe { libc::sysconf(libc::_SC_CLK_TCK) } as u64;
98        let cpu_time_ms = if clk_tck > 0 {
99            ((utime + stime) * 1000) / clk_tck
100        } else {
101            0
102        };
103
104        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as u64;
105        let rss_bytes = rss * page_size;
106        let memory_usage_mb = rss_bytes / (1024 * 1024);
107
108        Ok(ProcessStats {
109            pid,
110            vsize,
111            rss: rss_bytes,
112            memory_usage_mb,
113            cpu_time_ms,
114            num_threads,
115            state,
116            timestamp,
117        })
118    }
119}
120
121/// Process monitor for tracking sandbox resource usage
122pub struct ProcessMonitor {
123    pid: Pid,
124    creation_time: Instant,
125    peak_memory_mb: u64,
126    last_stats: Option<ProcessStats>,
127}
128
129impl ProcessMonitor {
130    /// Create new monitor for process
131    pub fn new(pid: Pid) -> Result<Self> {
132        let stat_path = format!("/proc/{}/stat", pid.as_raw());
133        if !Path::new(&stat_path).exists() {
134            return Err(SandboxError::ProcessMonitoring(format!(
135                "Process {} not found",
136                pid
137            )));
138        }
139
140        Ok(ProcessMonitor {
141            pid,
142            creation_time: Instant::now(),
143            peak_memory_mb: 0,
144            last_stats: None,
145        })
146    }
147
148    /// Collect current statistics
149    pub fn collect_stats(&mut self) -> Result<ProcessStats> {
150        let now = Instant::now();
151        let stats = ProcessStats::from_proc(self.pid.as_raw(), now)?;
152
153        if stats.memory_usage_mb > self.peak_memory_mb {
154            self.peak_memory_mb = stats.memory_usage_mb;
155        }
156
157        self.last_stats = Some(stats.clone());
158        Ok(stats)
159    }
160
161    /// Get peak memory usage since monitor creation (in MB)
162    pub fn peak_memory_mb(&self) -> u64 {
163        self.peak_memory_mb
164    }
165
166    /// Get elapsed time since monitor creation
167    pub fn elapsed(&self) -> Duration {
168        self.creation_time.elapsed()
169    }
170
171    /// Check if process is still alive
172    pub fn is_alive(&self) -> Result<bool> {
173        let stat_path = format!("/proc/{}/stat", self.pid.as_raw());
174        Ok(Path::new(&stat_path).exists())
175    }
176
177    /// Send SIGTERM (graceful shutdown)
178    pub fn send_sigterm(&self) -> Result<()> {
179        kill(self.pid, Signal::SIGTERM)
180            .map_err(|e| SandboxError::Syscall(format!("Failed to send SIGTERM: {}", e)))
181    }
182
183    /// Send SIGKILL (force termination)
184    pub fn send_sigkill(&self) -> Result<()> {
185        kill(self.pid, Signal::SIGKILL)
186            .map_err(|e| SandboxError::Syscall(format!("Failed to send SIGKILL: {}", e)))
187    }
188
189    /// Graceful shutdown: SIGTERM → wait → SIGKILL
190    pub fn graceful_shutdown(&self, wait_duration: Duration) -> Result<()> {
191        self.send_sigterm()?;
192
193        let start = Instant::now();
194        while start.elapsed() < wait_duration && self.is_alive()? {
195            std::thread::sleep(Duration::from_millis(10));
196        }
197
198        if self.is_alive()? {
199            self.send_sigkill()?;
200        }
201
202        Ok(())
203    }
204
205    /// Get last collected stats
206    pub fn last_stats(&self) -> Option<&ProcessStats> {
207        self.last_stats.as_ref()
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn test_process_state_from_char() {
217        assert_eq!(ProcessState::from_char('R'), ProcessState::Running);
218        assert_eq!(ProcessState::from_char('S'), ProcessState::Sleeping);
219        assert_eq!(ProcessState::from_char('Z'), ProcessState::Zombie);
220        assert_eq!(ProcessState::from_char('X'), ProcessState::Unknown);
221    }
222
223    #[test]
224    fn test_process_stats_creation() {
225        let pid = std::process::id() as i32;
226        let timestamp = Instant::now();
227        let result = ProcessStats::from_proc(pid, timestamp);
228        assert!(result.is_ok());
229
230        if let Ok(stats) = result {
231            assert_eq!(stats.pid, pid);
232            assert!(stats.memory_usage_mb > 0);
233        }
234    }
235
236    #[test]
237    fn test_process_monitor_new() {
238        let pid = Pid::from_raw(std::process::id() as i32);
239        let result = ProcessMonitor::new(pid);
240        assert!(result.is_ok());
241    }
242
243    #[test]
244    fn test_process_monitor_is_alive() {
245        let pid = Pid::from_raw(std::process::id() as i32);
246        let monitor = ProcessMonitor::new(pid).unwrap();
247        assert!(monitor.is_alive().unwrap());
248    }
249
250    #[test]
251    fn test_process_monitor_collect_stats() {
252        let pid = Pid::from_raw(std::process::id() as i32);
253        let mut monitor = ProcessMonitor::new(pid).unwrap();
254        let stats = monitor.collect_stats().unwrap();
255
256        assert_eq!(stats.pid, pid.as_raw());
257        assert!(stats.memory_usage_mb > 0);
258        assert_eq!(monitor.peak_memory_mb(), stats.memory_usage_mb);
259    }
260
261    #[test]
262    fn test_process_stats_from_proc_missing_file() {
263        let invalid_pid = 9_999_999i32;
264        let timestamp = Instant::now();
265        let result = ProcessStats::from_proc(invalid_pid, timestamp);
266        assert!(result.is_err());
267    }
268}