Skip to main content

pfx_run/
pin.rs

1use std::path::Path;
2use std::process::Command;
3
4const MARK: &str = "PFX_SLICE";
5
6const MACHINE_ENV: Option<&str> = option_env!("PFX_MACHINE_ENV");
7
8pub fn build_cpus() -> Option<String> {
9    let path = MACHINE_ENV?;
10    let cpus = std::fs::read_to_string(path)
11        .ok()
12        .and_then(|text| setting(&text, "BUILD_CPUS"));
13    if cpus.is_none() {
14        eprintln!(
15            "{}: {path} names no BUILD_CPUS; running unpinned",
16            crate::job::TOOL
17        );
18    }
19    cpus
20}
21
22pub fn setting(text: &str, key: &str) -> Option<String> {
23    text.lines()
24        .map(str::trim)
25        .filter(|line| !line.starts_with('#'))
26        .filter_map(|line| line.split_once('='))
27        .filter(|(name, _)| name.trim().trim_start_matches("export ").trim() == key)
28        .map(|(_, value)| {
29            value
30                .trim()
31                .trim_matches('"')
32                .trim_matches('\'')
33                .to_string()
34        })
35        .rfind(|value| !value.is_empty())
36}
37
38pub fn taskset_args(cpus: &str) -> [String; 3] {
39    ["taskset".into(), "-c".into(), cpus.to_string()]
40}
41
42pub fn should_reenter(direct: bool, inside: bool, systemd: bool, taskset: bool) -> bool {
43    !direct && !inside && systemd && taskset
44}
45
46pub fn reentry(exe: &Path, args: &[String], cpus: Option<&str>) -> Command {
47    let mut command = Command::new("systemd-run");
48    command
49        .args([
50            "--user",
51            "--quiet",
52            "--scope",
53            "--collect",
54            "--slice=build.slice",
55            "--",
56        ])
57        .args(cpus.map(taskset_args).into_iter().flatten())
58        .args(["nice", "-n", "10"])
59        .arg(exe)
60        .args(args)
61        .env(MARK, "1");
62    command
63}
64
65pub fn reenter_if_needed(args: &[String]) -> Option<i32> {
66    MACHINE_ENV?;
67    let direct = std::env::var("PFX_DIRECT").ok().as_deref() == Some("1");
68    let inside = std::env::var_os(MARK).is_some();
69    if !should_reenter(
70        direct,
71        inside,
72        crate::queue::on_path("systemd-run"),
73        crate::queue::on_path("taskset"),
74    ) {
75        return None;
76    }
77    let exe = std::env::current_exe().ok()?;
78    let status = reentry(&exe, args, build_cpus().as_deref()).status().ok()?;
79    Some(status.code().unwrap_or(1))
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn the_build_cpus_come_from_the_machine_file_and_the_last_value_wins() {
88        let text = "# tuned for this machine\nQUEUE_SLOTS=2\nBUILD_CPUS=2-3,6-7\nexport BUILD_CPUS=\"4-5\"\n";
89        assert_eq!(setting(text, "BUILD_CPUS").as_deref(), Some("4-5"));
90        assert_eq!(setting(text, "QUEUE_SLOTS").as_deref(), Some("2"));
91        assert_eq!(setting(text, "MISSING"), None);
92        assert_eq!(setting("BUILD_CPUS=\n", "BUILD_CPUS"), None);
93    }
94
95    #[test]
96    fn without_a_cpu_set_the_reentry_runs_unpinned() {
97        let command = reentry(Path::new("/opt/pfx"), &["list".into()], None);
98        let args: Vec<String> = command
99            .get_args()
100            .map(|a| a.to_string_lossy().into_owned())
101            .collect();
102        assert!(!args.iter().any(|a| a == "taskset"));
103        assert!(args.windows(3).any(|w| w == ["nice", "-n", "10"]));
104    }
105
106    #[test]
107    fn pfx_reenters_the_build_slice_once_and_never_in_tests() {
108        assert!(should_reenter(false, false, true, true));
109        assert!(!should_reenter(true, false, true, true));
110        assert!(!should_reenter(false, true, true, true));
111        assert!(!should_reenter(false, false, false, true));
112        assert!(!should_reenter(false, false, true, false));
113    }
114
115    #[test]
116    fn the_reentry_runs_pinned_and_niced_in_the_build_slice() {
117        let command = reentry(
118            Path::new("/opt/pfx"),
119            &["render".into(), "--app".into(), "alpha".into()],
120            Some("2-3,6-7"),
121        );
122        assert_eq!(command.get_program(), "systemd-run");
123        let args: Vec<String> = command
124            .get_args()
125            .map(|a| a.to_string_lossy().into_owned())
126            .collect();
127        assert_eq!(
128            args,
129            [
130                "--user",
131                "--quiet",
132                "--scope",
133                "--collect",
134                "--slice=build.slice",
135                "--",
136                "taskset",
137                "-c",
138                "2-3,6-7",
139                "nice",
140                "-n",
141                "10",
142                "/opt/pfx",
143                "render",
144                "--app",
145                "alpha"
146            ]
147        );
148        assert!(
149            command
150                .get_envs()
151                .any(|(k, v)| k == MARK && v.is_some_and(|v| v == "1"))
152        );
153    }
154}