Skip to main content

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}