Skip to main content

eska_tool_runtime/
capture.rs

1//! Bounded raw process capture shares the same ownership and cancellation rules as private IPC.
2
3use crate::{
4    Component, Error, Operation, QueryPolicy,
5    process::OwnedChild,
6    query::{Stream, reader},
7};
8#[cfg(unix)]
9const TERMINATE: i32 = signal_hook::consts::signal::SIGTERM;
10#[cfg(not(unix))]
11const TERMINATE: i32 = 0;
12
13use std::{
14    io::{self, Write},
15    process::{Command, Output, Stdio},
16    sync::atomic::{AtomicBool, Ordering},
17    sync::mpsc,
18    thread,
19    time::Instant,
20};
21
22/// Capture one owned subprocess, invoking native cleanup before the final kill/reap on failure.
23///
24/// # Errors
25/// Returns process, stream-limit, timeout or cancellation failures; the exact child is reaped.
26pub fn capture(
27    component: Component,
28    command: &mut Command,
29    scope: &Operation,
30    policy: QueryPolicy,
31    on_failure: impl FnOnce(),
32) -> Result<Output, Error> {
33    run(component, command, Some(scope), policy, on_failure)
34}
35
36/// Queries on Windows preserve their existing console controller rather than installing another.
37pub fn run(
38    component: Component,
39    command: &mut Command,
40    scope: Option<&Operation>,
41    policy: QueryPolicy,
42    on_failure: impl FnOnce(),
43) -> Result<Output, Error> {
44    if scope.is_some_and(Operation::is_cancelled) {
45        return Err(Error::Cancelled);
46    }
47    let child = command
48        .stdin(Stdio::null())
49        .stdout(Stdio::piped())
50        .stderr(Stdio::piped())
51        .spawn()
52        .map_err(|source| Error::Process { component, source })?;
53    let mut child = OwnedChild(child);
54    let result = receive(component, &mut child, scope, None, policy);
55    if result.is_err() && !matches!(child.0.try_wait(), Ok(Some(_))) {
56        on_failure();
57    }
58    result
59}
60
61/// Send one owned request while draining both child streams under timeout and RPC cancellation.
62///
63/// # Errors
64/// Returns process/pipe failures, cancellation and bounded-output failures; the exact child is reaped.
65pub fn capture_input(
66    component: Component,
67    command: &mut Command,
68    input: Vec<u8>,
69    policy: QueryPolicy,
70    cancelled: &AtomicBool,
71    closed: &AtomicBool,
72) -> Result<Output, Error> {
73    if cancelled.load(Ordering::Acquire) || closed.load(Ordering::Acquire) {
74        return Err(Error::Cancelled);
75    }
76    let child = command
77        .stdin(Stdio::piped())
78        .stdout(Stdio::piped())
79        .stderr(Stdio::piped())
80        .spawn()
81        .map_err(|source| Error::Process { component, source })?;
82    let mut child = OwnedChild(child);
83    let mut stdin = child.0.stdin.take().ok_or(Error::InvalidResponse)?;
84    let writer = thread::Builder::new()
85        .name("eska-component-input".into())
86        .spawn(move || stdin.write_all(&input))
87        .map_err(|source| Error::Process { component, source })?;
88    let result = receive(
89        component,
90        &mut child,
91        None,
92        Some((cancelled, closed)),
93        policy,
94    );
95    drop(child);
96    let written = writer
97        .join()
98        .map_err(|_| Error::Process {
99            component,
100            source: io::Error::other("input writer failed"),
101        })
102        .and_then(|result| result.map_err(|source| Error::Process { component, source }));
103    let output = result?;
104    if output.status.success() {
105        written?;
106    }
107    Ok(output)
108}
109
110/// Drain both pipes concurrently; a graceful owner can clean descendants before its own exit.
111fn receive(
112    component: Component,
113    child: &mut OwnedChild,
114    scope: Option<&Operation>,
115    request: Option<(&AtomicBool, &AtomicBool)>,
116    policy: QueryPolicy,
117) -> Result<Output, Error> {
118    let stdout = child.0.stdout.take().ok_or(Error::InvalidResponse)?;
119    let stderr = child.0.stderr.take().ok_or(Error::InvalidResponse)?;
120    let (sender, receiver) = mpsc::channel();
121    reader(
122        component,
123        stdout,
124        Stream::Stdout,
125        sender.clone(),
126        policy.output_limit,
127    )?;
128    reader(
129        component,
130        stderr,
131        Stream::Stderr,
132        sender,
133        policy.output_limit,
134    )?;
135    let started = Instant::now();
136    let mut stopping = None;
137    let mut status = None;
138    let mut stdout = None;
139    let mut stderr = None;
140    loop {
141        while let Ok((stream, result)) = receiver.try_recv() {
142            match stream {
143                Stream::Stdout => stdout = Some(result?),
144                Stream::Stderr => stderr = Some(result?),
145            }
146        }
147        let cancelled = scope.is_some_and(Operation::is_cancelled)
148            || request.is_some_and(|(cancelled, closed)| {
149                cancelled.load(Ordering::Acquire) || closed.load(Ordering::Acquire)
150            });
151        if stopping.is_none() && (cancelled || started.elapsed() >= policy.timeout) {
152            if policy.cancellation_grace.is_zero() {
153                return Err(reason(cancelled));
154            }
155            signal(
156                child,
157                scope
158                    .and_then(Operation::received_signal)
159                    .unwrap_or(TERMINATE),
160            )?;
161            stopping = Some((Instant::now(), cancelled));
162        }
163        if status.is_none() {
164            status = child
165                .0
166                .try_wait()
167                .map_err(|source| Error::Process { component, source })?;
168        }
169        if let Some((started, cancelled)) = stopping {
170            if status.is_some() || started.elapsed() >= policy.cancellation_grace {
171                return Err(reason(cancelled));
172            }
173        } else if let Some(status) = status
174            && stdout.is_some()
175            && stderr.is_some()
176        {
177            return Ok(Output {
178                status,
179                stdout: stdout.take().ok_or(Error::InvalidResponse)?,
180                stderr: stderr.take().ok_or(Error::InvalidResponse)?,
181            });
182        }
183        thread::sleep(std::time::Duration::from_millis(10));
184    }
185}
186
187/// Preserve the initiating failure even when the child exits successfully during cleanup.
188const fn reason(cancelled: bool) -> Error {
189    if cancelled {
190        Error::Cancelled
191    } else {
192        Error::Timeout
193    }
194}
195
196/// Signal only the owned process; its native owner handles any container-side PID.
197#[cfg(unix)]
198fn signal(child: &OwnedChild, signal: i32) -> Result<(), Error> {
199    use nix::{
200        errno::Errno,
201        sys::signal::{Signal, kill},
202        unistd::Pid,
203    };
204    let pid =
205        i32::try_from(child.0.id()).map_err(|error| Error::Signal(std::io::Error::other(error)))?;
206    let signal = Signal::try_from(signal).map_err(|error| Error::Signal(error.into()))?;
207    match kill(Pid::from_raw(pid), signal) {
208        Ok(()) | Err(Errno::ESRCH) => Ok(()),
209        Err(error) => Err(Error::Signal(error.into())),
210    }
211}
212
213/// Windows timeout fallback retains the previous direct-child termination behavior.
214#[cfg(not(unix))]
215fn signal(child: &mut OwnedChild, _: i32) -> Result<(), Error> {
216    child.0.kill().map_err(Error::Signal)
217}
218
219#[cfg(all(test, unix))]
220mod tests {
221    use super::*;
222    use crate::test_support::TestDir;
223    use std::{sync::Arc, time::Duration};
224
225    /// Cancellation also releases a writer blocked by a component that never consumes its input.
226    #[test]
227    fn cancelling_a_blocked_input_reaps_the_child_and_writer() {
228        let directory = TestDir::new();
229        let cancelled = Arc::new(AtomicBool::new(false));
230        let signal = Arc::clone(&cancelled);
231        let trigger = thread::spawn(move || {
232            thread::sleep(Duration::from_millis(50));
233            signal.store(true, Ordering::Release);
234        });
235        let mut command = Command::new("python3");
236        command
237            .args(["-c", "import time; time.sleep(30)"])
238            .current_dir(&directory.0);
239        let start = Instant::now();
240        let result = capture_input(
241            Component::Fmt,
242            &mut command,
243            vec![b'x'; 1024 * 1024],
244            QueryPolicy {
245                timeout: Duration::from_secs(3),
246                output_limit: 1024,
247                cancellation_grace: Duration::ZERO,
248            },
249            &cancelled,
250            &AtomicBool::new(false),
251        );
252        trigger.join().unwrap();
253        assert!(matches!(result, Err(Error::Cancelled)));
254        assert!(start.elapsed() < Duration::from_secs(2));
255    }
256
257    /// Both pipes drain while the complete input is sent; no partial response or broken-pipe success is accepted.
258    #[test]
259    fn input_and_output_are_complete_with_bounded_streams() {
260        let directory = TestDir::new();
261        let input = vec![b'a'; 256 * 1024];
262        let mut command = Command::new("python3");
263        command.args(["-c","import sys; data=sys.stdin.buffer.read(); sys.stderr.buffer.write(b'checked'); sys.stdout.buffer.write(data)"]).current_dir(&directory.0);
264        let output = capture_input(
265            Component::Fmt,
266            &mut command,
267            input.clone(),
268            QueryPolicy {
269                timeout: Duration::from_secs(3),
270                output_limit: 1024 * 1024,
271                cancellation_grace: Duration::ZERO,
272            },
273            &AtomicBool::new(false),
274            &AtomicBool::new(false),
275        )
276        .unwrap();
277        assert!(output.status.success());
278        assert_eq!(output.stdout, input);
279        assert_eq!(output.stderr, b"checked");
280    }
281}