Skip to main content

turbo_tasks/
scope_unbounded.rs

1//! Unbounded scoped parallelism: enables running jobs that can discover and enqueue more work
2
3use std::{
4    any::Any,
5    ops::ControlFlow,
6    panic::{self, AssertUnwindSafe, catch_unwind},
7    sync::{
8        Arc, OnceLock, SyncView,
9        atomic::{AtomicBool, AtomicUsize, Ordering},
10        mpmc,
11    },
12    time::Duration,
13};
14
15use fixedbitset::FixedBitSet;
16use parking_lot::{Condvar, Mutex, RwLock};
17use tokio::{runtime::Handle, task::AbortHandle};
18use tracing::{Span, info_span};
19
20use crate::{TurboTasksApi, manager::try_turbo_tasks, turbo_tasks_scope};
21
22/// How long a scope worker waits on an empty queue before exiting.
23///
24/// Optimizes respawning which triggers overhead managing WorkerSlots and the accumulator variables
25/// pausing for a short time is worthwhile to avoid that.
26const WORKER_IDLE_TIMEOUT: Duration = Duration::from_micros(100);
27
28/// Runs `run` over `initial` and everything it transitively spawns, returning once every item has
29/// been processed. No results are collected; jobs communicate through state captured in `run`. Use
30/// [`scope_unbounded_with`] to accumulate a value instead.
31//////
32/// Items must be `'static` (they sit in a queue drained by other threads); the `run` closure may
33/// borrow `'env` data.
34///
35/// # Aborting
36///
37/// Both [`ControlFlow::Break`] and a panic abandon all queued-but-unstarted items, so the scope
38/// returns as soon as the currently-running jobs finish. Jobs already in flight on other threads
39/// are **not** interrupted in either case.
40pub fn scope_unbounded<'env, T, F>(initial: impl IntoIterator<Item = T>, run: F)
41where
42    T: Send + 'static,
43    F: Fn(&Scope<'_, T, ()>, T) -> ControlFlow<()> + Send + Sync + 'env,
44{
45    scope_unbounded_with(
46        initial,
47        || (),
48        |spawner, item, ()| run(spawner, item),
49        |(), ()| (),
50    )
51}
52
53/// [`scope_unbounded`], plus a per-drainer accumulator folded into a single return value.
54///
55/// Each drainer builds its own accumulator with `init`, `run` mutates it in place while processing
56/// items, and the accumulators are combined pairwise with `merge` as drainers finish. `merge` must
57/// be associative and commutative — drainers finish in a nondeterministic order, so the grouping
58/// and ordering of the folds are not specified.
59///
60/// This exists so `run` can efficiently accumulate with minimal locking overhead managed by the
61/// scope.
62///
63/// `init` is called potentially many times for each thread context.
64///
65/// Returns `init()` when no item is ever processed (e.g. an empty `initial`).
66///
67/// # Panics and aborts
68///
69/// Both [`ControlFlow::Break`] and a panic abort the scope, abandoning every queued-but-unstarted
70/// item (see [`scope_unbounded`]). They differ in what comes back:
71///
72/// - On `Break`, results accumulated before the abort are returned as usual; the abandoned items
73///   simply never contributed.
74/// - On a panic, the panic is re-raised after the join and **all accumulated results are
75///   discarded** — the return value is only produced on the normal path.
76pub fn scope_unbounded_with<'env, T, R, F, Init, Merge>(
77    initial: impl IntoIterator<Item = T>,
78    init: Init,
79    run: F,
80    merge: Merge,
81) -> R
82where
83    T: Send + 'static,
84    R: Send + 'env,
85    F: Fn(&Scope<'_, T, R>, T, &mut R) -> ControlFlow<()> + Send + Sync + 'env,
86    Init: Fn() -> R + Send + Sync + 'env,
87    Merge: Fn(R, R) -> R + Send + Sync + 'env,
88{
89    let handle = Handle::current();
90    // One worker per runtime thread beyond the calling thread
91    let max_workers = handle.metrics().num_workers().saturating_sub(1);
92    let span = Span::current();
93
94    // `ScopeInner` is parameterized over the borrow lifetime, so these go in as ordinary
95    // references. The one erasure to `'static` is at the tokio hand-off in
96    // `spawn_worker_if_needed`.
97    let init_ref: &(dyn Fn() -> R + Send + Sync + '_) = &init;
98    let merge_ref: &(dyn Fn(R, R) -> R + Send + Sync + '_) = &merge;
99
100    let (sender, receiver) = mpmc::channel();
101    let mut inner = ScopeInner {
102        remaining_tasks: AtomicUsize::new(0),
103        panic: OnceLock::new(),
104        work_queue: receiver,
105        work_queue_sender: RwLock::new(Some(sender)),
106        aborted: AtomicBool::new(false),
107        available_slots: AtomicUsize::new(max_workers),
108        workers: Mutex::new(WorkerSlots::new(max_workers)),
109        handle: handle.clone(),
110        span: span.clone(),
111        workers_idle: Condvar::new(),
112        turbo_tasks: try_turbo_tasks(),
113        run: &run,
114        results: Mutex::new(None),
115        init: init_ref,
116        merge: merge_ref,
117    };
118
119    // Arm the join guard before anything can spawn, so a worker can never outlive the join.
120    let joiner = Joiner { inner: &inner };
121
122    // Increment remaining tasks to ensure the scope cannot exit before all tasks are enqueued
123    inner.remaining_tasks.fetch_add(1, Ordering::Relaxed);
124    for item in initial {
125        enqueue(&inner, item);
126    }
127
128    // Drain and join before checking for a panic. Every drainer has merged its accumulator by the
129    // time this returns.
130    drop(joiner);
131
132    if let Some(err) = inner.panic.take() {
133        panic::resume_unwind(err.into_inner());
134    }
135
136    inner.results.lock().take().unwrap_or_else(init)
137}
138
139/// Handle passed to the `run` closure of [`scope_unbounded`], used to enqueue additional items into
140/// the same scope.
141pub struct Scope<'scope, T: Send + 'static, R = ()> {
142    inner: &'scope ScopeInner<'scope, T, R>,
143}
144
145impl<T: Send + 'static, R: Send> Scope<'_, T, R> {
146    /// Enqueue another item to be processed by `run`. Callable any number of times from inside
147    /// `run`, on any drainer thread.
148    ///
149    /// Silently drops `item` once the scope has aborted.
150    pub fn spawn(&self, item: T) {
151        enqueue(self.inner, item);
152    }
153}
154
155/// A reference to the shared per-item closure for a [`scope_unbounded`] run. `'run` is the lifetime
156/// of the borrows it captures (`'env` at the call site, erased to `'static` when handed to tokio).
157/// `R` is the per-drainer accumulator threaded through by [`scope_unbounded_with`].
158type RunFn<'run, T, R> =
159    &'run (dyn Fn(&Scope<'_, T, R>, T, &mut R) -> ControlFlow<()> + Send + Sync + 'run);
160
161/// Shared state for a [`scope_unbounded`] run, living on the caller's stack.
162///
163/// `'run` is the lifetime of the borrows held by the `run`/`init`/`merge` closures (`'env` at the
164/// call site). It stays a real lifetime here rather than being pinned to `'static` so the fields
165/// don't each force `R: 'static`; the single erasure to `'static` happens at the [`Drainable`]
166/// hand-off to tokio.
167struct ScopeInner<'run, T: Send + 'static, R> {
168    /// Items enqueued but not yet finished. The scope is done exactly when this reaches zero; see
169    /// [`enqueue`] for the increment-before-push ordering that makes zero reliable.
170    remaining_tasks: AtomicUsize,
171    /// First panic raised while processing an item; propagated to the caller after the join.
172    panic: OnceLock<SyncView<Box<dyn Any + Send + 'static>>>,
173    /// Receiving end of the work queue, shared by every drainer.
174    work_queue: mpmc::Receiver<T>,
175    /// Sending end of the queue, modeled so we can `take` and thus close the queue
176    work_queue_sender: RwLock<Option<mpmc::Sender<T>>>,
177    aborted: AtomicBool,
178    /// Spawn budget, mutated under the workers slot, atomic so it can be read outside of it.
179    available_slots: AtomicUsize,
180    workers: Mutex<WorkerSlots>,
181    /// Triggered when the last worker in `workers` is cleared.
182    workers_idle: Condvar,
183    handle: Handle,
184    span: Span,
185    turbo_tasks: Option<Arc<dyn TurboTasksApi>>,
186    /// The per-item closure.
187    run: RunFn<'run, T, R>,
188    init: &'run (dyn Fn() -> R + Send + Sync + 'run),
189    merge: &'run (dyn Fn(R, R) -> R + Send + Sync + 'run),
190    // Accumulated results, workers aggregate into this using init/merge as workers exit their
191    // scope
192    results: Mutex<Option<R>>,
193}
194
195impl<T: Send + 'static, R> ScopeInner<'_, T, R> {
196    /// Closes the work queue by dropping the only sender. Every blocked `recv` returns `Err` once
197    /// this runs and the buffer is drained, which is how drainers learn the scope is finished.
198    /// Idempotent.
199    fn close(&self) {
200        drop(self.work_queue_sender.write().take());
201    }
202
203    /// Abandons all queued-but-unstarted work. Idempotent.
204    fn abort(&self) {
205        self.aborted.store(true, Ordering::Release);
206        self.close();
207    }
208
209    fn on_item_finished(&self) {
210        if self.remaining_tasks.fetch_sub(1, Ordering::Release) == 1 {
211            self.close();
212        }
213    }
214
215    /// Keeps the first panic seen; later ones are dropped.
216    fn record_panic(&self, err: Box<dyn Any + Send + 'static>) {
217        self.abort();
218        let _ = self.panic.set(SyncView::new(err));
219    }
220
221    /// Drain loop, run by both the workers and the calling thread until the scope terminates.
222    ///
223    /// - A scope worker (`is_worker`) exits after [`WORKER_IDLE_TIMEOUT`] on an empty queue
224    /// - The calling thread blocks until the queue closes, which requires every item to be
225    ///   finished.
226    fn drain(&self, is_worker: bool) {
227        if is_worker && let Some(turbo_tasks) = &self.turbo_tasks {
228            turbo_tasks_scope(turbo_tasks.clone(), || self.drain_loop(is_worker))
229        } else {
230            self.drain_loop(is_worker)
231        }
232    }
233
234    fn drain_loop(&self, is_worker: bool) {
235        let mut acc: Option<R> = None;
236        while let Some(item) = if is_worker {
237            self.work_queue.recv_timeout(WORKER_IDLE_TIMEOUT).ok()
238        } else {
239            self.work_queue.recv().ok()
240        } {
241            // Post-abort: discard without running, so the wind-down can't re-grow the queue.
242            if self.aborted.load(Ordering::Acquire) {
243                self.on_item_finished();
244                continue;
245            }
246            let spawner = Scope { inner: self };
247            let result = catch_unwind(AssertUnwindSafe(|| {
248                // Lazily init the thread local accumulator only when we are going to execute an
249                // item
250                let acc = acc.get_or_insert_with(self.init);
251                (self.run)(&spawner, item, acc)
252            }));
253
254            match result {
255                Ok(ControlFlow::Continue(())) => {}
256                Ok(ControlFlow::Break(())) => {
257                    self.abort();
258                }
259                // A panic aborts too; see `scope_unbounded`.
260                Err(panic) => {
261                    self.record_panic(panic);
262                }
263            };
264            self.on_item_finished();
265        }
266
267        // Fold this drainer's accumulator into the shared results slot
268        if let Some(acc) = acc {
269            let merged = catch_unwind(AssertUnwindSafe(|| {
270                let mut results = self.results.lock();
271                *results = Some(match results.take() {
272                    Some(existing) => (self.merge)(existing, acc),
273                    None => acc,
274                });
275            }));
276            if let Err(panic) = merged {
277                self.record_panic(panic);
278            }
279        }
280    }
281}
282
283/// Account for and enqueue one item. The increment must happen before the push: pushing first would
284/// let another drainer pop and finish the item before it is counted, so `remaining_tasks` could hit
285/// zero with work still live.
286fn enqueue<T: Send + 'static, R: Send>(inner: &ScopeInner<'_, T, R>, item: T) {
287    if inner.aborted.load(Ordering::Acquire) {
288        return;
289    }
290    let num_tasks = inner.remaining_tasks.fetch_add(1, Ordering::Relaxed) + 1;
291    let sent = {
292        let sender = inner.work_queue_sender.read();
293        match sender.as_ref() {
294            Some(sender) => sender.send(item).is_ok(),
295            // Closed: the scope is winding down (aborted, or already finished).
296            None => false,
297        }
298    };
299    if !sent {
300        inner.on_item_finished(); // since the item won't execute decrement now
301        return;
302    }
303    spawn_worker_if_needed(inner, num_tasks);
304}
305
306/// Re-arm one worker if the scope is running below its budget.
307fn spawn_worker_if_needed<T: Send + 'static, R: Send>(
308    inner: &ScopeInner<'_, T, R>,
309    num_enqueued_tasks: usize,
310) {
311    if num_enqueued_tasks <= 1
312        || inner.available_slots.load(Ordering::Relaxed) == 0
313        || inner.aborted.load(Ordering::Acquire)
314    {
315        return;
316    }
317
318    // SAFETY: `Joiner::drop` waits for every worker slot to be released before returning, and the
319    // slot for this worker is claimed below under the table lock *before* the spawn, so no erased
320    // reference can outlive `'env` or the `inner` stack slot.
321    let erased: &(dyn Drainable + Send + Sync + '_) = inner;
322    let erased: &'static (dyn Drainable + Send + Sync + 'static) = unsafe {
323        std::mem::transmute::<
324            &(dyn Drainable + Send + Sync + '_),
325            &'static (dyn Drainable + Send + Sync + 'static),
326        >(erased)
327    };
328
329    let mut slots = inner.workers.lock();
330    let Some(slot) = slots.free_slot() else {
331        return;
332    };
333    let span = inner.span.clone();
334    // capture before the spawn and move into it
335    let guard = erased.claim_worker_slot(slot);
336    let handle = inner
337        .handle
338        .spawn(async move {
339            let _span = span.entered();
340            let _guard = guard;
341            erased.drain(true);
342        })
343        .abort_handle();
344    slots.occupy(slot, handle, &inner.available_slots);
345}
346
347/// The drain loop with the accumulator type erased.
348///
349/// Worker tasks are spawned onto tokio and so must be `'static`, but the accumulator `R` borrows
350/// `'env`. A worker only ever needs to *run* the loop — it never names an `R` — so it holds the
351/// scope through this trait instead of the concrete [`ScopeInner`], keeping `R` out of the spawned
352/// future's type entirely.
353trait Drainable {
354    fn drain(&self, is_worker: bool);
355    /// Take ownership of `slot`'s release, which [`spawn_worker_if_needed`] has already claimed.
356    /// On the trait so a spawned worker can build its guard without naming `R`.
357    fn claim_worker_slot(&self, slot: usize) -> WorkerGuard<'_>;
358}
359
360impl<T: Send + 'static, R> Drainable for ScopeInner<'_, T, R> {
361    fn drain(&self, is_worker: bool) {
362        ScopeInner::drain(self, is_worker)
363    }
364
365    fn claim_worker_slot(&self, slot: usize) -> WorkerGuard<'_> {
366        WorkerGuard {
367            slots: &self.workers,
368            workers_idle: &self.workers_idle,
369            available_slots: &self.available_slots,
370            slot,
371        }
372    }
373}
374
375/// Fixed table of worker slots, indexed by slot number, with a bitset of which are occupied.
376///
377/// This is both the spawn budget and the join set, because a slot is occupied from *before* its
378/// task exists until *after* that task's last access to the caller's frame:
379struct WorkerSlots {
380    handles: Vec<Option<AbortHandle>>,
381    /// Bit `i` set means slot `i` is occupied — claimed by [`Self::occupy`] and not yet released
382    /// by a [`WorkerGuard`]. Tracks *occupancy*, not handle presence.
383    occupied: FixedBitSet,
384}
385
386impl WorkerSlots {
387    fn new(max_workers: usize) -> Self {
388        Self {
389            handles: vec![None; max_workers],
390            occupied: FixedBitSet::with_capacity(max_workers),
391        }
392    }
393
394    fn free_slot(&self) -> Option<usize> {
395        self.occupied.zeroes().next()
396    }
397
398    /// Record a newly spawned worker in `slot`, which must have come from [`Self::free_slot`].
399    fn occupy(&mut self, slot: usize, handle: AbortHandle, available_slots: &AtomicUsize) {
400        debug_assert!(
401            !self.occupied.contains(slot),
402            "slot {slot} already occupied"
403        );
404        available_slots.fetch_sub(1, Ordering::Relaxed);
405        self.occupied.insert(slot);
406        let previous = self.handles[slot].replace(handle);
407        debug_assert!(previous.is_none(), "slot {slot} held a live handle");
408    }
409
410    /// Dropping the handle here keeps the finished task's allocation from outliving the worker.
411    fn release(&mut self, slot: usize, available_slots: &AtomicUsize) {
412        available_slots.fetch_add(1, Ordering::Relaxed);
413        self.occupied.remove(slot);
414        self.handles[slot] = None;
415    }
416
417    /// Whether every slot is free, i.e. no task can still touch the caller's frame. The predicate
418    /// [`Joiner::drop`] waits on.
419    fn is_idle(&self) -> bool {
420        self.occupied.is_clear()
421    }
422}
423
424/// Releases a worker's slot and wakes [`Joiner::drop`] when it is the last one.
425struct WorkerGuard<'a> {
426    slots: &'a Mutex<WorkerSlots>,
427    workers_idle: &'a Condvar,
428    available_slots: &'a AtomicUsize,
429    slot: usize,
430}
431
432impl Drop for WorkerGuard<'_> {
433    fn drop(&mut self) {
434        let mut slots_guard = self.slots.lock();
435        slots_guard.release(self.slot, self.available_slots);
436        if slots_guard.is_idle() {
437            // Still holding the lock: the predicate the joiner waits on is read under this same
438            // lock, so it cannot go true between its check and its `wait`.
439            self.workers_idle.notify_all();
440        }
441    }
442}
443
444/// Drains the queue and joins the workers, on the return path and on an unwind alike.
445struct Joiner<'a, 'run, T: Send + 'static, R> {
446    inner: &'a ScopeInner<'run, T, R>,
447}
448
449impl<T: Send + 'static, R> Drop for Joiner<'_, '_, T, R> {
450    fn drop(&mut self) {
451        // Discharge the placeholder item that covered the seeding loop.
452        self.inner.on_item_finished();
453        // Returns only once the queue is closed, so no new work can arrive after this.
454        self.inner.drain(false);
455        // The queue is now closed so no workers can spawn. Capture all the handles.
456        // There should be no contention on this slot
457        let _span = info_span!("blocking: waiting for scope to end").entered();
458        let handles: Vec<_> = {
459            let mut slots = self.inner.workers.lock();
460            slots.handles.iter_mut().filter_map(Option::take).collect()
461        };
462
463        // Abort all workers, that way workers we have spawned but have never run get dropped and
464        // release their slots.  Otherwise a contended runtime could delay shutdown of the scope.
465        for handle in handles {
466            handle.abort();
467        }
468
469        // Wait for all workers to exit and drop their guards.  This should be fast, after the
470        // aborts and the queue is dropped each worker just needs to merge their accumulated
471        // stats.  So we wait for that.
472        let mut slots = self.inner.workers.lock();
473        while !slots.is_idle() {
474            self.inner.workers_idle.wait(&mut slots);
475        }
476    }
477}
478#[cfg(test)]
479mod tests {
480    use std::{
481        sync::{Arc, atomic::AtomicUsize},
482        thread,
483        time::Duration,
484    };
485
486    use super::*;
487
488    /// Runs `body` on a runtime with the I/O driver disabled, so the test also works under Miri.
489    ///
490    /// `#[tokio::test]` hardcodes `Builder::enable_all()`, whose I/O driver calls
491    /// `kqueue()`/`epoll_create1()` — Miri has no shim for either, so such a test aborts at runtime
492    /// construction before reaching any code under test. Nothing here needs I/O.
493    ///
494    /// `body` runs on a blocking thread, as production callers are expected to: `scope_unbounded`
495    /// blocks, and its `block_in_place` requires a multi-thread runtime.
496    fn with_runtime<F, T>(worker_threads: usize, body: F) -> T
497    where
498        F: FnOnce() -> T + Send + 'static,
499        T: Send + 'static,
500    {
501        let runtime = tokio::runtime::Builder::new_multi_thread()
502            .worker_threads(worker_threads)
503            .enable_time()
504            .build()
505            .unwrap();
506        runtime.block_on(async { tokio::task::spawn_blocking(body).await.unwrap() })
507    }
508
509    /// A single `run` call enqueues a large batch of leaves; every one must be processed.
510    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
511    async fn test_unbounded_wide_burst_of_leaves() {
512        const CHILDREN: usize = 1000;
513        let processed = Arc::new(AtomicUsize::new(0));
514        let processed_clone = processed.clone();
515        tokio::task::spawn_blocking(move || {
516            scope_unbounded(std::iter::once(0usize), move |spawner, item| {
517                processed_clone.fetch_add(1, Ordering::SeqCst);
518                if item == 0 {
519                    // The root fans out to CHILDREN leaves.
520                    for i in 0..CHILDREN {
521                        spawner.spawn(1 + i);
522                    }
523                }
524                ControlFlow::Continue(())
525            });
526        })
527        .await
528        .unwrap();
529        // 1 root + CHILDREN leaves.
530        assert_eq!(processed.load(Ordering::SeqCst), 1 + CHILDREN);
531    }
532
533    /// A slow seeding iterator must not let the scope finish early.
534    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
535    async fn test_unbounded_slow_seeding_iterator_completes() {
536        const SEEDS: usize = 16;
537        let processed = Arc::new(AtomicUsize::new(0));
538        let processed_clone = processed.clone();
539        tokio::task::spawn_blocking(move || {
540            // Each `next()` blocks briefly, so the queue drains to empty before the next seed
541            // arrives.
542            let slow_seeds = std::iter::from_fn({
543                let mut next = 0;
544                move || {
545                    if next == SEEDS {
546                        return None;
547                    }
548                    thread::sleep(Duration::from_millis(2));
549                    next += 1;
550                    Some(next - 1)
551                }
552            });
553            scope_unbounded(slow_seeds, move |_spawner, _item| {
554                processed_clone.fetch_add(1, Ordering::SeqCst);
555                ControlFlow::Continue(())
556            });
557        })
558        .await
559        .unwrap();
560        assert_eq!(
561            processed.load(Ordering::SeqCst),
562            SEEDS,
563            "seeds produced after the queue briefly drained must still be processed"
564        );
565    }
566
567    /// Aborting in the middle of a deep, still-growing cascade must terminate rather than hang:
568    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
569    async fn test_unbounded_abort_during_cascade() {
570        // Each item spawns two children until the id exceeds the bound, so the queue is still
571        // growing when the abort lands.
572        const MAX_ID: usize = 1 << 14;
573        let processed = Arc::new(AtomicUsize::new(0));
574        let processed_clone = processed.clone();
575        tokio::task::spawn_blocking(move || {
576            scope_unbounded(std::iter::once(1usize), move |spawner, id| {
577                let n = processed_clone.fetch_add(1, Ordering::SeqCst);
578                if n == 100 {
579                    return ControlFlow::Break(());
580                }
581                let (left, right) = (id * 2, id * 2 + 1);
582                if left <= MAX_ID {
583                    spawner.spawn(left);
584                }
585                if right <= MAX_ID {
586                    spawner.spawn(right);
587                }
588                ControlFlow::Continue(())
589            });
590        })
591        .await
592        .unwrap();
593        let count = processed.load(Ordering::SeqCst);
594        assert!(
595            count < MAX_ID,
596            "abort must cut the cascade short, but {count} items ran"
597        );
598    }
599
600    /// `spawn` issued *after* the abort has latched must be dropped, not enqueued — the case a job
601    /// finishing concurrently with another job's abort hits. A `spawn` that counted an item into
602    /// `remaining_tasks` without queueing it would never reach zero, so this hangs rather than
603    /// fails.
604    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
605    async fn test_unbounded_spawn_after_abort_is_dropped() {
606        let processed = Arc::new(AtomicUsize::new(0));
607        let processed_clone = processed.clone();
608        const SEEDS: usize = 64;
609        tokio::task::spawn_blocking(move || {
610            scope_unbounded(0..SEEDS, move |spawner, item| {
611                processed_clone.fetch_add(1, Ordering::SeqCst);
612                // Spawning *before* the `Break` is the point: these spawns race the abort latch and
613                // must be dropped rather than counted-but-unqueued.
614                for i in 0..1000 {
615                    spawner.spawn(SEEDS + item * 1000 + i);
616                }
617                ControlFlow::Break(())
618            });
619        })
620        .await
621        .unwrap();
622        // Only seeds may run: every spawned id is >= SEEDS, so processing even one would push the
623        // count past the seed total.
624        let count = processed.load(Ordering::SeqCst);
625        assert!(
626            count <= SEEDS,
627            "post-abort spawns must be dropped, but {count} items ran"
628        );
629    }
630
631    /// Abort on a `current_thread` runtime, where the calling thread is the only drainer.
632    #[tokio::test(flavor = "current_thread")]
633    async fn test_unbounded_abort_current_thread_runtime() {
634        let processed = Arc::new(AtomicUsize::new(0));
635        let processed_clone = processed.clone();
636        tokio::task::spawn_blocking(move || {
637            scope_unbounded(0..1000usize, move |spawner, _item| {
638                processed_clone.fetch_add(1, Ordering::SeqCst);
639                spawner.spawn(9999);
640                ControlFlow::Break(())
641            });
642        })
643        .await
644        .unwrap();
645        // With a single drainer the abort lands before any other item is picked up.
646        assert_eq!(processed.load(Ordering::SeqCst), 1);
647    }
648
649    /// A panic that happens while the scope is aborting still propagates rather than being
650    /// swallowed by the wind-down: the abort's queue-clear races the panic's unwind through
651    /// `catch_unwind` -> `on_item_finished`.
652    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
653    async fn test_unbounded_abort_then_panic() {
654        let result = catch_unwind(AssertUnwindSafe(|| {
655            scope_unbounded(0..1000usize, |_spawner, item| {
656                if item == 0 {
657                    panic!("Intentional panic");
658                }
659                ControlFlow::Break(())
660            });
661            unreachable!();
662        }));
663        let err = result.expect_err("the panic must propagate even though the scope aborted");
664        assert_eq!(err.downcast_ref::<&str>(), Some(&"Intentional panic"));
665    }
666
667    /// A panic in a `run` invocation propagates after all in-flight work is joined, and aborts the
668    /// scope: the queued-but-unstarted items are abandoned rather than run.
669    ///
670    /// Use a growing cascade and panic after enough work has run to guarantee that work remains
671    /// queued. A fixed seed set can drain completely before its first item panics, making it unable
672    /// to distinguish a missed abort from valid scheduling.
673    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
674    async fn test_unbounded_panic_propagates_and_abandons_queue() {
675        const MAX_ID: usize = 1 << 14;
676        let processed = Arc::new(AtomicUsize::new(0));
677        let processed_clone = processed.clone();
678        let result = catch_unwind(AssertUnwindSafe(|| {
679            scope_unbounded(std::iter::once(1usize), move |spawner, id| {
680                let n = processed_clone.fetch_add(1, Ordering::SeqCst);
681                if n == 100 {
682                    panic!("Intentional panic");
683                }
684                let (left, right) = (id * 2, id * 2 + 1);
685                if left <= MAX_ID {
686                    spawner.spawn(left);
687                }
688                if right <= MAX_ID {
689                    spawner.spawn(right);
690                }
691                ControlFlow::Continue(())
692            });
693            unreachable!();
694        }));
695        let err = result.expect_err("the panic must propagate");
696        assert_eq!(err.downcast_ref::<&str>(), Some(&"Intentional panic"));
697        let count = processed.load(Ordering::SeqCst);
698        assert!(
699            count < MAX_ID,
700            "a panic must cut the cascade short, but {count} items ran"
701        );
702    }
703
704    // -----------------------------------------------------------------------
705    // scope_unbounded_with (fold results)
706    // -----------------------------------------------------------------------
707
708    /// The accumulator must be per-drainer, not shared: collecting into a `Vec` and merging by
709    /// concatenation must preserve every element even with several drainers running.
710    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
711    async fn test_unbounded_with_collects_all_values() {
712        const ITEMS: usize = 500;
713        let mut collected = tokio::task::spawn_blocking(|| {
714            scope_unbounded_with(
715                0..ITEMS,
716                Vec::new,
717                |_spawner, item: usize, acc: &mut Vec<usize>| {
718                    acc.push(item);
719                    ControlFlow::Continue(())
720                },
721                |mut a: Vec<usize>, b| {
722                    a.extend(b);
723                    a
724                },
725            )
726        })
727        .await
728        .unwrap();
729        collected.sort_unstable();
730        assert_eq!(collected, (0..ITEMS).collect::<Vec<_>>());
731    }
732
733    /// With no items, no drainer builds an accumulator, so the result is exactly one `init()` —
734    /// not a fold of one per drainer that happened to start.
735    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
736    async fn test_unbounded_with_empty_returns_init() {
737        let total = tokio::task::spawn_blocking(|| {
738            scope_unbounded_with(
739                std::iter::empty::<usize>(),
740                || 42usize,
741                |_spawner, _item, _acc| ControlFlow::Continue(()),
742                |a, b| a + b,
743            )
744        })
745        .await
746        .unwrap();
747        assert_eq!(total, 42, "expected exactly one init(), got {total}");
748    }
749
750    /// A scope driven **directly** on a `current_thread` runtime — not via `spawn_blocking` — must
751    /// still complete.
752    #[tokio::test(flavor = "current_thread")]
753    async fn test_unbounded_current_thread_direct_call_completes() {
754        let processed = Arc::new(AtomicUsize::new(0));
755        let processed_clone = processed.clone();
756        scope_unbounded(0..8usize, move |spawner, item| {
757            processed_clone.fetch_add(1, Ordering::SeqCst);
758            if item < 3 {
759                spawner.spawn(100 + item);
760            }
761            ControlFlow::Continue(())
762        });
763        assert_eq!(processed.load(Ordering::SeqCst), 11);
764    }
765
766    /// Aborting returns the results accumulated up to that point rather than discarding them —
767    /// only the abandoned items are missing. The run must still terminate cleanly.
768    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
769    async fn test_unbounded_with_abort_returns_partial_results() {
770        let processed = tokio::task::spawn_blocking(|| {
771            scope_unbounded_with(
772                0..1000usize,
773                || 0usize,
774                |_spawner, item, acc| {
775                    *acc += 1;
776                    if item == 0 {
777                        return ControlFlow::Break(());
778                    }
779                    ControlFlow::Continue(())
780                },
781                |a, b| a + b,
782            )
783        })
784        .await
785        .unwrap();
786        // At least the aborting item ran, and the abort must have cut the run short.
787        assert!(processed >= 1, "expected the aborting item to be counted");
788        assert!(
789            processed < 1000,
790            "abort should abandon queued items, but all {processed} ran"
791        );
792    }
793
794    /// A panic must propagate through the fold path without deadlocking the join, which drainers
795    /// reach only after their merge.
796    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
797    async fn test_unbounded_with_panic_propagates() {
798        let result = catch_unwind(AssertUnwindSafe(|| {
799            scope_unbounded_with(
800                0..100usize,
801                || 0usize,
802                |_spawner, item, acc| {
803                    if item == 50 {
804                        panic!("Intentional panic");
805                    }
806                    *acc += 1;
807                    ControlFlow::Continue(())
808                },
809                |a, b| a + b,
810            );
811            unreachable!();
812        }));
813        let err = result.expect_err("the panic must propagate out of the fold API");
814        assert_eq!(err.downcast_ref::<&str>(), Some(&"Intentional panic"));
815    }
816
817    /// The accumulator may borrow `'env` data (it is not `'static`), mirroring how `run` may.
818    #[test]
819    fn test_unbounded_with_borrowed_accumulator() {
820        let label = String::from("item");
821        let count = with_runtime(4, move || {
822            let label = &label;
823            scope_unbounded_with(
824                0..32usize,
825                Vec::new,
826                |_spawner, item: usize, acc: &mut Vec<String>| {
827                    acc.push(format!("{label}-{item}"));
828                    ControlFlow::Continue(())
829                },
830                |mut a: Vec<String>, b| {
831                    a.extend(b);
832                    a
833                },
834            )
835            .len()
836        });
837        assert_eq!(count, 32);
838    }
839
840    // -----------------------------------------------------------------------
841    // worker exit / respawn tests
842    //
843    // `init` runs once per drainer that receives an item, so counting `init` calls counts *distinct
844    // drainer lifetimes* — the only externally visible signal that a worker exited and a later one
845    // replaced it.
846    // -----------------------------------------------------------------------
847
848    /// A worker exits once the queue is empty, and a later `spawn` re-arms one: the scope must
849    /// still finish work produced after every worker has gone away.
850    ///
851    /// Serialized by construction — one item in flight at a time, with a gap long enough that any
852    /// worker has certainly timed out — so reaching the second item at all exercises the respawn
853    /// path rather than a still-live worker.
854    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
855    async fn test_unbounded_worker_respawns_after_going_idle() {
856        let inits = Arc::new(AtomicUsize::new(0));
857        let counted = inits.clone();
858        let processed = Arc::new(AtomicUsize::new(0));
859        let ran = processed.clone();
860        tokio::task::spawn_blocking(move || {
861            scope_unbounded_with(
862                std::iter::once(0usize),
863                move || {
864                    counted.fetch_add(1, Ordering::SeqCst);
865                },
866                move |spawner, item, ()| {
867                    ran.fetch_add(1, Ordering::SeqCst);
868                    if item == 0 {
869                        // Let the queue sit empty long enough that any worker has exited, then
870                        // produce work again. A fresh drainer is the only thing that can pick it
871                        // up.
872                        thread::sleep(Duration::from_millis(100));
873                        spawner.spawn(1);
874                    }
875                    ControlFlow::Continue(())
876                },
877                |(), ()| (),
878            )
879        })
880        .await
881        .unwrap();
882        assert_eq!(
883            processed.load(Ordering::SeqCst),
884            2,
885            "work spawned after the pool went idle must still run"
886        );
887        // Both items ran (the scope only returns once the queue is drained), and at least one
888        // drainer built an accumulator. The exact count depends on which drainer wins each item.
889        let count = inits.load(Ordering::SeqCst);
890        assert!(
891            (1..=2).contains(&count),
892            "expected 1-2 drainer lifetimes, got {count}"
893        );
894    }
895
896    /// A scope that never has queued work must not occupy a worker at all.
897    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
898    async fn test_unbounded_empty_spawns_no_workers() {
899        let inits = Arc::new(AtomicUsize::new(0));
900        let counted = inits.clone();
901        tokio::task::spawn_blocking(move || {
902            scope_unbounded_with(
903                std::iter::empty::<usize>(),
904                move || counted.fetch_add(1, Ordering::SeqCst),
905                |_spawner, _item, _acc| ControlFlow::Continue(()),
906                |a, b| a + b,
907            )
908        })
909        .await
910        .unwrap();
911        assert_eq!(
912            inits.load(Ordering::SeqCst),
913            1,
914            "expected only the terminal identity init(), not a per-drainer one"
915        );
916    }
917
918    /// Sustained work keeps workers alive rather than churning them:
919    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
920    async fn test_unbounded_busy_queue_does_not_churn_workers() {
921        const ITEMS: usize = 20_000;
922        let inits = Arc::new(AtomicUsize::new(0));
923        let counted = inits.clone();
924        let processed = tokio::task::spawn_blocking(move || {
925            scope_unbounded_with(
926                0..ITEMS,
927                move || {
928                    counted.fetch_add(1, Ordering::SeqCst);
929                    0usize
930                },
931                |_spawner, _item, acc| {
932                    *acc += 1;
933                    ControlFlow::Continue(())
934                },
935                |a, b| a + b,
936            )
937        })
938        .await
939        .unwrap();
940        assert_eq!(processed, ITEMS, "every item must run");
941        // 4 runtime workers => 3 scope workers + the calling thread at any instant. The bound is
942        // loose because a worker can still lose a race to the last queued item and be re-armed, and
943        // the rest of the suite competes for the same threads — but a churning implementation lands
944        // in the thousands, so anything near the worker count proves the timeout is doing its job.
945        let count = inits.load(Ordering::SeqCst);
946        assert!(
947            count <= 16,
948            "a saturated queue should not churn drainers, got {count} lifetimes for {ITEMS} items"
949        );
950    }
951
952    /// The join must not depend on tokio scheduling, even when every runtime thread is contended
953    /// and workers are still mid-drain.
954    #[test]
955    fn test_unbounded_join_under_thread_starvation() {
956        let runtime = tokio::runtime::Builder::new_multi_thread()
957            .worker_threads(2)
958            .enable_time()
959            .build()
960            .unwrap();
961        runtime.block_on(async {
962            let mut scopes = Vec::new();
963            for _ in 0..8 {
964                scopes.push(tokio::task::spawn_blocking(|| {
965                    let processed = Arc::new(AtomicUsize::new(0));
966                    let counted = processed.clone();
967                    scope_unbounded(0..200usize, move |spawner, item| {
968                        counted.fetch_add(1, Ordering::SeqCst);
969                        // Keep the queue growing so workers are still draining at join time.
970                        if item < 200 {
971                            spawner.spawn(1000 + item);
972                        }
973                        thread::yield_now();
974                        ControlFlow::Continue(())
975                    });
976                    processed.load(Ordering::SeqCst)
977                }));
978            }
979            for scope in scopes {
980                assert_eq!(scope.await.unwrap(), 400, "every scope must drain fully");
981            }
982        });
983    }
984}