Skip to main content

xtask/
capture.rs

1//! File-backed child output keeps long generations out of memory while running.
2use crate::{metrics, util};
3use anyhow::{Result, ensure};
4use serde_json::{Value, json};
5use std::{
6    fs::{self, File},
7    io::{Read, Seek, SeekFrom},
8    path::Path,
9    process::{Child, Command},
10    sync::atomic::{AtomicBool, Ordering},
11    time::{Duration, Instant},
12};
13
14static INTERRUPTED: AtomicBool = AtomicBool::new(false);
15
16extern "C" fn interrupt(_: libc::c_int) {
17    INTERRUPTED.store(true, Ordering::Relaxed);
18}
19
20/// Let the polling loop unwind normally so its child is killed and reaped.
21pub fn install_interrupt_handler() -> Result<()> {
22    for signal in [libc::SIGINT, libc::SIGTERM] {
23        // The handler only writes a lock-free atomic; it allocates nothing.
24        let previous =
25            unsafe { libc::signal(signal, interrupt as *const () as libc::sighandler_t) };
26
27        ensure!(
28            previous != libc::SIG_ERR,
29            "could not install interrupt handler"
30        );
31    }
32
33    Ok(())
34}
35
36pub fn power() -> Value {
37    let detail = util::output(&["pmset", "-g", "batt"]).unwrap_or_else(|_| "unknown".into());
38    let source = if detail.contains("'AC Power'") {
39        "ac"
40    } else if detail.contains("'Battery Power'") {
41        "battery"
42    } else {
43        "unknown"
44    };
45
46    json!({"source":source,"detail":detail})
47}
48
49pub struct ChildGuard(pub Child);
50
51impl Drop for ChildGuard {
52    fn drop(&mut self) {
53        let _ = self.0.kill();
54        let _ = self.0.wait();
55    }
56}
57
58pub struct Captured {
59    pub code: i32,
60    pub stderr: String,
61    pub wall_seconds: f64,
62    pub power_before: Value,
63    pub power_after: Value,
64    pub power_samples: Vec<Value>,
65    pub cycle: Option<Value>,
66    pub timed_out: bool,
67}
68
69impl Captured {
70    pub fn stable_power(&self) -> bool {
71        self.power_samples
72            .iter()
73            .chain([&self.power_after])
74            .all(|p| p["source"] == self.power_before["source"])
75    }
76}
77
78pub fn run(
79    args: &[String],
80    output: &Path,
81    allow_battery: bool,
82    timeout: Option<Duration>,
83) -> Result<Captured> {
84    let before = power();
85
86    ensure!(
87        allow_battery || before["source"] == "ac",
88        "AC power is required; reconnect and resume, or use --allow-battery"
89    );
90
91    let mut diagnostics = tempfile::tempfile()?;
92    let mut command = Command::new(&args[0]);
93
94    command.args(&args[1..]).current_dir(util::root());
95
96    for (name, _) in std::env::vars_os() {
97        if name.to_string_lossy().starts_with("CHERENKOV_") {
98            command.env_remove(name);
99        }
100    }
101
102    let started = Instant::now();
103    let mut child = ChildGuard(
104        command
105            .stdout(File::create(output)?)
106            .stderr(diagnostics.try_clone()?)
107            .spawn()?,
108    );
109    let mut samples = vec![before.clone()];
110    let mut last_poll = Instant::now();
111    let mut cycle = None;
112    let mut timed_out = false;
113    let code = loop {
114        ensure!(
115            !INTERRUPTED.load(Ordering::Relaxed),
116            "interrupted; current output retained, resume to retry"
117        );
118
119        if let Some(status) = child.0.try_wait()? {
120            break status.code().unwrap_or(-1);
121        }
122
123        if timeout.is_some_and(|t| started.elapsed() >= t) {
124            timed_out = true;
125
126            child.0.kill()?;
127
128            break child.0.wait()?.code().unwrap_or(-1);
129        }
130
131        if last_poll.elapsed() >= Duration::from_secs(30) {
132            samples.push(power());
133
134            let bytes = fs::read(output)?;
135            cycle = metrics::repeated_tail(&String::from_utf8_lossy(&bytes));
136
137            if cycle.is_some() {
138                child.0.kill()?;
139
140                break child.0.wait()?.code().unwrap_or(-1);
141            }
142
143            eprintln!(
144                "{:.0}s elapsed: {}",
145                started.elapsed().as_secs_f64(),
146                output.display()
147            );
148
149            last_poll = Instant::now();
150        }
151
152        std::thread::sleep(Duration::from_millis(100));
153    };
154
155    diagnostics.seek(SeekFrom::Start(0))?;
156
157    let mut stderr = String::new();
158
159    diagnostics.read_to_string(&mut stderr)?;
160
161    Ok(Captured {
162        code,
163        stderr,
164        wall_seconds: started.elapsed().as_secs_f64(),
165        power_before: before,
166        power_after: power(),
167        power_samples: samples,
168        cycle,
169        timed_out,
170    })
171}