pinnacle_api/
process.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! Process management.
6//!
7//! This module provides ways to spawn processes and handle their output.
8
9use std::{
10    collections::HashMap,
11    os::fd::{FromRawFd, OwnedFd},
12};
13
14use passfd::FdPassingExt;
15use pinnacle_api_defs::pinnacle::process::v1::{SetEnvRequest, SpawnRequest, WaitOnSpawnRequest};
16use tokio_stream::StreamExt;
17
18use crate::{client::Client, BlockOnTokio};
19
20/// Adds an environment variable that all newly spawned [`Command`]s will inherit.
21pub fn set_env(key: impl ToString, value: impl ToString) {
22    Client::process()
23        .set_env(SetEnvRequest {
24            key: key.to_string(),
25            value: value.to_string(),
26        })
27        .block_on_tokio()
28        .unwrap();
29}
30
31/// A process builder that allows you to spawn programs.
32pub struct Command {
33    cmd: Vec<String>,
34    envs: HashMap<String, String>,
35    shell_cmd: Vec<String>,
36    unique: bool,
37    once: bool,
38    pipe_stdin: bool,
39    pipe_stdout: bool,
40    pipe_stderr: bool,
41}
42
43/// The result of spawning a [`Command`].
44#[derive(Debug)]
45pub struct Child {
46    pid: u32,
47    /// This process's standard input.
48    ///
49    /// This will only be `Some` if [`Command::pipe_stdin`] was called before spawning.
50    pub stdin: Option<tokio::process::ChildStdin>,
51    /// This process's standard output.
52    ///
53    /// This will only be `Some` if [`Command::pipe_stdout`] was called before spawning.
54    pub stdout: Option<tokio::process::ChildStdout>,
55    /// This process's standard error.
56    ///
57    /// This will only be `Some` if [`Command::pipe_stderr`] was called before spawning.
58    pub stderr: Option<tokio::process::ChildStderr>,
59}
60
61/// Information from an exited process.
62#[derive(Debug, Default)]
63pub struct ExitInfo {
64    /// The process's exit code.
65    pub exit_code: Option<i32>,
66    /// The process's exit message.
67    pub exit_msg: Option<String>,
68}
69
70impl Child {
71    /// Waits for this process to exit, blocking the current thread.
72    pub fn wait(self) -> ExitInfo {
73        self.wait_async().block_on_tokio()
74    }
75
76    /// Async impl for [`Self::wait`].
77    pub async fn wait_async(self) -> ExitInfo {
78        let mut exit_status = Client::process()
79            .wait_on_spawn(WaitOnSpawnRequest { pid: self.pid })
80            .await
81            .unwrap()
82            .into_inner();
83
84        let thing = exit_status.next().await;
85
86        let Some(Ok(response)) = thing else {
87            return Default::default();
88        };
89
90        ExitInfo {
91            exit_code: response.exit_code,
92            exit_msg: response.exit_msg,
93        }
94    }
95}
96
97impl Drop for Child {
98    fn drop(&mut self) {
99        let pid = self.pid;
100
101        // Wait on the process so it doesn't go zombie
102        tokio::spawn(async move {
103            Client::process()
104                .wait_on_spawn(WaitOnSpawnRequest { pid })
105                .await
106                .unwrap();
107        });
108    }
109}
110
111impl Command {
112    /// Creates a new [`Command`] that will spawn the provided `program`.
113    pub fn new(program: impl ToString) -> Self {
114        Self {
115            cmd: vec![program.to_string()],
116            envs: Default::default(),
117            shell_cmd: Vec::new(),
118            unique: false,
119            once: false,
120            pipe_stdin: false,
121            pipe_stdout: false,
122            pipe_stderr: false,
123        }
124    }
125
126    /// Creates a new [`Command`] that will spawn the provided `command` using the given shell and
127    /// its arguments.
128    ///
129    /// # Examples
130    ///
131    /// ```no_run
132    /// # use pinnacle_api::process::Command;
133    /// Command::with_shell(["bash", "-c"], "cat file.txt &> /dev/null").spawn();
134    /// ```
135    pub fn with_shell(
136        shell_args: impl IntoIterator<Item = impl ToString>,
137        command: impl ToString,
138    ) -> Self {
139        Self {
140            cmd: vec![command.to_string()],
141            envs: Default::default(),
142            shell_cmd: shell_args
143                .into_iter()
144                .map(|args| args.to_string())
145                .collect(),
146            unique: false,
147            once: false,
148            pipe_stdin: false,
149            pipe_stdout: false,
150            pipe_stderr: false,
151        }
152    }
153
154    /// Adds an argument to the command.
155    pub fn arg(&mut self, arg: impl ToString) -> &mut Self {
156        self.cmd.push(arg.to_string());
157        self
158    }
159
160    /// Adds multiple arguments to the command.
161    pub fn args(&mut self, args: impl IntoIterator<Item = impl ToString>) -> &mut Self {
162        self.cmd.extend(args.into_iter().map(|arg| arg.to_string()));
163        self
164    }
165
166    /// Sets an environment variable that the process will spawn with.
167    pub fn env(&mut self, key: impl ToString, value: impl ToString) -> &mut Self {
168        self.envs.insert(key.to_string(), value.to_string());
169        self
170    }
171
172    /// Sets multiple environment variables that the process will spawn with.
173    pub fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
174    where
175        I: IntoIterator<Item = (K, V)>,
176        K: ToString,
177        V: ToString,
178    {
179        self.envs.extend(
180            vars.into_iter()
181                .map(|(k, v)| (k.to_string(), v.to_string())),
182        );
183        self
184    }
185
186    /// Causes this command to only spawn the program if it is the only instance currently running.
187    pub fn unique(&mut self) -> &mut Self {
188        self.unique = true;
189        self
190    }
191
192    /// Causes this command to spawn the program exactly once in the compositor's lifespan.
193    pub fn once(&mut self) -> &mut Self {
194        self.once = true;
195        self
196    }
197
198    /// Sets up a pipe to allow the config to write to the process's stdin.
199    ///
200    /// The pipe will be available through the spawned child's [`stdin`][Child::stdin].
201    pub fn pipe_stdin(&mut self) -> &mut Self {
202        self.pipe_stdin = true;
203        self
204    }
205
206    /// Sets up a pipe to allow the config to read from the process's stdout.
207    ///
208    /// The pipe will be available through the spawned child's [`stdout`][Child::stdout].
209    pub fn pipe_stdout(&mut self) -> &mut Self {
210        self.pipe_stdout = true;
211        self
212    }
213
214    /// Sets up a pipe to allow the config to read from the process's stderr.
215    ///
216    /// The pipe will be available through the spawned child's [`stderr`][Child::stderr].
217    pub fn pipe_stderr(&mut self) -> &mut Self {
218        self.pipe_stderr = true;
219        self
220    }
221
222    /// Spawns this command, returning the spawned process's standard io, if any.
223    pub fn spawn(&mut self) -> Option<Child> {
224        let data = Client::process()
225            .spawn(SpawnRequest {
226                cmd: self.cmd.clone(),
227                unique: self.unique,
228                once: self.once,
229                shell_cmd: self.shell_cmd.clone(),
230                envs: self.envs.clone(),
231                pipe_stdin: self.pipe_stdin,
232                pipe_stdout: self.pipe_stdout,
233                pipe_stderr: self.pipe_stderr,
234            })
235            .block_on_tokio()
236            .unwrap()
237            .into_inner()
238            .spawn_data?;
239
240        let pid = data.pid;
241        let fd_socket_path = data.fd_socket_path;
242
243        let mut stdin = None;
244        let mut stdout = None;
245        let mut stderr = None;
246
247        let stream = std::os::unix::net::UnixStream::connect(fd_socket_path)
248            .expect("this should be set up by the compositor");
249
250        if data.has_stdin {
251            let fd = stream.recv_fd().unwrap();
252            // SAFETY: Fds are dup'd in over the socket
253            let child_stdin =
254                tokio::process::ChildStdin::from_std(std::process::ChildStdin::from(unsafe {
255                    OwnedFd::from_raw_fd(fd)
256                }))
257                .unwrap();
258            stdin = Some(child_stdin);
259        }
260
261        if data.has_stdout {
262            let fd = stream.recv_fd().unwrap();
263            let child_stdout =
264                tokio::process::ChildStdout::from_std(std::process::ChildStdout::from(unsafe {
265                    OwnedFd::from_raw_fd(fd)
266                }))
267                .unwrap();
268            stdout = Some(child_stdout);
269        }
270
271        if data.has_stderr {
272            let fd = stream.recv_fd().unwrap();
273            let child_stderr =
274                tokio::process::ChildStderr::from_std(std::process::ChildStderr::from(unsafe {
275                    OwnedFd::from_raw_fd(fd)
276                }))
277                .unwrap();
278            stderr = Some(child_stderr);
279        }
280
281        Some(Child {
282            pid,
283            stdin,
284            stdout,
285            stderr,
286        })
287    }
288}
pinnacle_api/process.rs

pinnacle_api/
process.rs
