use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::time::Duration;
use trg::sched::{Done, Job};
fn jobs(n: usize) -> Vec<Job> {
(0..n).map(|i| Job { index: i, path: format!("j{i}").into() }).collect()
}
fn buf() -> termcolor::Buffer {
termcolor::Buffer::no_color()
}
fn within<T: Send + 'static>(secs: u64, f: impl FnOnce() -> T + Send + 'static) -> T {
let (tx, rx) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let _ = tx.send(f());
});
match rx.recv_timeout(Duration::from_secs(secs)) {
Ok(v) => v,
Err(e) => panic!("did not finish within {secs}s ({e:?}) — deadlock or missing wakeup"),
}
}
#[derive(Default)]
struct InFlight {
now: AtomicUsize,
peak: AtomicUsize,
}
impl InFlight {
fn enter(&self) {
let n = self.now.fetch_add(1, Ordering::SeqCst) + 1;
self.peak.fetch_max(n, Ordering::SeqCst);
}
fn leave(&self) {
self.now.fetch_sub(1, Ordering::SeqCst);
}
}
#[test]
fn a_gate_of_one_runs_one_job_at_a_time_and_loses_nothing() {
let (peak, seen) = within(30, || {
let gate = AtomicUsize::new(1);
let fl = InFlight::default();
let seen = Mutex::new(Vec::new());
trg::sched::run_gated(
jobs(24),
trg::sched::Config { workers: 4, sorted: true, spill_bytes: 64 << 20 },
&gate,
buf,
|| (),
|_job, _c: &mut (), _b| {
fl.enter();
std::thread::sleep(Duration::from_millis(2));
fl.leave();
trg::archive::Outcome::default()
},
|d: &Done| {
seen.lock().unwrap().push(d.index);
Ok(())
},
)
.unwrap();
(fl.peak.load(Ordering::SeqCst), seen.into_inner().unwrap())
});
assert_eq!(peak, 1, "a gate of 1 let {peak} jobs run at once");
assert_eq!(seen, (0..24).collect::<Vec<_>>(), "order and completeness hold under the gate");
}
#[test]
fn raising_the_gate_mid_run_lets_more_workers_in() {
let peak = within(30, || {
let gate = Arc::new(AtomicUsize::new(1));
let fl = InFlight::default();
let done = AtomicUsize::new(0);
let g = Arc::clone(&gate);
trg::sched::run_gated(
jobs(40),
trg::sched::Config { workers: 4, sorted: true, spill_bytes: 64 << 20 },
&gate,
buf,
|| (),
|_job, _c: &mut (), _b| {
fl.enter();
if done.fetch_add(1, Ordering::SeqCst) == 4 {
g.store(4, Ordering::SeqCst);
}
std::thread::sleep(Duration::from_millis(20));
fl.leave();
trg::archive::Outcome::default()
},
|_d: &Done| Ok(()),
)
.unwrap();
fl.peak.load(Ordering::SeqCst)
});
assert!(peak > 1, "the gate opened to 4 but jobs still ran one at a time");
}
#[test]
fn lowering_the_gate_mid_run_holds_workers_back() {
let late_peak = within(30, || {
const FIRST_WAVE: usize = 4;
let gate = AtomicUsize::new(FIRST_WAVE);
let wave = (Mutex::new(0usize), Condvar::new());
let late = InFlight::default();
trg::sched::run_gated(
jobs(30),
trg::sched::Config { workers: FIRST_WAVE, sorted: false, spill_bytes: 64 << 20 },
&gate,
buf,
|| (),
|job, _c: &mut (), _b| {
if job.index < FIRST_WAVE {
let (m, cv) = &wave;
let mut n = m.lock().unwrap();
*n += 1;
if *n == FIRST_WAVE {
gate.store(1, Ordering::SeqCst);
cv.notify_all();
}
let (_n, timeout) = cv
.wait_timeout_while(n, Duration::from_secs(10), |n| *n < FIRST_WAVE)
.unwrap();
assert!(!timeout.timed_out(), "the first wave never filled");
} else {
late.enter();
std::thread::sleep(Duration::from_millis(5));
late.leave();
}
trg::archive::Outcome::default()
},
|_d: &Done| Ok(()),
)
.unwrap();
late.peak.load(Ordering::SeqCst)
});
assert_eq!(late_peak, 1, "the gate closed to 1 but {late_peak} jobs ran at once afterwards");
}
#[test]
fn an_early_stop_does_not_wait_on_held_back_workers() {
let r = within(10, || {
let gate = AtomicUsize::new(1);
trg::sched::run_gated(
jobs(50),
trg::sched::Config { workers: 8, sorted: true, spill_bytes: 64 << 20 },
&gate,
buf,
|| (),
|_job, _c: &mut (), _b| trg::archive::Outcome::default(),
|_d: &Done| Err(std::io::Error::from(std::io::ErrorKind::BrokenPipe)),
)
});
assert_eq!(r.unwrap_err().kind(), std::io::ErrorKind::BrokenPipe);
}
#[test]
fn a_dead_worker_does_not_hold_the_only_place_at_the_gate() {
let seen = within(10, || {
let gate = AtomicUsize::new(1);
let built = AtomicUsize::new(0);
let seen = Mutex::new(Vec::new());
trg::sched::run_gated(
jobs(12),
trg::sched::Config { workers: 4, sorted: true, spill_bytes: 64 << 20 },
&gate,
buf,
|| {
if built.fetch_add(1, Ordering::SeqCst) == 0 {
panic!("injected: the first worker fails to start");
}
},
|_job, _c: &mut (), _b| trg::archive::Outcome::default(),
|d: &Done| {
seen.lock().unwrap().push(d.index);
Ok(())
},
)
.unwrap();
seen.into_inner().unwrap()
});
assert_eq!(seen, (0..12).collect::<Vec<_>>());
}
#[test]
fn a_gate_of_zero_still_makes_progress() {
let n = within(10, || {
let gate = AtomicUsize::new(0);
let n = AtomicUsize::new(0);
trg::sched::run_gated(
jobs(6),
trg::sched::Config { workers: 3, sorted: true, spill_bytes: 64 << 20 },
&gate,
buf,
|| (),
|_job, _c: &mut (), _b| trg::archive::Outcome::default(),
|_d: &Done| {
n.fetch_add(1, Ordering::SeqCst);
Ok(())
},
)
.unwrap();
n.into_inner()
});
assert_eq!(n, 6);
}
#[test]
fn a_limit_below_the_load_changes_the_pace_not_the_result() {
#[path = "helpers/mod.rs"]
mod helpers;
let d = tempfile::tempdir().unwrap();
let mut args: Vec<String> = ["-j", "4", "--load-limit", "-1", "NEEDLE"].map(String::from).to_vec();
for i in 0..6 {
let body = format!("NEEDLE {i}\n");
let f = helpers::tgz("x.tgz", &[("logs/a.log", body.as_bytes())]);
let dest = d.path().join(format!("{i}.tgz"));
std::fs::copy(&f.path, &dest).unwrap();
args.push(dest.to_str().unwrap().to_owned());
}
let o = std::process::Command::new(env!("CARGO_BIN_EXE_trg")).args(&args).output().unwrap();
let out = String::from_utf8_lossy(&o.stdout);
let got: Vec<String> = out.lines().filter_map(|l| l.rsplit(':').next()).map(String::from).collect();
let want: Vec<String> = (0..6).map(|i| format!("NEEDLE {i}")).collect();
assert_eq!(got, want, "argument order, every archive");
assert_eq!(o.status.code(), Some(0));
assert!(o.stderr.is_empty(), "got: {}", String::from_utf8_lossy(&o.stderr));
}