1use std::io::Read;
2use std::path::PathBuf;
3use std::process::{Child, Command, Stdio};
4use std::time::{Duration, Instant};
5
6use serde::{Deserialize, Serialize};
7use tempfile::TempDir;
8
9#[derive(Debug, Clone)]
10pub struct SupervisorConfig {
11 pub max_capture_lines: usize,
12 pub max_line_bytes: usize,
13 pub startup_timeout: Duration,
14 pub scenario_timeout: Duration,
15 pub shutdown_timeout: Duration,
16 pub artifact_dir: Option<PathBuf>,
17}
18
19impl Default for SupervisorConfig {
20 fn default() -> Self {
21 Self {
22 max_capture_lines: 1000,
23 max_line_bytes: 4096,
24 startup_timeout: Duration::from_secs(5),
25 scenario_timeout: Duration::from_secs(15),
26 shutdown_timeout: Duration::from_secs(3),
27 artifact_dir: None,
28 }
29 }
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct ProcessExit {
34 pub exit_code: Option<i32>,
35 pub signal: Option<String>,
36 pub lifetime_ms: u64,
37}
38
39pub struct SupervisedProcess {
40 child: Option<Child>,
41 pid: u32,
42 stdout_lines: Vec<String>,
43 stderr_lines: Vec<String>,
44 start_time: Instant,
45 config: SupervisorConfig,
46 artifact_dir: TempDir,
47 killed: bool,
48}
49
50impl SupervisedProcess {
51 pub fn spawn(
52 config: SupervisorConfig,
53 program: &str,
54 args: &[&str],
55 ) -> Result<Self, std::io::Error> {
56 let artifact_dir = match &config.artifact_dir {
57 Some(dir) => {
58 std::fs::create_dir_all(dir)?;
59 TempDir::new_in(dir)?
60 }
61 None => TempDir::new()?,
62 };
63
64 let mut cmd = Command::new(program);
65 cmd.args(args);
66
67 #[cfg(unix)]
68 {
69 use std::os::unix::process::CommandExt;
70 cmd.process_group(0);
71 }
72
73 cmd.stdout(Stdio::piped());
74 cmd.stderr(Stdio::piped());
75
76 let child = cmd.spawn()?;
77 let pid = child.id();
78
79 Ok(Self {
80 child: Some(child),
81 pid,
82 stdout_lines: Vec::new(),
83 stderr_lines: Vec::new(),
84 start_time: Instant::now(),
85 config,
86 artifact_dir,
87 killed: false,
88 })
89 }
90
91 pub fn capture_output(&mut self) {
92 if let Some(ref mut child) = self.child {
93 let mut stdout_buf = Vec::new();
94 let mut stderr_buf = Vec::new();
95 if let Some(ref mut r) = child.stdout {
96 drain_reader(
97 r,
98 &mut stdout_buf,
99 self.config.max_capture_lines,
100 self.config.max_line_bytes,
101 );
102 }
103 if let Some(ref mut r) = child.stderr {
104 drain_reader(
105 r,
106 &mut stderr_buf,
107 self.config.max_capture_lines,
108 self.config.max_line_bytes,
109 );
110 }
111 self.stdout_lines.extend(stdout_buf);
112 self.stderr_lines.extend(stderr_buf);
113 }
114 }
115
116 pub fn wait(&mut self) -> ProcessExit {
117 self.capture_output();
118 let exit = self.child.as_mut().map(|c| c.wait());
119 self.capture_output();
120
121 let lifetime_ms = self.start_time.elapsed().as_millis() as u64;
122 match exit {
123 Some(Ok(status)) => ProcessExit {
124 exit_code: status.code(),
125 signal: None,
126 lifetime_ms,
127 },
128 _ => ProcessExit {
129 exit_code: None,
130 signal: Some("wait_failed".to_string()),
131 lifetime_ms,
132 },
133 }
134 }
135
136 pub fn kill_group(&mut self) {
137 if self.killed {
138 return;
139 }
140 self.killed = true;
141
142 if let Some(ref mut child) = self.child {
143 let _ = child.kill();
144 }
145
146 if let Some(ref mut child) = self.child {
147 let _ = child.wait();
148 }
149 }
150
151 pub fn save_artifacts(&self, label: &str) -> std::io::Result<()> {
152 let stdout_path = self.artifact_dir.path().join(format!("{label}.stdout.log"));
153 let stderr_path = self.artifact_dir.path().join(format!("{label}.stderr.log"));
154
155 std::fs::write(&stdout_path, self.stdout_lines.join("\n"))?;
156 std::fs::write(&stderr_path, self.stderr_lines.join("\n"))?;
157
158 Ok(())
159 }
160
161 pub fn artifact_dir(&self) -> &std::path::Path {
162 self.artifact_dir.path()
163 }
164
165 pub fn stdout_lines(&self) -> &[String] {
166 &self.stdout_lines
167 }
168
169 pub fn stderr_lines(&self) -> &[String] {
170 &self.stderr_lines
171 }
172
173 pub fn pid(&self) -> u32 {
174 self.pid
175 }
176}
177
178fn drain_reader(
179 reader: &mut impl Read,
180 lines: &mut Vec<String>,
181 max_lines: usize,
182 max_bytes: usize,
183) {
184 let mut buf = [0u8; 8192];
185 loop {
186 match reader.read(&mut buf) {
187 Ok(0) => break,
188 Ok(n) => {
189 let chunk = String::from_utf8_lossy(&buf[..n]);
190 for line in chunk.lines() {
191 if lines.len() < max_lines {
192 let truncated = if line.len() > max_bytes {
193 format!("{}... (truncated)", &line[..max_bytes])
194 } else {
195 line.to_string()
196 };
197 lines.push(truncated);
198 }
199 }
200 }
201 Err(_) => break,
202 }
203 }
204}
205
206impl Drop for SupervisedProcess {
207 fn drop(&mut self) {
208 self.kill_group();
209 self.capture_output();
210 let _ = self.save_artifacts(&format!("pid_{}", self.pid));
211 }
212}
213
214#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
215#[serde(rename_all = "snake_case")]
216pub enum ReadinessProbe {
217 TcpPort,
218 StdoutPattern(String),
219 FixedDelay(u64),
220 FileExists(String),
221}
222
223pub async fn probe_readiness(
224 addr: std::net::SocketAddr,
225 probe: &ReadinessProbe,
226 timeout: Duration,
227) -> bool {
228 match probe {
229 ReadinessProbe::TcpPort => probe_tcp_port(addr, timeout).await,
230 ReadinessProbe::StdoutPattern(_) => {
231 tokio::time::sleep(Duration::from_millis(200)).await;
232 true
233 }
234 ReadinessProbe::FixedDelay(ms) => {
235 tokio::time::sleep(Duration::from_millis(*ms)).await;
236 true
237 }
238 ReadinessProbe::FileExists(path) => {
239 let start = Instant::now();
240 loop {
241 if std::path::Path::new(path).exists() {
242 return true;
243 }
244 if start.elapsed() >= timeout {
245 return false;
246 }
247 tokio::time::sleep(Duration::from_millis(50)).await;
248 }
249 }
250 }
251}
252
253async fn probe_tcp_port(addr: std::net::SocketAddr, timeout: Duration) -> bool {
254 let start = Instant::now();
255 loop {
256 match tokio::net::TcpStream::connect(addr).await {
257 Ok(_) => return true,
258 Err(_) => {
259 if start.elapsed() >= timeout {
260 return false;
261 }
262 tokio::time::sleep(Duration::from_millis(50)).await;
263 }
264 }
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 #[test]
273 fn supervised_process_spawn_and_wait() {
274 let config = SupervisorConfig::default();
275 let mut proc = SupervisedProcess::spawn(config, "echo", &["hello"]).expect("spawn failed");
276 let exit = proc.wait();
277 assert_eq!(exit.exit_code, Some(0));
278 proc.capture_output();
279 assert!(proc.stdout_lines.iter().any(|l| l.contains("hello")));
280 }
281
282 #[test]
283 fn supervised_process_kill_group() {
284 let config = SupervisorConfig::default();
285 let mut proc = SupervisedProcess::spawn(config, "sleep", &["60"]).expect("spawn failed");
286 proc.kill_group();
287 let exit = proc.wait();
288 assert!(exit.exit_code.is_none() || exit.exit_code != Some(0));
289 }
290
291 #[cfg(unix)]
292 #[test]
293 fn bounded_capture_limits_output() {
294 let config = SupervisorConfig {
295 max_capture_lines: 3,
296 ..Default::default()
297 };
298 let mut proc = SupervisedProcess::spawn(
299 config,
300 "bash",
301 &["-c", "for i in $(seq 1 10); do echo \"line $i\"; done"],
302 )
303 .expect("spawn failed");
304 let exit = proc.wait();
305 assert_eq!(exit.exit_code, Some(0));
306 proc.capture_output();
307 assert!(proc.stdout_lines.len() <= 3);
308 }
309
310 #[test]
311 fn artifact_save() {
312 let config = SupervisorConfig::default();
313 let mut proc = SupervisedProcess::spawn(config, "echo", &["test"]).expect("spawn failed");
314 let exit = proc.wait();
315 assert_eq!(exit.exit_code, Some(0));
316 proc.save_artifacts("test").expect("save failed");
317 let stdout_path = proc.artifact_dir().join("test.stdout.log");
318 assert!(stdout_path.exists());
319 }
320
321 #[tokio::test]
322 async fn probe_tcp_port_success() {
323 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
324 let addr = listener.local_addr().unwrap();
325 assert!(probe_tcp_port(addr, Duration::from_secs(1)).await);
326 }
327
328 #[tokio::test]
329 async fn probe_tcp_port_timeout() {
330 let addr: std::net::SocketAddr = "127.0.0.1:1".parse().unwrap();
331 assert!(!probe_tcp_port(addr, Duration::from_millis(100)).await);
332 }
333}