atagen

roundabout

public
Star0
Watch0
9117440444eb609ecea8476b5a62b03917077591
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/run.rsRust392 lines12.0 KB
1
use std::{
2
    env,
3
    fmt::Write as _,
4
    os::unix::{net::UnixStream, process::CommandExt},
5
    process::{Command, ExitStatus},
6
    sync::{
7
        Arc,
8
        atomic::{AtomicBool, AtomicI32, Ordering},
9
    },
10
    thread,
11
    time::{Duration, Instant},
12
};
13
14
use crate::{
15
    artifacts, client,
16
    config::{Config, OverCap},
17
    history, paths,
18
    proto::{self, Acquire, Event, Reader, Request},
19
    sched::Need,
20
    sys,
21
    units::{Amount, Bytes, Dur},
22
};
23
24
pub const EXIT_TIMEOUT: i32 = 124;
25
pub const EXIT_GAVE_UP: i32 = 75;
26
27
#[derive(Default)]
28
pub struct Args {
29
    pub size: Option<String>,
30
    pub ram: Option<Bytes>,
31
    pub disk: Option<Bytes>,
32
    pub cores: Option<Amount<u32>>,
33
    pub time: Option<Dur>,
34
    pub deadline: Option<Dur>,
35
    pub max_wait: Option<Dur>,
36
    pub session: Option<String>,
37
    pub exact: bool,
38
    pub cmd: Vec<String>,
39
}
40
41
static PGID: AtomicI32 = AtomicI32::new(0);
42
#[cfg(target_os = "linux")]
43
static FORWARDED: AtomicI32 = AtomicI32::new(0);
44
45
extern "C" fn forward(sig: libc::c_int) {
46
    #[cfg(target_os = "linux")]
47
    FORWARDED.store(sig, Ordering::Relaxed);
48
    let pgid = PGID.load(Ordering::Relaxed);
49
    if pgid > 1 {
50
        unsafe { libc::kill(-pgid, sig) };
51
    }
52
}
53
54
fn say(msg: &str) {
55
    eprintln!("roundabout: {msg}");
56
}
57
58
pub fn need(cfg: &Config, a: &Args) -> Result<Need, String> {
59
    let base = cfg.preset(a.size.as_deref().unwrap_or(&cfg.enforce.default_size))?;
60
    Ok(Need {
61
        ram: a.ram.unwrap_or(base.ram),
62
        disk: a.disk.unwrap_or(base.disk),
63
        cores: a.cores.unwrap_or(base.cores).resolve(sys::cores()),
64
        time: a.time.unwrap_or(base.time).0,
65
    })
66
}
67
68
pub fn base(cfg: &Config) -> Result<Need, String> {
69
    let p = cfg.preset(&cfg.enforce.default_size)?;
70
    Ok(Need {
71
        ram: p.ram,
72
        disk: p.disk,
73
        cores: p.cores.resolve(sys::cores()),
74
        time: p.time.0,
75
    })
76
}
77
78
struct Grant {
79
    stream: UnixStream,
80
    lease: u64,
81
    cores: u32,
82
    ram: Bytes,
83
}
84
85
fn acquire(req: &Acquire, max_wait: Option<Dur>) -> Result<Grant, i32> {
86
    let started = Instant::now();
87
    // a daemon on its way out can accept us then hang up, so go round again
88
    for _ in 0..3 {
89
        let mut stream = client::connect_or_spawn().map_err(|e| {
90
            say(&format!("can't reach the daemon: {e}"));
91
            1
92
        })?;
93
        if proto::send(&mut stream, &Request::Acquire(req.clone())).is_err() {
94
            continue;
95
        }
96
        let _ = stream.set_read_timeout(Some(Duration::from_secs(1)));
97
        let mut reader = Reader::new(&stream).map_err(|_| 1)?;
98
        loop {
99
            match reader.recv::<Event>() {
100
                Ok(Some(Event::Granted {
101
                    lease,
102
                    cores,
103
                    ram,
104
                    heavy,
105
                })) => {
106
                    let bonus = cores.saturating_sub(req.need.cores);
107
                    say(&format!(
108
                        "lease {lease} go, cores {cores}{}{}",
109
                        if bonus > 0 {
110
                            format!(" ({bonus} spare)")
111
                        } else {
112
                            String::new()
113
                        },
114
                        if heavy { ", heavy" } else { "" }
115
                    ));
116
                    let _ = stream.set_read_timeout(None);
117
                    return Ok(Grant {
118
                        stream,
119
                        lease,
120
                        cores,
121
                        ram,
122
                    });
123
                }
124
                Ok(Some(Event::Queued {
125
                    lease,
126
                    position,
127
                    reason,
128
                })) => {
129
                    let why = reason.map(|r| format!(", {r}")).unwrap_or_default();
130
                    say(&format!("lease {lease} queued at #{position}{why}"));
131
                }
132
                Ok(Some(Event::Kill)) => {
133
                    say("killed while queued");
134
                    return Err(143);
135
                }
136
                Ok(Some(_)) => {}
137
                Err(e)
138
                    if matches!(
139
                        e.kind(),
140
                        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
141
                    ) =>
142
                {
143
                    if max_wait.is_some_and(|m| started.elapsed() > m.0) {
144
                        say("gave up waiting");
145
                        return Err(EXIT_GAVE_UP);
146
                    }
147
                }
148
                Ok(None) | Err(_) => break,
149
            }
150
        }
151
    }
152
    say("daemon keeps hanging up");
153
    Err(1)
154
}
155
156
fn command(cfg: &Config, a: &Args, grant: &Grant) -> Command {
157
    let n = grant.cores.to_string();
158
    let mut cmd = Command::new(&a.cmd[0]);
159
    cmd.args(&a.cmd[1..])
160
        .process_group(0)
161
        .env("ROUNDABOUT_LEASE", grant.lease.to_string())
162
        .env("CARGO_BUILD_JOBS", &n)
163
        .env("CMAKE_BUILD_PARALLEL_LEVEL", &n);
164
165
    let make = env::var("MAKEFLAGS").unwrap_or_default();
166
    if !make.contains("-j") {
167
        cmd.env("MAKEFLAGS", format!("-j{n} {make}").trim_end());
168
    }
169
170
    if cfg.nix.inject {
171
        let per = cfg.nix.cores_per_job.max(1);
172
        let jobs = (grant.cores / per).max(1);
173
        let cores = grant.cores.min(per);
174
        let mut nix = env::var("NIX_CONFIG").unwrap_or_default();
175
        let _ = write!(nix, "\ncores = {cores}\nmax-jobs = {jobs}\n");
176
        cmd.env("NIX_CONFIG", nix);
177
    }
178
179
    cmd
180
}
181
182
fn cpu_time() -> Duration {
183
    let mut ru = std::mem::MaybeUninit::<libc::rusage>::uninit();
184
    unsafe {
185
        if libc::getrusage(libc::RUSAGE_CHILDREN, ru.as_mut_ptr()) != 0 {
186
            return Duration::ZERO;
187
        }
188
        let ru = ru.assume_init();
189
        let tv = |t: libc::timeval| {
190
            Duration::from_secs(t.tv_sec as u64) + Duration::from_micros(t.tv_usec as u64)
191
        };
192
        tv(ru.ru_utime) + tv(ru.ru_stime)
193
    }
194
}
195
196
fn code(status: ExitStatus) -> i32 {
197
    use std::os::unix::process::ExitStatusExt;
198
    status
199
        .code()
200
        .or_else(|| status.signal().map(|s| 128 + s))
201
        .unwrap_or(1)
202
}
203
204
pub fn run(cfg: &Config, a: &Args) -> i32 {
205
    if a.cmd.is_empty() {
206
        say("nothing to run, put the command after --");
207
        return 2;
208
    }
209
    if a.size.is_none() && a.ram.is_none() {
210
        let sizes: Vec<_> = cfg.sizes.keys().map(String::as_str).collect();
211
        say(&format!(
212
            "how big is it? pass --size {} or --ram and friends. `roundabout suggest -- <cmd>` has a guess",
213
            sizes.join("|")
214
        ));
215
        return 2;
216
    }
217
    let mut need = match need(cfg, a) {
218
        Ok(n) => n,
219
        Err(e) => {
220
            say(&e);
221
            return 2;
222
        }
223
    };
224
    let cwd = env::current_dir().unwrap_or_else(|_| "/".into());
225
    let root = paths::project_root(&cwd);
226
227
    let report = artifacts::check(cfg, Some(&root), false);
228
    if report.over() {
229
        let msg = report.describe(cfg);
230
        if cfg.artifacts.over_cap == OverCap::Block {
231
            say(&format!("{msg}. clean up first, `{}`", report.remedy()));
232
            return EXIT_GAVE_UP;
233
        }
234
        say(&format!("{msg}. `{}` when you get a sec", report.remedy()));
235
    }
236
237
    let req = Acquire {
238
        need,
239
        cwd,
240
        cmd: a.cmd.clone(),
241
        session: a.session.clone(),
242
        deadline: a.deadline,
243
        exact: a.exact,
244
    };
245
246
    let grant = match acquire(&req, a.max_wait.or(cfg.queue.max_wait)) {
247
        Ok(g) => g,
248
        Err(code) => return code,
249
    };
250
    need.ram = need.ram.max(grant.ram);
251
252
    let command = command(cfg, a, &grant);
253
    #[cfg(target_os = "linux")]
254
    let (mut child, scope) =
255
        match crate::systemd::spawn(&command, grant.lease, grant.cores, grant.ram, &cfg.overrun) {
256
            Ok(spawned) => spawned,
257
            Err(error) => {
258
                say(&format!("couldn't start {}: {error}", a.cmd[0]));
259
                return 127;
260
            }
261
        };
262
    #[cfg(not(target_os = "linux"))]
263
    let mut command = command;
264
    #[cfg(not(target_os = "linux"))]
265
    let mut child = match command.spawn() {
266
        Ok(c) => c,
267
        Err(e) => {
268
            say(&format!("couldn't start {}: {e}", a.cmd[0]));
269
            return 127;
270
        }
271
    };
272
    let pgid = i32::try_from(child.id()).expect("pid fits in pid_t");
273
    PGID.store(pgid, Ordering::Relaxed);
274
    for sig in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP] {
275
        unsafe { libc::signal(sig, forward as *const () as libc::sighandler_t) };
276
    }
277
278
    let mut out = grant.stream.try_clone().ok();
279
    if let Some(s) = &mut out {
280
        let _ = proto::send(s, &Request::Started { pid: pgid });
281
    }
282
283
    let killed = Arc::new(AtomicBool::new(false));
284
    if let Ok(mut reader) = Reader::new(&grant.stream) {
285
        let killed = killed.clone();
286
        thread::spawn(move || {
287
            while let Ok(Some(ev)) = reader.recv::<Event>() {
288
                if matches!(ev, Event::Kill) {
289
                    killed.store(true, Ordering::Relaxed);
290
                }
291
            }
292
        });
293
    }
294
295
    let started = Instant::now();
296
    let grace = need.time.mul_f64(cfg.overrun.time_grace);
297
    let loud_ram = Bytes((need.ram.0 as f64 * cfg.overrun.ram_slack) as u64);
298
    let mut peak = Bytes::ZERO;
299
    let mut last_sample: Option<Instant> = None;
300
    let mut term_sent: Option<Instant> = None;
301
    let mut timed_out = false;
302
    let mut warned_ram = false;
303
    let mut warned_time = false;
304
305
    let status = loop {
306
        #[cfg(target_os = "linux")]
307
        if let signal @ 1.. = FORWARDED.swap(0, Ordering::Relaxed) {
308
            let _ = scope.signal(signal);
309
        }
310
        match child.try_wait() {
311
            Ok(Some(status)) => break status,
312
            Ok(None) => {}
313
            Err(e) => {
314
                say(&format!("lost track of the child: {e}"));
315
                return 1;
316
            }
317
        }
318
        let elapsed = started.elapsed();
319
320
        if term_sent.is_none() {
321
            if killed.load(Ordering::Relaxed) {
322
                say("killed on request");
323
                term_sent = Some(Instant::now());
324
            } else if a.deadline.is_some_and(|d| elapsed > d.0) {
325
                say("deadline hit, killing it");
326
                timed_out = true;
327
                term_sent = Some(Instant::now());
328
            }
329
            if term_sent.is_some() {
330
                #[cfg(target_os = "linux")]
331
                let _ = scope.signal(libc::SIGTERM);
332
                unsafe { libc::kill(-pgid, libc::SIGTERM) };
333
            }
334
        } else if term_sent.is_some_and(|t| t.elapsed() > Duration::from_secs(10)) {
335
            #[cfg(target_os = "linux")]
336
            let _ = scope.signal(libc::SIGKILL);
337
            unsafe { libc::kill(-pgid, libc::SIGKILL) };
338
        }
339
340
        if last_sample.is_none_or(|t| t.elapsed() >= Duration::from_secs(2)) {
341
            last_sample = Some(Instant::now());
342
            let rss = sys::group_rss(pgid);
343
            peak = peak.max(rss);
344
            if let Some(s) = &mut out {
345
                let _ = proto::send(s, &Request::Usage { rss });
346
            }
347
            if !warned_ram && rss > loud_ram {
348
                warned_ram = true;
349
                say(&format!("using {rss} against a {} estimate", need.ram));
350
            }
351
        }
352
        if !warned_time && elapsed > grace {
353
            warned_time = true;
354
            say(&format!(
355
                "{} past its {} estimate and still going",
356
                Dur(elapsed),
357
                Dur(need.time)
358
            ));
359
        }
360
        thread::sleep(Duration::from_millis(250));
361
    };
362
363
    PGID.store(0, Ordering::Relaxed);
364
    #[cfg(target_os = "linux")]
365
    drop(scope);
366
    let wall = started.elapsed();
367
    if let Some(s) = &mut out {
368
        let _ = proto::send(s, &Request::Done);
369
    }
370
    history::record(
371
        history::key(&root, &a.cmd),
372
        peak,
373
        wall,
374
        cpu_time(),
375
        grant.cores,
376
    );
377
    artifacts::register(cfg, &root, true);
378
    if peak > loud_ram || wall > grace {
379
        say(&format!(
380
            "took {} and peaked at {}, asked for {} and {}. ask for more next time",
381
            Dur(wall),
382
            peak,
383
            Dur(need.time),
384
            need.ram
385
        ));
386
    }
387
    if timed_out {
388
        EXIT_TIMEOUT
389
    } else {
390
        code(status)
391
    }
392
}