trg/sched.rs
1//! Bounded worker pool. Concurrency lives at the archive boundary and nowhere
2//! else, because gzip has no index: member N requires inflating 1..N-1.
3//!
4//! # Ordering
5//!
6//! Output stays in argument order by default: these archives are days, and
7//! argument order is chronological order. That is achieved by holding a
8//! completed archive's buffer until every earlier archive has been emitted —
9//! **not** by joining all workers first. A finished job whose predecessors are
10//! already out is written and freed immediately, so nothing waits on the
11//! slowest archive to be handed to `emit`.
12//!
13//! # What actually bounds memory
14//!
15//! Three separate things, and it is worth being precise about which does what,
16//! because `spill_bytes` alone does far less than its name suggests:
17//!
18//! 1. **Buffer recycling.** Output buffers are owned by this module and pooled.
19//! A worker takes one, fills it, and the collector `clear()`s it after
20//! emitting and returns it to the pool. Allocating a fresh buffer per
21//! archive instead makes peak RSS climb with the *archive count*: each one
22//! is grown by a realloc chain to that archive's whole match volume and then
23//! freed, and the allocator does not hand the same span back for the next
24//! one. Measured at `-j1` on six archives of 120k matching lines: 84 MB
25//! per-job vs. 22 MB pooled.
26//!
27//! At `workers == 1` exactly one buffer exists for the whole run, which is
28//! the case that matters most since `-j1` is the default. In the pool path
29//! the *count* of live buffers is not capped, and deliberately so: holding
30//! `k` results in order requires `k` buffers, so a hard cap would be a cap
31//! on the ordering window. It is not needed either — an unwritten
32//! `termcolor::Buffer` is an empty `Vec` that has allocated nothing, and the
33//! ones that did allocate are bounded in *bytes* by `spill_bytes` below.
34//! What the pool guarantees is that a buffer's allocation is reused rather
35//! than rebuilt, which is what the realloc chain cost.
36//! 2. **Per-worker state.** `make_worker` is called once per worker thread, not
37//! once per job, for the same reason — see [`run`].
38//! 3. **`spill_bytes`.** This bounds only the buffers *parked* in `slots`
39//! waiting for their turn in the order. It does **not** bound peak RSS: the
40//! `workers` buffers currently being filled, and the batch the collector has
41//! drained but not yet emitted, are outside the count. Treat it as "how far
42//! out of order are we willing to hold results", not as a memory ceiling. A
43//! real run peaked at 286 MB under a 64 MB cap.
44//!
45//! When parked buffers do exceed `spill_bytes`, they are flushed out of order
46//! with a note on stderr. Holding results in order is worth a lot, but not
47//! unboundedly much when a single slow archive is pinning the rest.
48//!
49//! # Design
50//!
51//! One `Mutex` guards the slot vector, the buffer pool, and a little
52//! bookkeeping; one `Condvar` wakes the collector. Workers only ever *deposit*:
53//! pull the next index with a `fetch_add`, do the work outside the lock, then
54//! take the lock just long enough to move the finished [`Done`] into its slot
55//! and notify. Nothing a worker does can block on another worker.
56//!
57//! The calling thread is the collector. It owns `head` and `emitted`, neither
58//! of which is shared, and it calls `emit` *outside* the lock so a slow stdout
59//! — a pipe into `head(1)`, say — never stalls a worker mid-archive.
60//!
61//! Emission order is correct because a slot is written exactly once (position
62//! `p` is handed to exactly one worker by the atomic counter) and read exactly
63//! once (only the collector takes, and only from a slot it then marks emitted).
64//! In sorted mode the collector walks `order`, so it emits the job with the
65//! next-lowest [`Job::index`] only once every lower index has gone out. The one
66//! deliberate exception is the spill above.
67//!
68//! # The load gate
69//!
70//! [`run_gated`] takes a gate: how many workers may be inside a job at once,
71//! read before every claim and free to move while the run goes on. That is
72//! what makes `--load-limit` adaptive — every `-j` worker is started, and the
73//! ones the gate leaves out wait between archives, holding nothing.
74//!
75//! Places at the gate are counted, not assigned to workers: a gate of 1 admits
76//! whichever worker gets there first, never "worker 0". An assigned place
77//! would belong to a worker that might be dead, with the rest waiting behind
78//! it for good. A place is released by a drop guard, so a worker that unwinds
79//! gives its place back too.
80//!
81//! A worker held at the gate polls rather than waiting on a condvar. The gate
82//! is moved by a thread outside the pool that knows nothing of its locks, and
83//! a poll cannot miss a wakeup: it sees a raised gate, a stop, or the end of
84//! the job list within [`GATE_POLL`], whoever caused it. It costs nothing when
85//! no gate is given, and a few wakeups a second while one is closed.
86//!
87//! # Panics and early stop
88//!
89//! A panicking job must not be able to hang the process — during an incident a
90//! hang looks like a slow search and gets waited on, which is strictly worse
91//! than a crash. Four independent mechanisms guarantee it cannot, one per place
92//! a thread can leave unexpectedly:
93//!
94//! - [`catch_unwind`](std::panic::catch_unwind) around the call to `work`,
95//! which turns a panic into an ordinary error on that archive, and a second
96//! one around `make_worker`, which is likewise the caller's closure.
97//! - a `Completion` drop guard that advances the finished count and wakes the
98//! collector even if a worker unwinds anyway, once it has claimed a job.
99//! - a `WorkerLife` drop guard covering the rest of the thread's life —
100//! everything before and between jobs, where `Completion` does not exist yet.
101//! When the last worker goes with jobs still unclaimed, the collector stops
102//! waiting and reports those archives as unread, so the run is a 2.
103//! - a `Teardown` drop guard on the collector, releasing workers parked on
104//! backpressure when it leaves, including by unwinding out of `emit`.
105//!
106//! Each guards a different exit; none of them is redundant with the others, and
107//! the module has hit this failure shape three times, so a new path out of a
108//! thread needs a guard rather than an argument that it cannot be taken.
109//!
110//! `emit` returns `io::Result`, and the first error stops the run: no further
111//! job is started and [`run`] returns the error. That is what makes
112//! `trg ... | head -5` stop searching instead of grinding through a whole
113//! month of archives writing into a closed pipe.
114
115use std::collections::VecDeque;
116use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
117use std::sync::{Condvar, Mutex};
118use std::time::Duration;
119
120/// How long a worker held at the gate sleeps between looks. Short against an
121/// archive, long against a busy loop.
122const GATE_POLL: Duration = Duration::from_millis(50);
123
124pub struct Job {
125 pub index: usize,
126 pub path: std::path::PathBuf,
127}
128
129pub struct Done {
130 pub index: usize,
131 pub buffer: termcolor::Buffer,
132 pub outcome: crate::archive::Outcome,
133}
134
135/// How the pool runs, as opposed to what it runs.
136pub struct Config {
137 /// Archives searched concurrently. `1` takes a dedicated sequential path.
138 pub workers: usize,
139 /// Emit in [`Job::index`] order rather than completion order.
140 pub sorted: bool,
141 /// How many bytes of *parked* output may wait for their turn before the
142 /// order is abandoned. Not a memory ceiling — see the module docs.
143 pub spill_bytes: usize,
144}
145
146/// Everything the workers and the collector share. The lock is held only for
147/// the moves in and out of `slots` and `pool`.
148struct Shared {
149 /// One slot per job position. `Some` means finished and not yet emitted.
150 slots: Vec<Option<Done>>,
151 /// Positions in completion order, for `--no-sort`.
152 ready: VecDeque<usize>,
153 /// Output buffers not currently in use, kept so their allocations are
154 /// reused instead of being rebuilt per archive.
155 pool: Vec<termcolor::Buffer>,
156 /// Bytes currently parked in `slots`, waiting for their turn in the order.
157 held: usize,
158 /// Bytes the collector has drained out of `slots` but not yet emitted and
159 /// recycled. Counted separately and added to `held` for backpressure,
160 /// because draining moves bytes out of `slots` without freeing anything —
161 /// missing this is what let live buffers track the job count.
162 batch_bytes: usize,
163 /// Jobs finished so far, so the collector knows when to stop waiting.
164 finished: usize,
165 /// Worker threads still running. Decremented by [`WorkerLife`] on every
166 /// exit path, normal or unwinding, so that "nobody is left to finish the
167 /// outstanding jobs" is a state the collector can observe rather than a
168 /// wait it never wakes from.
169 live_workers: usize,
170}
171
172/// Advances the finished count and wakes the collector on drop, so that a
173/// worker which unwinds still counts as having completed its job.
174///
175/// Without this, a panicking worker would leave `finished` short of `n`
176/// forever: the collector would wait on a notification that never comes, and
177/// `thread::scope` would never return. Belt and braces alongside the
178/// `catch_unwind` in [`run_job`], which already stops a panic in `work` itself;
179/// this covers everything else on the path.
180struct Completion<'a> {
181 state: &'a Mutex<Shared>,
182 wake: &'a Condvar,
183 armed: bool,
184}
185
186impl Drop for Completion<'_> {
187 fn drop(&mut self) {
188 if !self.armed {
189 return;
190 }
191 // Never `unwrap` here: panicking inside a `Drop` that is itself running
192 // during an unwind aborts the process. A poisoned lock is handled by
193 // waking the collector anyway, which will then fail on its own lock and
194 // propagate a panic — noisy, but not a hang.
195 if let Ok(mut st) = self.state.lock() {
196 st.finished += 1;
197 }
198 self.wake.notify_all();
199 }
200}
201
202/// Decrements the live-worker count and wakes the collector on drop.
203///
204/// [`Completion`] covers a worker that unwinds *after* it has claimed a job.
205/// This covers the rest of the thread's life, which [`Completion`] cannot: it
206/// is constructed as the thread's first act, before `make_worker()` and before
207/// the first `state.lock()`, so a panic in either — or in `space.wait()` on a
208/// poisoned lock — still tells the collector that this worker is gone.
209///
210/// Without it the collector waits at `wake.wait` for a notification no living
211/// thread will send, and `thread::scope` never returns: a hang, which during an
212/// incident looks like a slow search and gets waited on. This is the third
213/// instance of that shape in this module, after the panicking worker and the
214/// collector unwinding out of `emit`, so it is guarded structurally rather than
215/// by arguing that the current caller does not panic — `run` is public API and
216/// `make_worker` is the caller's closure.
217struct WorkerLife<'a> {
218 state: &'a Mutex<Shared>,
219 wake: &'a Condvar,
220}
221
222impl Drop for WorkerLife<'_> {
223 fn drop(&mut self) {
224 // Never `unwrap`: same discipline as `Completion` — this `Drop` runs
225 // during an unwind, and panicking there aborts the process.
226 if let Ok(mut st) = self.state.lock() {
227 st.live_workers -= 1;
228 }
229 self.wake.notify_all();
230 }
231}
232
233/// Releases every worker parked on backpressure when the collector leaves,
234/// *including* when it leaves by unwinding out of `emit`.
235///
236/// The exact mirror of [`Completion`] at the other end of the pipe. Without it,
237/// a panic in `emit` skips the teardown, workers stay parked on `space`, and
238/// `thread::scope` blocks forever joining them — a hang, which is the one
239/// outcome worse than a crash. Fixing that on the worker side and not the
240/// collector side would not be defensible.
241struct Teardown<'a> {
242 state: &'a Mutex<Shared>,
243 space: &'a Condvar,
244 wake: &'a Condvar,
245 stop: &'a AtomicBool,
246}
247
248impl Drop for Teardown<'_> {
249 fn drop(&mut self) {
250 // Never `unwrap` here: this `Drop` runs during an unwind, and a panic
251 // inside it would abort the process. Same discipline as `Completion`.
252 if let Ok(mut st) = self.state.lock() {
253 // The `Done`s still in the collector's batch are dropped without
254 // passing through `recycle`, so their bytes would be counted
255 // against the backpressure predicate forever. Zeroing here makes
256 // the predicate correct rather than merely short-circuited, so the
257 // wakeup below does not depend on `stop` alone.
258 st.batch_bytes = 0;
259 // Stored while the lock is held so the mutex's release/acquire edge
260 // carries it too; the `Release` ordering covers the lock-free check
261 // at the top of the worker loop.
262 self.stop.store(true, Ordering::Release);
263 } else {
264 self.stop.store(true, Ordering::Release);
265 }
266 self.space.notify_all();
267 self.wake.notify_all();
268 }
269}
270
271/// One place at the gate, given back on drop — unwinding included, so a dying
272/// worker cannot keep a place the others are waiting for.
273struct Permit<'a>(&'a AtomicUsize);
274
275impl Drop for Permit<'_> {
276 fn drop(&mut self) {
277 self.0.fetch_sub(1, Ordering::AcqRel);
278 }
279}
280
281/// Wait for a place at the gate. `None` means the worker should exit instead:
282/// the run has been stopped, or every job is already claimed. Checking those
283/// here, and not only after admission, is what keeps a held-back worker from
284/// outliving the run.
285fn admit<'a>(
286 gate: &AtomicUsize,
287 inside: &'a AtomicUsize,
288 stop: &AtomicBool,
289 next: &AtomicUsize,
290 n: usize,
291) -> Option<Permit<'a>> {
292 loop {
293 if stop.load(Ordering::Acquire) || next.load(Ordering::Relaxed) >= n {
294 return None;
295 }
296 // Clamped: a gate of 0 would admit nobody and the run would never end.
297 let cap = gate.load(Ordering::Relaxed).max(1);
298 let now = inside.load(Ordering::Acquire);
299 if now < cap {
300 if inside
301 .compare_exchange_weak(now, now + 1, Ordering::AcqRel, Ordering::Acquire)
302 .is_ok()
303 {
304 return Some(Permit(inside));
305 }
306 continue;
307 }
308 std::thread::sleep(GATE_POLL);
309 }
310}
311
312/// Call `work`, converting a panic into an error on that archive.
313///
314/// The buffer is cleared on panic: a job that died mid-write may have left a
315/// partial line behind, and half a log line in the middle of ordered output is
316/// worse than none.
317///
318/// This is a no-op under `panic = "abort"`, which the release profile sets. It
319/// is still worth having: an abort is a crash, and a crash is a diagnosable
320/// outcome. The failure mode being prevented is the *hang*, and the
321/// `Completion` guard prevents that on both profiles.
322fn run_job<C, F>(
323 job: &Job,
324 ctx: &mut C,
325 buf: &mut termcolor::Buffer,
326 work: &F,
327) -> crate::archive::Outcome
328where
329 F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome,
330{
331 // `AssertUnwindSafe` is honest here rather than a silencer: everything
332 // reachable is either this job's own scratch state, which is discarded on
333 // the panic path, or the shared slot vector, which this job has not touched
334 // yet — it deposits only after `work` returns.
335 let called = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| work(job, ctx, buf)));
336 match called {
337 Ok(o) => o,
338 Err(_) => {
339 buf.clear();
340 let mut o = crate::archive::Outcome::default();
341 o.errors.push(format!(
342 "{}: panicked while searching; skipped",
343 job.path.display()
344 ));
345 o
346 }
347 }
348}
349
350/// Run `work` over `jobs`, emitting each result through `emit`.
351///
352/// `make_worker` is called **once per worker thread**, not once per job. That
353/// is the correct granularity for anything a job needs but cannot share:
354/// `grep_searcher::Searcher` is not `Sync`, so it cannot be hoisted out of the
355/// pool entirely, but building one per *archive* is a real cost — its internal
356/// line buffer grows toward the configured heap limit and is thrown away every
357/// archive. Measured on 60 small archives: 288 MB per-job vs. 30 MB per-worker.
358/// At `workers == 1` this means exactly one is ever built.
359///
360/// `emit` borrows its [`Done`] rather than taking it, so the buffer can be
361/// recycled afterwards. The first `Err` it returns stops the run: no further
362/// job is started, and that error is returned.
363pub fn run<C, F, E>(
364 jobs: Vec<Job>,
365 cfg: Config,
366 make_buffer: impl Fn() -> termcolor::Buffer + Send + Sync,
367 make_worker: impl Fn() -> C + Send + Sync,
368 work: F,
369 emit: E,
370) -> std::io::Result<()>
371where
372 F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome + Send + Sync,
373 E: FnMut(&Done) -> std::io::Result<()>,
374{
375 run_with_spawn_hook(jobs, cfg, None, make_buffer, make_worker, work, emit, |_| {})
376}
377
378/// [`run`], with at most `gate` workers inside a job at any moment. The gate
379/// is re-read before every claim, so it may be moved during the run; values
380/// below 1 count as 1. See "The load gate" in the module docs.
381///
382/// At `workers == 1` the gate has nothing to hold back and is ignored.
383pub fn run_gated<C, F, E>(
384 jobs: Vec<Job>,
385 cfg: Config,
386 gate: &AtomicUsize,
387 make_buffer: impl Fn() -> termcolor::Buffer + Send + Sync,
388 make_worker: impl Fn() -> C + Send + Sync,
389 work: F,
390 emit: E,
391) -> std::io::Result<()>
392where
393 F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome + Send + Sync,
394 E: FnMut(&Done) -> std::io::Result<()>,
395{
396 run_with_spawn_hook(jobs, cfg, Some(gate), make_buffer, make_worker, work, emit, |_| {})
397}
398
399/// [`run`], with `before_spawn(i)` called just before worker `i` is spawned.
400///
401/// A fault-injection point and nothing else: a panic from it stands in for
402/// `scope.spawn` panicking when the OS refuses a thread, which no test can
403/// provoke for real. `run` passes a no-op, which compiles away.
404#[allow(clippy::too_many_arguments)]
405fn run_with_spawn_hook<C, F, E>(
406 jobs: Vec<Job>,
407 cfg: Config,
408 gate: Option<&AtomicUsize>,
409 make_buffer: impl Fn() -> termcolor::Buffer + Send + Sync,
410 make_worker: impl Fn() -> C + Send + Sync,
411 work: F,
412 mut emit: E,
413 before_spawn: impl Fn(usize),
414) -> std::io::Result<()>
415where
416 F: Fn(&Job, &mut C, &mut termcolor::Buffer) -> crate::archive::Outcome + Send + Sync,
417 E: FnMut(&Done) -> std::io::Result<()>,
418{
419 let Config { workers, sorted, spill_bytes } = cfg;
420 let n = jobs.len();
421 if n == 0 {
422 return Ok(());
423 }
424
425 // The sequential path holds nothing and needs no lock, which is already the
426 // bounded-memory ideal. Kept separate rather than folded into the pool:
427 // `-j1` is the production default and deserves no threads at all. One
428 // worker context and one buffer for the whole run.
429 if workers <= 1 {
430 let mut ctx = make_worker();
431 let mut buf = make_buffer();
432 for job in &jobs {
433 buf.clear();
434 let outcome = run_job(job, &mut ctx, &mut buf, &work);
435 let done = Done { index: job.index, buffer: buf, outcome };
436 let r = emit(&done);
437 buf = done.buffer;
438 r?;
439 // One pathological archive should not leave its whole match volume
440 // reserved for the rest of the run.
441 if buf.as_slice().len() > spill_bytes {
442 buf = make_buffer();
443 }
444 }
445 return Ok(());
446 }
447
448 // Emission order. Sorting by `Job::index` rather than by slot position
449 // means a caller that hands us indices which are not `0..n` still gets the
450 // order it asked for, and that `Done::index` is the one ordering key
451 // everywhere — including the spill loop. Stable, so equal indices keep
452 // argument order. For `main`, where index == position, this is the identity.
453 let mut order: Vec<usize> = (0..n).collect();
454 if sorted {
455 order.sort_by_key(|&p| jobs[p].index);
456 }
457
458 let live = workers.min(n);
459
460 let next = AtomicUsize::new(0);
461 let stop = AtomicBool::new(false);
462 // Workers currently holding a place at the gate. Unused without one.
463 let inside = AtomicUsize::new(0);
464 let state = Mutex::new(Shared {
465 slots: (0..n).map(|_| None).collect(),
466 ready: VecDeque::new(),
467 pool: Vec::new(),
468 held: 0,
469 batch_bytes: 0,
470 finished: 0,
471 // Seeded before a single thread is spawned, so it only ever counts
472 // down: the collector cannot mistake "not started yet" for "all gone".
473 live_workers: live,
474 });
475 let wake = Condvar::new();
476 // Workers wait here for the collector to free output bytes. A second
477 // condvar on the *same* mutex, so a worker's check-and-wait is atomic
478 // against the collector's decrement and no wakeup can be missed.
479 let space = Condvar::new();
480
481 let mut write_err: Option<std::io::Error> = None;
482
483 std::thread::scope(|scope| {
484 // Armed before a single worker is spawned, so a panic from
485 // `scope.spawn` itself — the OS refusing a thread — still runs
486 // teardown instead of leaving spawned workers parked on
487 // backpressure forever. The normal exit and the unwinding exit
488 // run identical teardown either way, since this lives until the
489 // end of the same scope regardless of where it's declared.
490 let _teardown = Teardown { state: &state, space: &space, wake: &wake, stop: &stop };
491
492 for i in 0..live {
493 before_spawn(i);
494 let (next, stop, state, wake) = (&next, &stop, &state, &wake);
495 let (space, inside) = (&space, &inside);
496 let (jobs, work, make_buffer, make_worker) =
497 (&jobs, &work, &make_buffer, &make_worker);
498 scope.spawn(move || {
499 // First act of the thread, so every later line is covered.
500 let _life = WorkerLife { state, wake };
501 // `make_worker` is the caller's closure and may panic. Catching
502 // it here keeps the panic from propagating out of the scope —
503 // which would turn a per-worker fault into a process-level one
504 // — and lets the surviving workers claim the whole job list.
505 // If none survive, `_life` above is what stops the collector
506 // waiting for them.
507 let mut ctx = match std::panic::catch_unwind(
508 std::panic::AssertUnwindSafe(&make_worker),
509 ) {
510 Ok(c) => c,
511 Err(_) => return,
512 };
513 loop {
514 // Producer-side backpressure. Waiting *before* claiming a
515 // job means a blocked worker holds nothing at all — no
516 // buffer, no slot, no job — which is what makes the wait
517 // deadlock-free. See the module docs.
518 {
519 let mut st = state.lock().unwrap();
520 while !stop.load(Ordering::Acquire)
521 && st.held + st.batch_bytes > spill_bytes
522 {
523 st = space.wait(st).unwrap();
524 }
525 }
526
527 // After backpressure, so a worker waiting for output to
528 // drain does not also hold a place someone else could use.
529 // Held until the job is deposited, at the end of this pass.
530 let _permit = match gate {
531 Some(g) => match admit(g, inside, stop, next, n) {
532 Some(p) => Some(p),
533 None => return,
534 },
535 None => None,
536 };
537
538 // Checked before claiming work, so a write failure or a
539 // closed pipe stops the sweep at the next archive boundary
540 // instead of at the end of the month.
541 if stop.load(Ordering::Acquire) {
542 return;
543 }
544 let p = next.fetch_add(1, Ordering::Relaxed);
545 if p >= n {
546 return;
547 }
548
549 let mut guard = Completion { state, wake, armed: true };
550
551 let mut buf = {
552 let mut st = state.lock().unwrap();
553 st.pool.pop()
554 }
555 .unwrap_or_else(make_buffer);
556
557 let outcome = run_job(&jobs[p], &mut ctx, &mut buf, work);
558 let bytes = buf.as_slice().len();
559 let done = Done { index: jobs[p].index, buffer: buf, outcome };
560
561 {
562 let mut st = state.lock().unwrap();
563 st.slots[p] = Some(done);
564 st.ready.push_back(p);
565 st.held += bytes;
566 st.finished += 1;
567 guard.armed = false;
568 }
569 wake.notify_all();
570 }
571 });
572 }
573
574 // The collector. Runs on the calling thread so `emit` needs to be
575 // neither `Send` nor `Sync`, and so ordering decisions live in exactly
576 // one place. `head` walks `order`, not raw positions.
577
578 let mut head = 0usize;
579 let mut emitted = vec![false; n];
580 let mut emitted_count = 0usize;
581 let mut batch: Vec<Done> = Vec::new();
582
583 'collect: while emitted_count < n {
584 let mut spill_notices = 0usize;
585 {
586 let mut st = state.lock().unwrap();
587 loop {
588 if sorted {
589 // Drain the in-order prefix: everything from `head`
590 // onwards, in index order, that has landed.
591 while head < n {
592 if emitted[head] {
593 head += 1;
594 continue;
595 }
596 match st.slots[order[head]].take() {
597 Some(d) => {
598 let len = d.buffer.as_slice().len();
599 st.held -= len;
600 st.batch_bytes += len;
601 emitted[head] = true;
602 head += 1;
603 batch.push(d);
604 }
605 None => break,
606 }
607 }
608 // Checked every pass, not only when blocked, so a fast
609 // producer cannot outrun the cap while we are emitting.
610 if st.held > spill_bytes {
611 // Counted here, printed after the lock is dropped.
612 // Writing to stderr under the state mutex stalls
613 // every worker behind a slow terminal, and a failed
614 // stderr write panics *inside* the lock: poisoned
615 // under unwinding, an abort under `panic = "abort"`.
616 // The lock is for moves in and out of `slots` and
617 // `pool`, and nothing else.
618 spill_notices += 1;
619 flush_parked(&mut st, &order, &mut emitted, head, &mut batch);
620 }
621 } else {
622 while let Some(p) = st.ready.pop_front() {
623 if let Some(d) = st.slots[p].take() {
624 let len = d.buffer.as_slice().len();
625 st.held -= len;
626 st.batch_bytes += len;
627 // `order` is the identity here, so position and
628 // order-index coincide. Marked in this mode too
629 // so `emitted` means "taken" on both paths —
630 // the stall sweep below needs to know which
631 // positions are still outstanding.
632 emitted[p] = true;
633 batch.push(d);
634 }
635 }
636 }
637
638 if !batch.is_empty() {
639 break;
640 }
641 // Not a spurious-wakeup hazard: `finished == n` implies
642 // every unemitted slot is `Some`, so the drain above would
643 // have filled `batch`. Reaching here means work is still
644 // outstanding and some worker will notify — the `Completion`
645 // guard makes that true even for a worker that unwinds.
646 debug_assert!(st.finished < n, "collector stalled with all work finished");
647 if st.finished >= n {
648 break;
649 }
650 // Jobs outstanding and not one worker left to claim them.
651 // Only reachable when a worker died on a path neither
652 // `run_job`'s `catch_unwind` nor `Completion` covers — a
653 // poisoned lock, say. Waiting here is the hang; instead,
654 // take whatever landed and report the rest as unread, which
655 // `main` turns into a 2. Never a silent short read.
656 if st.live_workers == 0 {
657 for (h, flag) in emitted.iter_mut().enumerate() {
658 if *flag {
659 continue;
660 }
661 *flag = true;
662 let p = order[h];
663 match st.slots[p].take() {
664 Some(d) => {
665 let len = d.buffer.as_slice().len();
666 st.held -= len;
667 st.batch_bytes += len;
668 batch.push(d);
669 }
670 // `make_buffer` under the lock is safe *here*
671 // and nowhere else: there are no workers left
672 // to stall. The buffer stays empty, so it costs
673 // `recycle` nothing.
674 None => batch.push(Done {
675 index: jobs[p].index,
676 buffer: make_buffer(),
677 outcome: unclaimed(&jobs[p]),
678 }),
679 }
680 }
681 break;
682 }
683 st = wake.wait(st).unwrap();
684 }
685 }
686
687 for _ in 0..spill_notices {
688 eprintln!(
689 "trg: held output exceeded {spill_bytes} bytes; flushing out of order"
690 );
691 }
692
693 // Unreachable per the invariant above; a `break` rather than a spin
694 // so a bug degrades into short output instead of a hung process.
695 if batch.is_empty() {
696 break;
697 }
698
699 // Emit outside the lock: stdout may block, workers must not.
700 for d in batch.drain(..) {
701 emitted_count += 1;
702 let r = emit(&d);
703 recycle(&state, &space, d.buffer, spill_bytes);
704 if let Err(e) = r {
705 stop.store(true, Ordering::Release);
706 write_err = Some(e);
707 break 'collect;
708 }
709 }
710 }
711
712 // Anything still parked is dropped on the way out of the scope; only
713 // reachable on the early-stop path, where it is output nobody wants.
714 // Waking the blocked producers is `Teardown`'s job, on this path and on
715 // the unwinding one alike.
716 });
717
718 match write_err {
719 Some(e) => Err(e),
720 None => Ok(()),
721 }
722}
723
724/// The outcome for a job no worker ever claimed. An `errors` entry is the only
725/// partial-run signal `main` reads, so this is what makes the stalled run a 2
726/// rather than a confident "no matches, everything read".
727fn unclaimed(job: &Job) -> crate::archive::Outcome {
728 let mut o = crate::archive::Outcome::default();
729 o.errors.push(format!(
730 "{}: every worker exited before this archive was searched; not read",
731 job.path.display()
732 ));
733 o
734}
735
736/// Return a buffer to the pool so its allocation is reused, unless it grew past
737/// the spill cap — one pathological archive should not reserve its whole match
738/// volume for the rest of the run. `termcolor::Buffer` exposes no `capacity`,
739/// so the length it reached stands in for it.
740///
741/// This is the only place output bytes are actually freed, so it is also where
742/// backpressure is released. The pool itself needs no size cap: the number of
743/// buffers in existence is now bounded by the backpressure invariant, and
744/// capping it would only throw away the allocations the pool exists to keep.
745fn recycle(
746 state: &Mutex<Shared>,
747 space: &Condvar,
748 mut buffer: termcolor::Buffer,
749 spill_bytes: usize,
750) {
751 let len = buffer.as_slice().len();
752 let oversized = len > spill_bytes;
753 buffer.clear();
754 if let Ok(mut st) = state.lock() {
755 st.batch_bytes -= len;
756 if !oversized {
757 st.pool.push(buffer);
758 }
759 }
760 space.notify_all();
761}
762
763/// Emit everything that has landed but is out of turn, in index order. The
764/// escape valve behind both the byte cap and buffer starvation.
765fn flush_parked(
766 st: &mut Shared,
767 order: &[usize],
768 emitted: &mut [bool],
769 head: usize,
770 batch: &mut Vec<Done>,
771) {
772 for (h, flag) in emitted.iter_mut().enumerate().skip(head) {
773 if let Some(d) = st.slots[order[h]].take() {
774 // Note this does not relieve backpressure: the bytes move from
775 // `held` to `batch_bytes`, and only `recycle` frees them.
776 let len = d.buffer.as_slice().len();
777 st.held -= len;
778 st.batch_bytes += len;
779 *flag = true;
780 batch.push(d);
781 }
782 }
783}
784
785#[cfg(test)]
786mod tests {
787 use super::*;
788 use std::time::Duration;
789
790 /// Deadlock #4: the teardown guard was once armed *after* the spawn loop.
791 /// If `scope.spawn` panics for worker 1 — the OS refusing a thread — while
792 /// worker 0 is parked on backpressure, nothing would wake worker 0, and
793 /// `thread::scope` would wait forever to join it. The hook stands in for
794 /// the refused spawn, since the real one cannot be provoked on demand.
795 ///
796 /// The run happens on its own thread so a regression fails here after
797 /// `PATIENCE` instead of hanging the whole suite.
798 #[test]
799 fn a_spawn_failure_does_not_strand_a_parked_worker() {
800 const PATIENCE: Duration = Duration::from_secs(10);
801 let (tx, rx) = std::sync::mpsc::channel();
802 std::thread::spawn(move || {
803 let jobs = (0..4)
804 .map(|i| Job { index: i, path: std::path::PathBuf::from(format!("j{i}")) })
805 .collect();
806 let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
807 run_with_spawn_hook(
808 jobs,
809 // Worker 0's first job overfills the cap, so it parks on
810 // backpressure: no collector is running to drain it.
811 Config { workers: 4, sorted: true, spill_bytes: 8 },
812 None,
813 termcolor::Buffer::no_color,
814 || (),
815 |_job, _c: &mut (), b| {
816 use std::io::Write;
817 write!(b, "{}", "y".repeat(100)).unwrap();
818 crate::archive::Outcome::default()
819 },
820 |_done: &Done| Ok(()),
821 |i| {
822 if i == 1 {
823 panic!("injected: the OS refused thread {i}");
824 }
825 },
826 )
827 }));
828 let _ = tx.send(r.is_err());
829 });
830 match rx.recv_timeout(PATIENCE) {
831 Ok(panicked) => assert!(panicked, "the injected spawn failure must propagate"),
832 Err(_) => panic!("run hung after a spawn failure: a parked worker was never released"),
833 }
834 }
835}