use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
pub struct Job {
pub index: usize,
pub path: std::path::PathBuf,
}
pub struct Done {
pub index: usize,
pub buffer: termcolor::Buffer,
pub outcome: crate::archive::Outcome,
}
pub struct Config {
pub workers: usize,
pub sorted: bool,
pub spill_bytes: usize,
}
struct Shared {
slots: Vec<Option<Done>>,
ready: VecDeque<usize>,
pool: Vec<termcolor::Buffer>,
held: usize,
batch_bytes: usize,
finished: usize,
live_workers: usize,
}
struct Completion<'a> {
state: &'a Mutex<Shared>,
wake: &'a Condvar,
armed: bool,
}
impl Drop for Completion<'_> {
fn drop(&mut self) {
if !self.armed {
return;
}
if let Ok(mut st) = self.state.lock() {
st.finished += 1;
}
self.wake.notify_all();
}
}
struct WorkerLife<'a> {
state: &'a Mutex<Shared>,
wake: &'a Condvar,
}
impl Drop for WorkerLife<'_> {
fn drop(&mut self) {
if let Ok(mut st) = self.state.lock() {
st.live_workers -= 1;
}
self.wake.notify_all();
}
}
struct Teardown<'a> {
state: &'a Mutex<Shared>,
space: &'a Condvar,
wake: &'a Condvar,
stop: &'a AtomicBool,
}
impl Drop for Teardown<'_> {
fn drop(&mut self) {
if let Ok(mut st) = self.state.lock() {
st.batch_bytes = 0;
self.stop.store(true, Ordering::Release);
} else {
self.stop.store(true, Ordering::Release);
}
self.space.notify_all();
self.wake.notify_all();
}
}
fn run_job<C, F>(
job: &Job,
ctx: &mut C,
buf: &mut termcolor::Buffer,
work: &F,
) -> crate::archive::Outcome
where
F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome,
{
let called = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| work(job, ctx, buf)));
match called {
Ok(o) => o,
Err(_) => {
buf.clear();
let mut o = crate::archive::Outcome::default();
o.errors.push(format!(
"{}: panicked while searching; skipped",
job.path.display()
));
o
}
}
}
pub fn run<C, F, E>(
jobs: Vec<Job>,
cfg: Config,
make_buffer: impl Fn() -> termcolor::Buffer + Send + Sync,
make_worker: impl Fn() -> C + Send + Sync,
work: F,
emit: E,
) -> std::io::Result<()>
where
F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome + Send + Sync,
E: FnMut(&Done) -> std::io::Result<()>,
{
run_with_spawn_hook(jobs, cfg, make_buffer, make_worker, work, emit, |_| {})
}
fn run_with_spawn_hook<C, F, E>(
jobs: Vec<Job>,
cfg: Config,
make_buffer: impl Fn() -> termcolor::Buffer + Send + Sync,
make_worker: impl Fn() -> C + Send + Sync,
work: F,
mut emit: E,
before_spawn: impl Fn(usize),
) -> std::io::Result<()>
where
F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome + Send + Sync,
E: FnMut(&Done) -> std::io::Result<()>,
{
let Config { workers, sorted, spill_bytes } = cfg;
let n = jobs.len();
if n == 0 {
return Ok(());
}
if workers <= 1 {
let mut ctx = make_worker();
let mut buf = make_buffer();
for job in &jobs {
buf.clear();
let outcome = run_job(job, &mut ctx, &mut buf, &work);
let done = Done { index: job.index, buffer: buf, outcome };
let r = emit(&done);
buf = done.buffer;
r?;
if buf.as_slice().len() > spill_bytes {
buf = make_buffer();
}
}
return Ok(());
}
let mut order: Vec<usize> = (0..n).collect();
if sorted {
order.sort_by_key(|&p| jobs[p].index);
}
let live = workers.min(n);
let next = AtomicUsize::new(0);
let stop = AtomicBool::new(false);
let state = Mutex::new(Shared {
slots: (0..n).map(|_| None).collect(),
ready: VecDeque::new(),
pool: Vec::new(),
held: 0,
batch_bytes: 0,
finished: 0,
live_workers: live,
});
let wake = Condvar::new();
let space = Condvar::new();
let mut write_err: Option<std::io::Error> = None;
std::thread::scope(|scope| {
let _teardown = Teardown { state: &state, space: &space, wake: &wake, stop: &stop };
for i in 0..live {
before_spawn(i);
let (next, stop, state, wake) = (&next, &stop, &state, &wake);
let space = &space;
let (jobs, work, make_buffer, make_worker) =
(&jobs, &work, &make_buffer, &make_worker);
scope.spawn(move || {
let _life = WorkerLife { state, wake };
let mut ctx = match std::panic::catch_unwind(
std::panic::AssertUnwindSafe(&make_worker),
) {
Ok(c) => c,
Err(_) => return,
};
loop {
{
let mut st = state.lock().unwrap();
while !stop.load(Ordering::Acquire)
&& st.held + st.batch_bytes > spill_bytes
{
st = space.wait(st).unwrap();
}
}
if stop.load(Ordering::Acquire) {
return;
}
let p = next.fetch_add(1, Ordering::Relaxed);
if p >= n {
return;
}
let mut guard = Completion { state, wake, armed: true };
let mut buf = {
let mut st = state.lock().unwrap();
st.pool.pop()
}
.unwrap_or_else(make_buffer);
let outcome = run_job(&jobs[p], &mut ctx, &mut buf, work);
let bytes = buf.as_slice().len();
let done = Done { index: jobs[p].index, buffer: buf, outcome };
{
let mut st = state.lock().unwrap();
st.slots[p] = Some(done);
st.ready.push_back(p);
st.held += bytes;
st.finished += 1;
guard.armed = false;
}
wake.notify_all();
}
});
}
let mut head = 0usize;
let mut emitted = vec![false; n];
let mut emitted_count = 0usize;
let mut batch: Vec<Done> = Vec::new();
'collect: while emitted_count < n {
let mut spill_notices = 0usize;
{
let mut st = state.lock().unwrap();
loop {
if sorted {
while head < n {
if emitted[head] {
head += 1;
continue;
}
match st.slots[order[head]].take() {
Some(d) => {
let len = d.buffer.as_slice().len();
st.held -= len;
st.batch_bytes += len;
emitted[head] = true;
head += 1;
batch.push(d);
}
None => break,
}
}
if st.held > spill_bytes {
spill_notices += 1;
flush_parked(&mut st, &order, &mut emitted, head, &mut batch);
}
} else {
while let Some(p) = st.ready.pop_front() {
if let Some(d) = st.slots[p].take() {
let len = d.buffer.as_slice().len();
st.held -= len;
st.batch_bytes += len;
emitted[p] = true;
batch.push(d);
}
}
}
if !batch.is_empty() {
break;
}
debug_assert!(st.finished < n, "collector stalled with all work finished");
if st.finished >= n {
break;
}
if st.live_workers == 0 {
for (h, flag) in emitted.iter_mut().enumerate() {
if *flag {
continue;
}
*flag = true;
let p = order[h];
match st.slots[p].take() {
Some(d) => {
let len = d.buffer.as_slice().len();
st.held -= len;
st.batch_bytes += len;
batch.push(d);
}
None => batch.push(Done {
index: jobs[p].index,
buffer: make_buffer(),
outcome: unclaimed(&jobs[p]),
}),
}
}
break;
}
st = wake.wait(st).unwrap();
}
}
for _ in 0..spill_notices {
eprintln!(
"trg: held output exceeded {spill_bytes} bytes; flushing out of order"
);
}
if batch.is_empty() {
break;
}
for d in batch.drain(..) {
emitted_count += 1;
let r = emit(&d);
recycle(&state, &space, d.buffer, spill_bytes);
if let Err(e) = r {
stop.store(true, Ordering::Release);
write_err = Some(e);
break 'collect;
}
}
}
});
match write_err {
Some(e) => Err(e),
None => Ok(()),
}
}
fn unclaimed(job: &Job) -> crate::archive::Outcome {
let mut o = crate::archive::Outcome::default();
o.errors.push(format!(
"{}: every worker exited before this archive was searched; not read",
job.path.display()
));
o
}
fn recycle(
state: &Mutex<Shared>,
space: &Condvar,
mut buffer: termcolor::Buffer,
spill_bytes: usize,
) {
let len = buffer.as_slice().len();
let oversized = len > spill_bytes;
buffer.clear();
if let Ok(mut st) = state.lock() {
st.batch_bytes -= len;
if !oversized {
st.pool.push(buffer);
}
}
space.notify_all();
}
fn flush_parked(
st: &mut Shared,
order: &[usize],
emitted: &mut [bool],
head: usize,
batch: &mut Vec<Done>,
) {
for (h, flag) in emitted.iter_mut().enumerate().skip(head) {
if let Some(d) = st.slots[order[h]].take() {
let len = d.buffer.as_slice().len();
st.held -= len;
st.batch_bytes += len;
*flag = true;
batch.push(d);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[test]
fn a_spawn_failure_does_not_strand_a_parked_worker() {
const PATIENCE: Duration = Duration::from_secs(10);
let (tx, rx) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let jobs = (0..4)
.map(|i| Job { index: i, path: std::path::PathBuf::from(format!("j{i}")) })
.collect();
let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
run_with_spawn_hook(
jobs,
Config { workers: 4, sorted: true, spill_bytes: 8 },
termcolor::Buffer::no_color,
|| (),
|_job, _c: &mut (), b| {
use std::io::Write;
write!(b, "{}", "y".repeat(100)).unwrap();
crate::archive::Outcome::default()
},
|_done: &Done| Ok(()),
|i| {
if i == 1 {
panic!("injected: the OS refused thread {i}");
}
},
)
}));
let _ = tx.send(r.is_err());
});
match rx.recv_timeout(PATIENCE) {
Ok(panicked) => assert!(panicked, "the injected spawn failure must propagate"),
Err(_) => panic!("run hung after a spawn failure: a parked worker was never released"),
}
}
}