sandbox_rs/monitoring/
monitor.rs1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum ProcessState {
18 Running,
20 Sleeping,
22 Zombie,
24 Unknown,
26}
27
28impl ProcessState {
29 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#[derive(Debug, Clone)]
42pub struct ProcessStats {
43 pub pid: i32,
45 pub vsize: u64,
47 pub rss: u64,
49 pub memory_usage_mb: u64,
51 pub cpu_time_ms: u64,
53 pub num_threads: u32,
55 pub state: ProcessState,
57 pub timestamp: Instant,
59}
60
61impl ProcessStats {
62 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
121pub struct ProcessMonitor {
123 pid: Pid,
124 creation_time: Instant,
125 peak_memory_mb: u64,
126 last_stats: Option<ProcessStats>,
127}
128
129impl ProcessMonitor {
130 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 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 pub fn peak_memory_mb(&self) -> u64 {
163 self.peak_memory_mb
164 }
165
166 pub fn elapsed(&self) -> Duration {
168 self.creation_time.elapsed()
169 }
170
171 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 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 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 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 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}