soma_codemode/pool/
runner_handle.rs

1use std::path::PathBuf;
2use std::process::Stdio;
3use std::sync::Arc;
4
5use futures::StreamExt;
6use tokio::process::{Child, ChildStdin, Command};
7use tokio::sync::{Mutex, Notify};
8use tokio_util::codec::{FramedRead, LinesCodec};
9
10use crate::runner::limits::MAX_STDIO_LINE_BYTES;
11use crate::ToolError;
12
13use super::job_guard::JobGuard;
14
15pub type RunnerLines = FramedRead<tokio::process::ChildStdout, LinesCodec>;
16
17#[derive(Debug, Clone)]
18pub struct RunnerSpawn {
19    pub program: PathBuf,
20    pub args: Vec<String>,
21}
22
23impl RunnerSpawn {
24    pub fn current_exe() -> Result<Self, ToolError> {
25        Ok(Self {
26            program: crate::runner_exe::resolve_runner_exe()?,
27            args: Vec::new(),
28        })
29    }
30}
31
32#[derive(Clone, Debug)]
33pub struct StderrBuffer {
34    lines: Arc<Mutex<Vec<String>>>,
35    notify: Arc<Notify>,
36}
37
38impl StderrBuffer {
39    fn new() -> Self {
40        Self {
41            lines: Arc::new(Mutex::new(Vec::new())),
42            notify: Arc::new(Notify::new()),
43        }
44    }
45
46    pub async fn mark(&self) -> usize {
47        self.lines.lock().await.len()
48    }
49
50    pub async fn take_since_and_clear(&self, start_index: usize) -> Vec<String> {
51        let mut guard = self.lines.lock().await;
52        let captured = guard
53            .get(start_index..)
54            .map(<[String]>::to_vec)
55            .unwrap_or_default();
56        guard.clear();
57        captured
58    }
59
60    pub async fn clear(&self) {
61        self.lines.lock().await.clear();
62    }
63
64    pub async fn flush_settle(&self) {
65        let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(50);
66        let mut last_len = self.lines.lock().await.len();
67        while let Ok(()) = tokio::time::timeout_at(deadline, self.notify.notified()).await {
68            let len = self.lines.lock().await.len();
69            if len == last_len {
70                break;
71            }
72            last_len = len;
73        }
74    }
75}
76
77#[derive(Debug)]
78pub struct RunnerHandle {
79    pub child: Child,
80    pub child_pid: Option<u32>,
81    pub stdin: ChildStdin,
82    pub lines: RunnerLines,
83    pub stderr: StderrBuffer,
84    pub success_count: u64,
85    _job_guard: JobGuard,
86    drain_task: tokio::task::JoinHandle<()>,
87    _temp_dir: tempfile::TempDir,
88}
89
90impl RunnerHandle {
91    pub fn spawn(spawn: &RunnerSpawn) -> Result<Self, ToolError> {
92        let temp_dir = tempfile::TempDir::new().map_err(|err| {
93            ToolError::internal_message(format!("failed to create runner temp dir: {err}"))
94        })?;
95        let mut command = Command::new(&spawn.program);
96        command
97            .args(&spawn.args)
98            .current_dir(temp_dir.path())
99            .env_clear()
100            .kill_on_drop(true)
101            .stdin(Stdio::piped())
102            .stdout(Stdio::piped())
103            .stderr(Stdio::piped());
104        #[cfg(unix)]
105        command.process_group(0);
106
107        let mut child = command.spawn().map_err(|err| {
108            ToolError::internal_message(format!(
109                "failed to spawn Code Mode runner from `{}`: {err}",
110                spawn.program.display()
111            ))
112        })?;
113        let child_pid = child.id();
114        let stdin = child
115            .stdin
116            .take()
117            .ok_or_else(|| ToolError::internal_message("runner stdin unavailable"))?;
118        let stdout = child
119            .stdout
120            .take()
121            .ok_or_else(|| ToolError::internal_message("runner stdout unavailable"))?;
122        let stderr_pipe = child
123            .stderr
124            .take()
125            .ok_or_else(|| ToolError::internal_message("runner stderr unavailable"))?;
126        let stderr = StderrBuffer::new();
127        let drain_task = spawn_stderr_drain(stderr_pipe, stderr.clone());
128        Ok(Self {
129            child,
130            child_pid,
131            stdin,
132            lines: FramedRead::new(
133                stdout,
134                LinesCodec::new_with_max_length(MAX_STDIO_LINE_BYTES),
135            ),
136            stderr,
137            success_count: 0,
138            _job_guard: JobGuard::new(child_pid),
139            drain_task,
140            _temp_dir: temp_dir,
141        })
142    }
143
144    #[cfg(test)]
145    pub fn spawn_stub_command(program: &str, args: &[&str]) -> Result<Self, ToolError> {
146        let spawn = RunnerSpawn {
147            program: PathBuf::from(program),
148            args: args.iter().map(|value| value.to_string()).collect(),
149        };
150        Self::spawn(&spawn)
151    }
152}
153
154impl Drop for RunnerHandle {
155    fn drop(&mut self) {
156        self.drain_task.abort();
157        #[cfg(unix)]
158        if let Some(pid) = self.child_pid {
159            use nix::sys::signal::Signal;
160            use nix::unistd::Pid;
161            let _ = nix::sys::signal::killpg(Pid::from_raw(pid as i32), Signal::SIGKILL);
162        }
163    }
164}
165
166fn spawn_stderr_drain(
167    stderr: tokio::process::ChildStderr,
168    buffer: StderrBuffer,
169) -> tokio::task::JoinHandle<()> {
170    tokio::spawn(async move {
171        const CAP_ENTRIES: usize = 100_000;
172        const CAP_BYTES: usize = 8 * 1024 * 1024;
173        let mut lines = FramedRead::new(
174            stderr,
175            LinesCodec::new_with_max_length(MAX_STDIO_LINE_BYTES),
176        );
177        let mut total_bytes = 0usize;
178        while let Some(line) = lines.next().await {
179            let line = match line {
180                Ok(line) => line,
181                Err(_) => "[soma] runner stderr truncated".to_string(),
182            };
183            total_bytes = total_bytes.saturating_add(line.len() + 1);
184            let mut buf = buffer.lines.lock().await;
185            if buf.len() < CAP_ENTRIES && total_bytes <= CAP_BYTES {
186                buf.push(line);
187            } else if buf
188                .last()
189                .is_none_or(|last| last != "[soma] runner stderr truncated")
190            {
191                buf.push("[soma] runner stderr truncated".to_string());
192            }
193            drop(buf);
194            buffer.notify.notify_waiters();
195        }
196    })
197}
soma_codemode/pool/runner_handle.rs

soma_codemode/pool/
runner_handle.rs
