soma_fleet/
process_driver.rs

1use std::process::Stdio;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use tokio::io::AsyncWriteExt;
6use tokio::process::Command;
7use tokio_util::sync::CancellationToken;
8
9use crate::{
10    CommandExecutor, CommandOutput, CommandRequest, FleetError, FleetResult, HostEndpoint,
11    HostRecord, io::drain_bounded,
12};
13
14/// Local process-backed command driver used for conformance and local hosts.
15#[derive(Debug, Clone, Copy, Default)]
16pub struct LocalProcessDriver;
17
18#[async_trait]
19impl CommandExecutor for LocalProcessDriver {
20    async fn execute(
21        &self,
22        host: &HostRecord,
23        request: &CommandRequest,
24        cancellation: &CancellationToken,
25    ) -> FleetResult<CommandOutput> {
26        if !matches!(host.endpoint(), HostEndpoint::Local) {
27            return Err(FleetError::Command {
28                host: host.id().clone(),
29                message: "local process driver requires a local endpoint".into(),
30            });
31        }
32        if cancellation.is_cancelled() {
33            return Err(FleetError::Cancelled);
34        }
35        request.validate_at(soma_ops::Timestamp::now())?;
36        let timeout = remaining(request.deadline())?;
37
38        let mut command = Command::new(request.program());
39        command
40            .args(request.args())
41            .stdin(if request.stdin().is_some() {
42                Stdio::piped()
43            } else {
44                Stdio::null()
45            })
46            .stdout(Stdio::piped())
47            .stderr(Stdio::piped())
48            .kill_on_drop(true);
49        if let Some(directory) = request.working_dir() {
50            command.current_dir(directory);
51        }
52        let mut child = command.spawn().map_err(|error| FleetError::Command {
53            host: host.id().clone(),
54            message: format!("spawn failed: {error}"),
55        })?;
56        let input = match request.stdin() {
57            Some(bytes) => Some((
58                child.stdin.take().ok_or_else(|| FleetError::Command {
59                    host: host.id().clone(),
60                    message: "stdin pipe unavailable".into(),
61                })?,
62                bytes.to_vec(),
63            )),
64            None => None,
65        };
66        let stdout = child.stdout.take().ok_or_else(|| FleetError::Command {
67            host: host.id().clone(),
68            message: "stdout pipe unavailable".into(),
69        })?;
70        let stderr = child.stderr.take().ok_or_else(|| FleetError::Command {
71            host: host.id().clone(),
72            message: "stderr pipe unavailable".into(),
73        })?;
74
75        let input = async move {
76            if let Some((mut stdin, bytes)) = input {
77                stdin.write_all(&bytes).await?;
78                stdin.shutdown().await?;
79            }
80            Ok::<_, std::io::Error>(())
81        };
82        let completion = async {
83            let (status, (stdout, stderr), ()) = tokio::try_join!(
84                child.wait(),
85                async {
86                    tokio::try_join!(
87                        drain_bounded(stdout, request.max_stdout_bytes()),
88                        drain_bounded(stderr, request.max_stderr_bytes())
89                    )
90                },
91                input
92            )?;
93            Ok::<_, std::io::Error>((status, stdout, stderr))
94        };
95
96        tokio::select! {
97            () = cancellation.cancelled() => {
98                let _ = child.kill().await;
99                let _ = child.wait().await;
100                Err(FleetError::Cancelled)
101            }
102            result = tokio::time::timeout(timeout, completion) => {
103                match result {
104                    Err(_) => {
105                        let _ = child.kill().await;
106                        let _ = child.wait().await;
107                        Err(FleetError::DeadlineExceeded)
108                    }
109                    Ok(Err(error)) => Err(FleetError::Command {
110                        host: host.id().clone(),
111                        message: format!("process I/O failed: {error}"),
112                    }),
113                    Ok(Ok((status, (stdout, stdout_truncated), (stderr, stderr_truncated)))) => {
114                        Ok(CommandOutput::new(
115                            stdout,
116                            stderr,
117                            status.code(),
118                            stdout_truncated || stderr_truncated,
119                        ))
120                    }
121                }
122            }
123        }
124    }
125}
126
127fn remaining(deadline: soma_ops::Timestamp) -> FleetResult<Duration> {
128    let millis = deadline
129        .unix_millis()
130        .saturating_sub(soma_ops::Timestamp::now().unix_millis());
131    if millis <= 0 {
132        Err(FleetError::DeadlineExceeded)
133    } else {
134        Ok(Duration::from_millis(millis as u64))
135    }
136}
137
138#[cfg(test)]
139#[path = "process_driver_tests.rs"]
140mod tests;
soma_fleet/process_driver.rs

soma_fleet/
process_driver.rs
