optative_process_pool/
process.rs

1use std::collections::BTreeMap;
2use std::io::{BufRead, Write as IoWrite};
3use std::path::PathBuf;
4use std::process::Stdio;
5use std::sync::mpsc;
6use std::thread;
7
8use optative::Lifecycle;
9
10use super::{StreamItem, StreamKind};
11
12/// Stable identity for a process: uniquely identifies which process to manage.
13/// Used as the key in `Lifecycle` so that `OptativeSet` can track processes by identity.
14#[derive(Hash, Eq, PartialEq, Clone, Debug, serde::Serialize, serde::Deserialize)]
15pub struct ProcessIdentity {
16    pub bin: String,
17    pub key: String,
18}
19
20// NOTE: env uses BTreeMap (not HashMap) for deterministic ordering; HashMap doesn't implement Hash.
21#[derive(Clone, Debug)]
22pub struct ProcessSource {
23    pub identity: ProcessIdentity,
24    pub args: Vec<String>,
25    pub env: BTreeMap<String, String>,
26    pub current_dir: Option<PathBuf>,
27    pub props: Option<serde_json::Value>,
28}
29
30pub struct ProcessState {
31    pub child: std::process::Child,
32    pub event_tx: mpsc::Sender<serde_json::Value>,
33    pub last_sent_props: Option<serde_json::Value>,
34}
35
36/// Error type for process spawning failures.
37#[derive(Debug, thiserror::Error)]
38pub enum SpawnError {
39    #[error("failed to spawn {bin}: {source}")]
40    ProcessSpawnFailed {
41        bin: String,
42        #[source]
43        source: std::io::Error,
44    },
45}
46
47fn spawn_stdout_thread(
48    stdout: std::process::ChildStdout,
49    identity: ProcessIdentity,
50    tx: mpsc::Sender<StreamItem>,
51) {
52    thread::spawn(move || {
53        let reader = std::io::BufReader::new(stdout);
54        for line in reader.lines() {
55            match line {
56                Ok(l) => {
57                    let item = StreamItem {
58                        key: identity.clone(),
59                        stream: StreamKind::Stdout,
60                        line: l,
61                    };
62                    if tx.send(item).is_err() {
63                        break;
64                    }
65                }
66                Err(_) => break,
67            }
68        }
69    });
70}
71
72fn spawn_stderr_thread(stderr: std::process::ChildStderr, bin_name: String) {
73    thread::spawn(move || {
74        let reader = std::io::BufReader::new(stderr);
75        for line in reader.lines() {
76            match line {
77                Ok(l) => tracing::warn!(module = %bin_name, "{l}"),
78                Err(_) => break,
79            }
80        }
81    });
82}
83
84fn spawn_stdin_thread(
85    mut stdin: std::process::ChildStdin,
86    event_rx: mpsc::Receiver<serde_json::Value>,
87) {
88    thread::spawn(move || {
89        while let Ok(event) = event_rx.recv() {
90            let line = serde_json::to_string(&event).unwrap_or_default() + "\n";
91            if stdin.write_all(line.as_bytes()).is_err() {
92                break;
93            }
94        }
95    });
96}
97
98fn expand_tilde(path: &str) -> String {
99    if path.starts_with("~/") {
100        let home = std::env::var("HOME").unwrap_or_default();
101        format!("{}{}", home, &path[1..])
102    } else if path == "~" {
103        std::env::var("HOME").unwrap_or_default()
104    } else {
105        path.to_string()
106    }
107}
108
109pub(super) fn spawn_process(
110    spec: ProcessSource,
111    tx: &mpsc::Sender<StreamItem>,
112) -> Result<ProcessState, SpawnError> {
113    let bin = expand_tilde(&spec.identity.bin);
114    let mut cmd = std::process::Command::new(&bin);
115    cmd.args(&spec.args);
116    for (k, v) in &spec.env {
117        cmd.env(k, v);
118    }
119    if let Some(ref dir) = spec.current_dir {
120        cmd.current_dir(dir);
121    }
122
123    cmd.stdout(Stdio::piped());
124    cmd.stderr(Stdio::piped());
125    cmd.stdin(Stdio::piped());
126
127    let mut child = match cmd.spawn() {
128        Ok(c) => c,
129        Err(e) => {
130            return Err(SpawnError::ProcessSpawnFailed { bin, source: e });
131        }
132    };
133
134    if let Some(stdout) = child.stdout.take() {
135        spawn_stdout_thread(stdout, spec.identity.clone(), tx.clone());
136    }
137    if let Some(stderr) = child.stderr.take() {
138        spawn_stderr_thread(stderr, spec.identity.bin.clone());
139    }
140    let (event_tx, event_rx) = mpsc::channel::<serde_json::Value>();
141    if let Some(stdin) = child.stdin.take() {
142        spawn_stdin_thread(stdin, event_rx);
143    }
144
145    Ok(ProcessState {
146        child,
147        event_tx,
148        last_sent_props: None,
149    })
150}
151
152impl std::fmt::Display for ProcessSource {
153    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        write!(f, "{}", self.identity.bin)
155    }
156}
157
158impl Lifecycle for ProcessSource {
159    type Key = ProcessIdentity;
160    type State = ProcessState;
161    type Context = ();
162    type Output = mpsc::Sender<StreamItem>;
163    type Error = SpawnError;
164
165    fn key(&self) -> ProcessIdentity {
166        self.identity.clone()
167    }
168
169    fn enter(self, _ctx: &mut (), output: &mut Self::Output) -> Result<Self::State, Self::Error> {
170        let props = self.props.clone();
171        let mut state = spawn_process(self, output)?;
172        if let Some(p) = props {
173            let _ = state.event_tx.send(p.clone());
174            state.last_sent_props = Some(p);
175        }
176        Ok(state)
177    }
178
179    #[allow(clippy::collapsible_if)]
180    fn reconcile_self(
181        self,
182        state: &mut Self::State,
183        _ctx: &mut (),
184        output: &mut Self::Output,
185    ) -> Result<(), Self::Error> {
186        if matches!(state.child.try_wait(), Ok(Some(_))) {
187            tracing::warn!(bin = %self.identity.bin, "process exited");
188            let props = self.props.clone();
189            let mut new_state = spawn_process(self, output)?;
190            if let Some(p) = props {
191                let _ = new_state.event_tx.send(p.clone());
192                new_state.last_sent_props = Some(p);
193            }
194            *state = new_state;
195        } else if let Some(p) = self.props {
196            if state.last_sent_props.as_ref() != Some(&p) {
197                let _ = state.event_tx.send(p.clone());
198                state.last_sent_props = Some(p);
199            }
200        }
201        Ok(())
202    }
203
204    fn exit(
205        mut state: Self::State,
206        _ctx: &mut (),
207        _output: &mut Self::Output,
208    ) -> Result<(), Self::Error> {
209        let _ = state.child.kill();
210        let _ = state.child.wait();
211        Ok(())
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::{ProcessIdentity, ProcessSource};
218    use optative::Lifecycle;
219    use std::collections::BTreeMap;
220
221    fn make_source(bin: &str) -> ProcessSource {
222        ProcessSource {
223            identity: ProcessIdentity {
224                bin: bin.to_string(),
225                key: bin.to_string(),
226            },
227            args: vec![],
228            env: BTreeMap::new(),
229            current_dir: None,
230            props: None,
231        }
232    }
233
234    #[test]
235    fn process_identity_has_bin_and_key_fields() {
236        let id = ProcessIdentity {
237            bin: "mybin".to_string(),
238            key: "mykey".to_string(),
239        };
240        assert_eq!(id.bin, "mybin");
241        assert_eq!(id.key, "mykey");
242    }
243
244    #[test]
245    fn process_identity_derives_hash_eq_partialeq_clone() {
246        use std::collections::HashSet;
247        let a = ProcessIdentity {
248            bin: "bin".to_string(),
249            key: "k".to_string(),
250        };
251        let b = a.clone();
252        assert_eq!(a, b);
253        let mut set = HashSet::new();
254        set.insert(a);
255        assert!(!set.insert(b));
256    }
257
258    #[test]
259    fn process_source_has_identity_fields() {
260        let spec = ProcessSource {
261            identity: ProcessIdentity {
262                bin: "/bin/sh".to_string(),
263                key: "my-key".to_string(),
264            },
265            args: vec!["--flag".to_string()],
266            env: BTreeMap::new(),
267            current_dir: None,
268            props: None,
269        };
270        assert_eq!(spec.identity.bin, "/bin/sh");
271        assert_eq!(spec.identity.key, "my-key");
272    }
273
274    #[test]
275    fn lifecycle_key_returns_identity() {
276        let id = ProcessIdentity {
277            bin: "/usr/bin/cat".to_string(),
278            key: "cat-key".to_string(),
279        };
280        let returned: ProcessIdentity = make_source("/usr/bin/cat").key();
281        assert_eq!(returned.bin, id.bin);
282    }
283
284    mod spawn_process {
285        use super::super::{SpawnError, spawn_process};
286        use std::sync::mpsc;
287
288        #[test]
289        fn nonexistent_binary_returns_process_spawn_failed() {
290            let (tx, _rx) = mpsc::channel();
291            let result = spawn_process(
292                super::make_source("/nonexistent/binary/that/cannot/exist"),
293                &tx,
294            );
295            match result {
296                Err(SpawnError::ProcessSpawnFailed { bin, .. }) => {
297                    assert_eq!(bin, "/nonexistent/binary/that/cannot/exist");
298                }
299                Ok(_) => panic!("expected Err, got Ok"),
300            }
301        }
302
303        #[test]
304        fn tilde_bin_is_expanded_to_home_dir() {
305            let home = std::env::var("HOME").expect("HOME must be set");
306            let (tx, _rx) = mpsc::channel();
307            let result = spawn_process(super::make_source("~/nonexistent-tilde-test-binary"), &tx);
308            match result {
309                Err(SpawnError::ProcessSpawnFailed { bin, .. }) => {
310                    assert!(
311                        !bin.starts_with('~'),
312                        "bin must not contain literal ~; got: {bin}"
313                    );
314                    assert!(
315                        bin.starts_with(&home),
316                        "bin must start with HOME ({home}); got: {bin}"
317                    );
318                }
319                Ok(_) => panic!("expected Err, got Ok"),
320            }
321        }
322    }
323
324    mod lifecycle {
325        use super::super::{ProcessIdentity, ProcessSource, SpawnError};
326        use optative::Lifecycle;
327        use std::collections::BTreeMap;
328        use std::sync::mpsc;
329
330        #[test]
331        fn reconcile_self_propagates_err_when_restart_spawn_fails() {
332            let (mut tx, _rx) = mpsc::channel();
333
334            let mut state = ProcessSource {
335                identity: ProcessIdentity {
336                    bin: "/bin/sh".to_string(),
337                    key: "t".to_string(),
338                },
339                args: vec!["-c".to_string(), "exit 0".to_string()],
340                env: BTreeMap::new(),
341                current_dir: None,
342                props: None,
343            }
344            .enter(&mut (), &mut tx)
345            .expect("enter must succeed with /bin/sh");
346
347            std::thread::sleep(std::time::Duration::from_millis(200));
348            assert!(
349                matches!(state.child.try_wait(), Ok(Some(_))),
350                "child should have exited"
351            );
352
353            let result = super::make_source("/nonexistent/binary/that/cannot/exist")
354                .reconcile_self(&mut state, &mut (), &mut tx);
355            match result {
356                Err(SpawnError::ProcessSpawnFailed { .. }) => {}
357                Ok(_) => panic!("expected Err, got Ok"),
358            }
359        }
360    }
361}
optative_process_pool/process.rs

optative_process_pool/
process.rs
