1use std::ffi::OsStr;
2use std::process::Command;
3
4pub const HEADLESS_VAR: &str = "PFX_HEADLESS";
5
6pub const DISPLAY_VARS: [&str; 4] = ["WAYLAND_DISPLAY", "WAYLAND_SOCKET", "DISPLAY", "XAUTHORITY"];
7
8pub fn headless(command: &mut Command) -> &mut Command {
9 for name in DISPLAY_VARS {
10 command.env_remove(name);
11 }
12 command.env(HEADLESS_VAR, "1")
13}
14
15pub struct Launch {
16 pub command: Command,
17 pub queued: bool,
18}
19
20pub fn launch(class: &str, as_name: &str, program: impl AsRef<OsStr>) -> Launch {
21 let direct = direct_set(std::env::var("PFX_DIRECT").ok().as_deref());
22 let cpus = on_path("taskset").then(crate::pin::build_cpus).flatten();
23 launch_with(
24 class,
25 as_name,
26 program,
27 direct,
28 on_path("pgpu"),
29 on_path("systemd-run"),
30 cpus.as_deref(),
31 )
32}
33
34fn direct_set(value: Option<&str>) -> bool {
35 value == Some("1")
36}
37
38fn launch_with(
39 class: &str,
40 as_name: &str,
41 program: impl AsRef<OsStr>,
42 direct: bool,
43 pgpu: bool,
44 systemd: bool,
45 cpus: Option<&str>,
46) -> Launch {
47 let picked = choice(direct, pgpu, systemd);
48 if matches!(picked, Choice::Missing) {
49 eprintln!(
50 "{}: pgpu is not installed here; running the adapter directly",
51 crate::job::TOOL
52 );
53 }
54 let mut launch = prepare(picked, class, as_name, program, cpus);
55 headless(&mut launch.command);
56 launch
57}
58
59enum Choice {
60 Direct,
61 Missing,
62 Systemd,
63 Pgpu,
64}
65
66fn choice(direct: bool, pgpu: bool, systemd: bool) -> Choice {
67 if direct {
68 Choice::Direct
69 } else if !pgpu {
70 Choice::Missing
71 } else if systemd {
72 Choice::Systemd
73 } else {
74 Choice::Pgpu
75 }
76}
77
78fn prepare(
79 picked: Choice,
80 class: &str,
81 as_name: &str,
82 program: impl AsRef<OsStr>,
83 cpus: Option<&str>,
84) -> Launch {
85 match picked {
86 Choice::Direct | Choice::Missing => Launch {
87 command: Command::new(program),
88 queued: false,
89 },
90 Choice::Systemd => {
91 let mut command = Command::new("systemd-run");
92 command.args([
93 "--user",
94 "--quiet",
95 "--scope",
96 "--collect",
97 "--slice=build.slice",
98 "-p",
99 "MemoryMax=12G",
100 "-p",
101 "MemorySwapMax=0",
102 "--",
103 ]);
104 if let Some(cpus) = cpus {
105 command.args(crate::pin::taskset_args(cpus));
106 }
107 command.args([
108 "nice", "-n", "10", "pgpu", "run", "--class", class, "--as", as_name, "--",
109 ]);
110 command.arg(program);
111 Launch {
112 command,
113 queued: true,
114 }
115 }
116 Choice::Pgpu => {
117 let mut command = Command::new("pgpu");
118 command
119 .args(["run", "--class", class, "--as", as_name, "--"])
120 .arg(program);
121 Launch {
122 command,
123 queued: true,
124 }
125 }
126 }
127}
128
129pub fn on_path(name: &str) -> bool {
130 std::env::var_os("PATH").is_some_and(|paths| {
131 std::env::split_paths(&paths)
132 .any(|dir| dir.join(name).is_file() || dir.join(format!("{name}.exe")).is_file())
133 })
134}
135
136#[cfg(test)]
137mod tests {
138 use super::*;
139
140 #[test]
141 fn pfx_direct_runs_the_child_and_the_wrap_names_the_tool() {
142 assert!(direct_set(Some("1")));
143 assert!(!direct_set(Some("off")));
144 assert!(!direct_set(Some("0")));
145 assert!(!direct_set(None));
146
147 let direct = launch_with("clip", "pfx", "adapter", true, true, true, None);
148 assert!(!direct.queued);
149 let envs: Vec<_> = direct.command.get_envs().collect();
150 for name in DISPLAY_VARS {
151 assert!(envs.contains(&(OsStr::new(name), None)), "{name} kept");
152 }
153 assert!(envs.contains(&(OsStr::new(HEADLESS_VAR), Some(OsStr::new("1")))));
154 assert_eq!(direct.command.get_program(), "adapter");
155 assert_eq!(direct.command.get_args().count(), 0);
156 assert!(
157 launch_with("clip", "pfx", "adapter", false, true, true, None)
158 .command
159 .get_envs()
160 .any(|pair| pair == (OsStr::new(HEADLESS_VAR), Some(OsStr::new("1"))))
161 );
162
163 let missing = launch_with("clip", "pfx", "render", false, false, true, None);
164 assert!(!missing.queued);
165 assert_eq!(missing.command.get_program(), "render");
166
167 let wrapped = launch_with("clip", "pfx", "adapter", false, true, true, Some("2-3,6-7"));
168 assert!(wrapped.queued);
169 assert_eq!(wrapped.command.get_program(), "systemd-run");
170 let args: Vec<_> = wrapped
171 .command
172 .get_args()
173 .map(|a| a.to_string_lossy().into_owned())
174 .collect();
175 assert_eq!(
176 args,
177 [
178 "--user",
179 "--quiet",
180 "--scope",
181 "--collect",
182 "--slice=build.slice",
183 "-p",
184 "MemoryMax=12G",
185 "-p",
186 "MemorySwapMax=0",
187 "--",
188 "taskset",
189 "-c",
190 "2-3,6-7",
191 "nice",
192 "-n",
193 "10",
194 "pgpu",
195 "run",
196 "--class",
197 "clip",
198 "--as",
199 "pfx",
200 "--",
201 "adapter",
202 ]
203 );
204
205 let bare = launch_with("truth", "pfx", "render", false, true, false, None);
206 assert!(bare.queued);
207 assert_eq!(bare.command.get_program(), "pgpu");
208 let args: Vec<_> = bare
209 .command
210 .get_args()
211 .map(|a| a.to_string_lossy().into_owned())
212 .collect();
213 assert_eq!(
214 args,
215 ["run", "--class", "truth", "--as", "pfx", "--", "render",]
216 );
217 }
218
219 #[test]
220 fn the_job_runs_as_the_name_the_caller_gives() {
221 for name in ["alpha", "beta", "delta", "render", "pfx"] {
222 let wrapped = launch_with("clip", name, "adapter", false, true, true, Some("2-3"));
223 let args: Vec<_> = wrapped
224 .command
225 .get_args()
226 .map(|a| a.to_string_lossy().into_owned())
227 .collect();
228 let at = args.iter().position(|a| a == "--as").unwrap();
229 assert_eq!(args[at + 1], name);
230 let bare = launch_with("truth", name, "render", false, true, false, None);
231 let args: Vec<_> = bare
232 .command
233 .get_args()
234 .map(|a| a.to_string_lossy().into_owned())
235 .collect();
236 assert_eq!(
237 args,
238 ["run", "--class", "truth", "--as", name, "--", "render"]
239 );
240 }
241 }
242}