Skip to main content

fast_steal/
task_queue.rs

1//! A concurrent work-stealing queue.
2//!
3//! [`TaskQueue`] holds pending and running [`Task`](crate::Task)s and lets worker
4//! threads pull fresh work or steal a sub-range from a busy peer via
5//! [`steal`](TaskQueue::steal). The number of running workers can be adjusted at
6//! runtime with [`set_threads`](TaskQueue::set_threads).
7
8#![allow(clippy::significant_drop_tightening)]
9extern crate alloc;
10use crate::{Executor, Handle, Task, WeakTask};
11use alloc::{collections::vec_deque::VecDeque, sync::Arc, vec::Vec};
12use core::ops::Range;
13// The queue's critical sections are short and synchronous (a few `VecDeque`
14// operations, never held across an await), so both locks are safe. With `std`
15// the queue uses `parking_lot` (parks waiting threads instead of spinning); on
16// `no_std` (the `spin` feature) it uses `spin`. Exactly one backend is required.
17#[cfg(feature = "std")]
18use parking_lot::Mutex;
19#[cfg(all(not(feature = "std"), feature = "spin"))]
20use spin::Mutex;
21#[cfg(not(any(feature = "std", feature = "spin")))]
22compile_error!(
23    "fast-steal needs a lock backend: enable the default `std` feature, or `spin` for `no_std`"
24);
25
26/// A concurrent work-stealing queue that manages a set of [`Task`]s.
27///
28/// Workers created by [`Executor::execute`] call [`steal`](TaskQueue::steal) to obtain
29/// new work when their current task is exhausted. The queue supports splitting,
30/// speculative execution, and dynamic thread adjustment.
31#[derive(Debug)]
32pub struct TaskQueue<H: Handle> {
33    inner: Arc<Mutex<TaskQueueInner<H>>>,
34}
35impl<H: Handle> Clone for TaskQueue<H> {
36    fn clone(&self) -> Self {
37        Self {
38            inner: self.inner.clone(),
39        }
40    }
41}
42#[derive(Debug)]
43struct TaskQueueInner<H: Handle> {
44    running: VecDeque<(WeakTask, H)>,
45    waiting: VecDeque<Task>,
46}
47impl<H: Handle> TaskQueue<H> {
48    /// Creates a queue from an iterator of `start..end` ranges, each wrapped in its
49    /// own [`Task`].
50    pub fn new(tasks: impl Iterator<Item = Range<u64>>) -> Self {
51        let waiting: VecDeque<_> = tasks.map(Task::new).collect();
52        Self {
53            inner: Arc::new(Mutex::new(TaskQueueInner {
54                running: VecDeque::with_capacity(waiting.len()),
55                waiting,
56            })),
57        }
58    }
59    /// Appends a [`Task`] to the waiting queue so a future
60    /// [`steal`](TaskQueue::steal) or [`set_threads`](TaskQueue::set_threads) can
61    /// pick it up.
62    ///
63    /// Returns `true` if at least one worker is currently registered and live
64    /// (so the task will be picked up on that worker's next [`steal`](TaskQueue::steal)), or
65    /// `false` if no live worker exists — in which case the task stays stranded
66    /// in `waiting` until a [`set_threads`](TaskQueue::set_threads) call spawns a
67    /// worker to rescue it.
68    #[must_use]
69    pub fn add(&self, task: Task) -> bool {
70        let mut guard = self.inner.lock();
71        let live = guard.running.iter().any(|w| w.0.is_alive());
72        guard.waiting.push_back(task);
73        live
74    }
75    /// Tries to refill `task` with more work for the worker identified by `id`.
76    ///
77    /// The caller must pass its own currently-held [`Task`] plus `id` (compared via
78    /// [`Handle::is_self`](crate::Handle::is_self)). The function first hands out a
79    /// pending task from the waiting queue; if none is available it steals a half
80    /// range from the busiest running task via [`Task::split_two`](crate::Task::split_two)
81    /// (when at least `min_chunk_size * 2` work remains), or, if `max_speculative > 1`
82    /// and the stolen task has few enough strong references, shares that same task
83    /// speculatively.
84    ///
85    /// Returns `true` if `task` was refilled, or `false` if the worker is not
86    /// registered or no work could be found.
87    pub fn steal(
88        &self,
89        id: &H::Id,
90        task: &mut Task,
91        min_chunk_size: u64,
92        max_speculative: usize,
93    ) -> bool {
94        let min_chunk_size = min_chunk_size.max(1);
95        let mut guard = self.inner.lock();
96        let mut worker_idx = None;
97        for (i, (_, handle)) in guard.running.iter().enumerate() {
98            if handle.is_self(id) {
99                worker_idx = Some(i);
100                break;
101            }
102        }
103        let Some(worker_idx) = worker_idx else {
104            return false;
105        };
106        let mut found = false;
107        while let Some(new_task) = guard.waiting.pop_front() {
108            // A task whose range invariant is broken (`start > end`) yields
109            // `Err` and is skipped, never handed to a worker. This keeps steal's
110            // policy toward corrupted tasks uniform with the speculative branch
111            // below, which likewise discards `split_two`'s `Err`. Whether steal
112            // should instead surface such corruption is deliberately left open.
113            if let Ok(Some(range)) = new_task.take() {
114                *task = Task::new(range);
115                found = true;
116                break;
117            }
118        }
119        if !found
120            && let Some(steal_task) = guard
121                .running
122                .iter()
123                .filter_map(|w| w.0.upgrade())
124                .filter(|w| w != task)
125                .max_by_key(Task::remain)
126        {
127            if let Ok(Some(range)) = steal_task.split_two(min_chunk_size) {
128                *task = Task::new(range);
129                found = true;
130            } else if max_speculative > 1
131                && steal_task.sharer_count() < max_speculative
132                && steal_task.remain() > 0
133            {
134                task.share_state(&steal_task);
135                found = true;
136            }
137        }
138        if found {
139            guard.running[worker_idx].0 = task.downgrade();
140        } else {
141            guard.running.remove(worker_idx);
142        }
143        found
144    }
145    /// Returns `None` when threads need to be increased but the executor is `None`
146    #[must_use]
147    #[allow(clippy::significant_drop_tightening)]
148    pub fn set_threads<E: Executor<Handle = H>>(
149        &self,
150        threads: usize,
151        min_chunk_size: u64,
152        executor: Option<&E>,
153    ) -> Option<()> {
154        let threads = threads.max(1);
155        let min_chunk_size = min_chunk_size.max(1);
156        let mut guard = self.inner.lock();
157        guard.running.retain(|t| t.0.is_alive());
158        let len = guard.running.len();
159        if len < threads {
160            let executor = executor?;
161            let need = guard.waiting.len().min(threads - len);
162            let mut temp = Vec::with_capacity(need);
163            let iter = guard.waiting.drain(..need);
164            for task in iter {
165                let weak = task.downgrade();
166                let handle = executor.execute(task, self.clone());
167                temp.push((weak, handle));
168            }
169            guard.running.extend(temp);
170            while guard.running.len() < threads
171                && let Some(steal_task) = guard
172                    .running
173                    .iter()
174                    .filter_map(|w| w.0.upgrade())
175                    .max_by_key(Task::remain)
176                && let Ok(Some(range)) = steal_task.split_two(min_chunk_size)
177            {
178                let task = Task::new(range);
179                let weak = task.downgrade();
180                let handle = executor.execute(task, self.clone());
181                guard.running.push_back((weak, handle));
182            }
183        } else if len > threads {
184            let mut temp = Vec::with_capacity(len - threads);
185            let iter = guard.running.drain(threads..);
186            for (task, mut handle) in iter {
187                if let Some(task) = task.upgrade() {
188                    temp.push(task);
189                }
190                handle.abort();
191            }
192            guard.waiting.extend(temp);
193        }
194        Some(())
195    }
196    /// Provides mutable access to the handles of all running tasks, e.g. to abort
197    /// or inspect them.
198    ///
199    /// # Liveness / deadlock contract
200    /// The closure `f` is invoked *while the queue lock is held*. It must **not**
201    /// re-enter `TaskQueue` (e.g. call [`steal`](TaskQueue::steal),
202    /// [`add`](TaskQueue::add), [`set_threads`](TaskQueue::set_threads), or
203    /// [`handles`](TaskQueue::handles) again) — doing so deadlocks. Keep `f`
204    /// short: it blocks every other queue operation until it returns.
205    pub fn handles<F, R>(&self, f: F) -> R
206    where
207        F: FnOnce(&mut dyn Iterator<Item = &mut H>) -> R,
208    {
209        #![allow(clippy::significant_drop_tightening)]
210        let mut guard = self.inner.lock();
211        let mut iter = guard.running.iter_mut().map(|w| &mut w.1);
212        f(&mut iter)
213    }
214
215    /// Aborts every running task equal to `task` that does not belong to the
216    /// worker `id`.
217    ///
218    /// The call is a no-op unless `id` identifies a currently registered worker.
219    /// An unregistered caller matches no entry in `running`, which makes the
220    /// `is_self` guard vacuous: *every* twin would be aborted, including the one
221    /// that should survive, leaving work in `waiting` with no worker to claim it.
222    /// A caller can legitimately reach this state after a
223    /// [`set_threads`](TaskQueue::set_threads) shrink deregisters it, because the
224    /// abort that follows is cooperative and the worker keeps running until it
225    /// observes the signal.
226    ///
227    /// Aborted twins are *deregistered* (removed from `running`) so a later
228    /// [`set_threads`](TaskQueue::set_threads) liveness sweep does not mistake
229    /// them for live workers. Their remaining work is **not** reclaimed into
230    /// `waiting`: the aborted twins matched `t == *task`, i.e. they alias the
231    /// caller's cursor, so the caller's own still-live task already owns and
232    /// advances that remaining range. Reclaiming would only add a redundant
233    /// `waiting` entry — the `take` handshake that `steal` applies to a shared
234    /// cursor partitions it atomically, so a reclaimed twin would *not* execute
235    /// the same bytes twice. The production caller `fast-pull` therefore invokes
236    /// this only after the shared range has finished, letting the caller's task
237    /// carry the work to completion.
238    pub fn cancel_task(&self, task: &Task, id: &H::Id) {
239        let mut guard = self.inner.lock();
240        // Abort every twin whose task matches but is not the caller's own, then
241        // *deregister* it (drop it from `running`). We rebuild `running` from the
242        // survivors because removing in place would require mutating through a
243        // shared `&` handed to `retain`'s closure.
244        if !guard.running.iter().any(|(_, h)| h.is_self(id)) {
245            return;
246        }
247        let mut kept: VecDeque<(WeakTask, H)> = VecDeque::with_capacity(guard.running.len());
248        for (weak, mut handle) in guard.running.drain(..) {
249            let is_twin = weak
250                .upgrade()
251                .is_some_and(|t| t == *task && !handle.is_self(id));
252            if is_twin {
253                handle.abort();
254            } else {
255                kept.push_back((weak, handle));
256            }
257        }
258        guard.running = kept;
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    #![allow(clippy::unwrap_used)]
265    extern crate std;
266    use crate::{Executor, Handle, Task, TaskQueue};
267    use std::{
268        collections::{HashMap, HashSet},
269        dbg, println,
270        sync::{
271            Arc, Mutex,
272            atomic::{AtomicBool, Ordering},
273        },
274        vec::Vec,
275    };
276    use tokio::{sync::mpsc, task::AbortHandle};
277
278    struct TokioExecutor {
279        tx: mpsc::UnboundedSender<(u64, u64)>,
280        speculative: usize,
281    }
282    #[derive(Clone)]
283    struct TokioHandle(AbortHandle);
284
285    impl Handle for TokioHandle {
286        type Id = ();
287        fn abort(&mut self) {
288            self.0.abort();
289        }
290        fn is_self(&self, (): &Self::Id) -> bool {
291            false
292        }
293    }
294
295    impl Executor for TokioExecutor {
296        type Handle = TokioHandle;
297        fn execute(&self, mut task: Task, task_queue: TaskQueue<Self::Handle>) -> Self::Handle {
298            println!("execute");
299            let tx = self.tx.clone();
300            let speculative = self.speculative;
301            let handle = tokio::spawn(async move {
302                loop {
303                    // Keep the worker alive briefly so the shrink-mid-run test can
304                    // observe in-flight work without paying the recursive-fib cost.
305                    std::thread::sleep(std::time::Duration::from_millis(100));
306                    while task.start() < task.end() {
307                        let i = task.start();
308                        let res = fib_fast(i);
309                        let Ok(_) = task.safe_add_start(i, 1) else {
310                            println!("task-failed: {i} = {res}");
311                            continue;
312                        };
313                        println!("task: {i} = {res}");
314                        tx.send((i, res)).unwrap();
315                    }
316                    if !task_queue.steal(&(), &mut task, 1, speculative) {
317                        break;
318                    }
319                }
320            });
321            let abort_handle = handle.abort_handle();
322            TokioHandle(abort_handle)
323        }
324    }
325
326    fn fib_fast(n: u64) -> u64 {
327        let mut a = 0;
328        let mut b = 1;
329        for _ in 0..n {
330            (a, b) = (b, a + b);
331        }
332        a
333    }
334
335    #[tokio::test(flavor = "multi_thread")]
336    async fn test_task_queue() {
337        let (tx, mut rx) = mpsc::unbounded_channel();
338        let executor = TokioExecutor { tx, speculative: 1 };
339        let pre_data = [1..20, 41..48];
340        let task_queue = TaskQueue::new(pre_data.iter().cloned());
341        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
342        drop(executor);
343        let mut data = HashMap::new();
344        while let Some((i, res)) = rx.recv().await {
345            println!("main: {i} = {res}");
346            assert!(
347                data.insert(i, res).is_none(),
348                "number {i} with value {res} was computed twice"
349            );
350        }
351        dbg!(&data);
352        for range in pre_data {
353            for i in range {
354                assert_eq!((i, data.get(&i)), (i, Some(&fib_fast(i))));
355                data.remove(&i);
356            }
357        }
358        assert_eq!(data.len(), 0);
359    }
360
361    #[tokio::test(flavor = "multi_thread")]
362    async fn test_task_queue2() {
363        let (tx, mut rx) = mpsc::unbounded_channel();
364        let executor = TokioExecutor { tx, speculative: 2 };
365        let pre_data = [1..20, 41..48];
366        let task_queue = TaskQueue::new(pre_data.iter().cloned());
367        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
368        drop(executor);
369        let mut data = HashMap::new();
370        while let Some((i, res)) = rx.recv().await {
371            println!("main: {i} = {res}");
372            assert!(
373                data.insert(i, res).is_none(),
374                "number {i} with value {res} was computed twice"
375            );
376        }
377        dbg!(&data);
378        for range in pre_data {
379            for i in range {
380                assert_eq!((i, data.get(&i)), (i, Some(&fib_fast(i))));
381                data.remove(&i);
382            }
383        }
384        assert_eq!(data.len(), 0);
385    }
386
387    /// End-to-end proof that shrinking the worker pool mid-run does not lose work.
388    ///
389    /// Eight workers are started, allowed to make progress, then the pool is cut to
390    /// two. The aborted workers' in-progress tasks are reclaimed into `waiting` and
391    /// picked up by the survivors via [`steal`], so every number must still be
392    /// computed exactly once with no gaps.
393    #[tokio::test(flavor = "multi_thread")]
394    async fn test_set_threads_decrease_keeps_all_work() {
395        let (tx, mut rx) = mpsc::unbounded_channel();
396        let released = Arc::new(std::sync::atomic::AtomicBool::new(false));
397        let executor = StealExecutor {
398            tx,
399            speculative: 1,
400            next_id: Arc::new(Mutex::new(0)),
401            released: released.clone(),
402            steals: Arc::new(Mutex::new(0)),
403        };
404        let pre_data = [1..20, 41..48];
405        let task_queue = TaskQueue::new(pre_data.iter().cloned());
406        // Spin up 8 workers and hold them in-flight on the `released` gate, then
407        // shrink to 2 mid-run. The 6 excess workers are genuinely cancelled (the
408        // worker loop's `.await` point makes `abort` effective) and their
409        // remaining ranges are reclaimed into `waiting`.
410        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
411        task_queue.set_threads(2, 1, Some(&executor)).unwrap();
412        // The decrease branch must have actually reduced the running pool.
413        assert_eq!(task_queue.inner.lock().running.len(), 2);
414        // Release the survivors so they drain `waiting` via a working `steal`
415        // and finish every number exactly once.
416        released.store(true, std::sync::atomic::Ordering::Relaxed);
417        drop(executor);
418        let mut seen = HashSet::new();
419        while let Some((_, i, res)) = rx.recv().await {
420            assert!(seen.insert(i), "number {i} was computed twice");
421            assert_eq!(res, i);
422        }
423        // Every number must be present despite the mid-run shrink: the reclaimed
424        // ranges were picked up by the survivors via a working `steal` (this is
425        // the invariant `TokioExecutor`'s broken `is_self` could never prove).
426        for range in pre_data {
427            for i in range {
428                assert!(seen.contains(&i), "number {i} was never computed");
429            }
430        }
431        assert_eq!(seen.len(), 26);
432    }
433
434    /// A *correct* executor used to genuinely exercise work-stealing and
435    /// mid-run reclaim. Unlike [`TokioExecutor`], its handle carries a real
436    /// worker id so [`Handle::is_self`] resolves, and the worker loop has an
437    /// `.await` point so [`Handle::abort`] actually cancels in-flight tasks.
438    ///
439    /// This is what lets a shrink truly reclaim a busy worker's remaining range
440    /// into `waiting` and have the survivors pick it up via [`steal`], instead
441    /// of the existing tests' brute-force completion + `safe_add_start`
442    /// deduplication.
443    struct StealExecutor {
444        tx: mpsc::UnboundedSender<(usize, u64, u64)>,
445        speculative: usize,
446        next_id: Arc<Mutex<usize>>,
447        released: Arc<std::sync::atomic::AtomicBool>,
448        steals: Arc<Mutex<usize>>,
449    }
450    #[derive(Clone)]
451    struct StealHandle {
452        abort: AbortHandle,
453        id: usize,
454    }
455    impl Handle for StealHandle {
456        type Id = usize;
457        fn abort(&mut self) {
458            self.abort.abort();
459        }
460        fn is_self(&self, id: &usize) -> bool {
461            self.id == *id
462        }
463    }
464    impl Executor for StealExecutor {
465        type Handle = StealHandle;
466        fn execute(&self, mut task: Task, q: TaskQueue<Self::Handle>) -> Self::Handle {
467            let id = {
468                let mut g = self.next_id.lock().unwrap();
469                let i = *g;
470                *g += 1;
471                i
472            };
473            let tx = self.tx.clone();
474            let speculative = self.speculative;
475            let released = self.released.clone();
476            let steals = self.steals.clone();
477            // Stay in-flight (and keep `abort` effective via the `.await` point)
478            // until the test flips `released`, so a mid-run shrink sees this
479            // worker as still running.
480            let handle = tokio::spawn(async move {
481                while !released.load(std::sync::atomic::Ordering::Relaxed) {
482                    tokio::task::yield_now().await;
483                }
484                loop {
485                    while task.start() < task.end() {
486                        // Yield between numbers so a busy worker stays
487                        // schedulable while its peers steal from it, instead of
488                        // finishing its whole range in one uninterruptible burst.
489                        tokio::task::yield_now().await;
490                        let i = task.start();
491                        let res = i;
492                        if task.safe_add_start(i, 1).is_err() {
493                            continue;
494                        }
495                        tx.send((id, i, res)).unwrap();
496                    }
497                    tokio::task::yield_now().await;
498                    if !q.steal(&id, &mut task, 1, speculative) {
499                        break;
500                    }
501                    *steals.lock().unwrap() += 1;
502                }
503            });
504            StealHandle {
505                abort: handle.abort_handle(),
506                id,
507            }
508        }
509    }
510
511    /// Genuinely verifies work-stealing: one worker gets a huge range while the
512    /// other seven get single-number crumbs, so the crumb workers *must* `steal`
513    /// from the busy peer to finish. With a working `steal` every crumb worker
514    /// ends up computing more than its initial 1-number crumb; if `steal` were a
515    /// no-op (e.g. `is_self` broken) only the single big worker does >1 number.
516    #[tokio::test(flavor = "multi_thread")]
517    async fn test_steal_distributes_work() {
518        let (tx, mut rx) = mpsc::unbounded_channel();
519        let released = Arc::new(std::sync::atomic::AtomicBool::new(false));
520        let executor = StealExecutor {
521            tx,
522            speculative: 1,
523            next_id: Arc::new(Mutex::new(0)),
524            released: released.clone(),
525            steals: Arc::new(Mutex::new(0)),
526        };
527        // One big task [7..1000] plus seven single-number crumbs.
528        let pre_data = [0..1, 1..2, 2..3, 3..4, 4..5, 5..6, 6..7, 7..1000];
529        let task_queue = TaskQueue::new(pre_data.iter().cloned());
530        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
531        drop(executor);
532        released.store(true, std::sync::atomic::Ordering::Relaxed);
533        let mut seen = std::collections::HashSet::new();
534        let mut per_worker = HashMap::new();
535        while let Some((wid, i, res)) = rx.recv().await {
536            assert!(seen.insert(i), "number {i} computed twice");
537            assert_eq!(res, i);
538            *per_worker.entry(wid).or_insert(0) += 1;
539        }
540        assert_eq!(seen.len(), 1000, "not all numbers were computed");
541        // The discriminating check: with working steal, crumb workers steal from
542        // the big peer, so at least one of them ends up doing far more than its
543        // initial 1-number crumb. A broken `is_self` makes `steal` a no-op (no
544        // `worker_idx` is found), leaving *exactly one* worker (the big one)
545        // above 1.
546        //
547        // The threshold is therefore `>= 2`: it is the sound invariant, not a
548        // statistical guess. A working steal *always* yields at least 2 workers
549        // above 1 (the big worker plus at least one stealer) — the only way to
550        // land at 1 is zero steals, i.e. the broken case. Asserting a higher
551        // count (>= 3) was flaky: how many crumb workers get a bite depends on
552        // scheduling, and on a fast or loaded runner two hot workers can drain
553        // the whole range before their peers win the race to steal,
554        // legitimately leaving exactly 2 workers above 1. That is still correct
555        // stealing, so it must not fail the test.
556        let multi = per_worker.values().filter(|&&c| c > 1).count();
557        assert!(
558            multi >= 2,
559            "steal did not distribute work; only {multi} workers exceeded their \
560             initial crumb (per-worker counts: {per_worker:?})"
561        );
562    }
563
564    /// Genuinely verifies mid-run reclaim: 8 workers split a big task, then the
565    /// pool is cut to 2. The 6 aborted workers are truly cancelled (the worker
566    /// loop's `.await` point makes `abort` effective) and their remaining ranges
567    /// are reclaimed into `waiting`, where the 2 survivors pick them up via
568    /// `steal`. If reclaim or `steal` failed, those reclaimed ranges would be
569    /// lost and the count would fall short of 1000.
570    #[tokio::test(flavor = "multi_thread")]
571    async fn test_set_threads_decrease_reclaims_via_steal() {
572        let (tx, mut rx) = mpsc::unbounded_channel();
573        let released = Arc::new(std::sync::atomic::AtomicBool::new(false));
574        let executor = StealExecutor {
575            tx,
576            speculative: 1,
577            next_id: Arc::new(Mutex::new(0)),
578            released: released.clone(),
579            steals: Arc::new(Mutex::new(0)),
580        };
581        let task_queue = TaskQueue::new(std::iter::once(0..1000));
582        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
583        // Workers are spinning on `released` (in-flight), so the shrink sees
584        // them as still running.
585        task_queue.set_threads(2, 1, Some(&executor)).unwrap();
586        assert_eq!(task_queue.inner.lock().running.len(), 2);
587        released.store(true, std::sync::atomic::Ordering::Relaxed);
588        drop(executor);
589        let mut seen = std::collections::HashSet::new();
590        while let Some((_, i, res)) = rx.recv().await {
591            assert!(seen.insert(i), "number {i} computed twice (reclaim failed)");
592            assert_eq!(res, i);
593        }
594        // All 1000 must be present: the reclaimed ranges were picked up by the
595        // survivors via steal. A broken `is_self` (steal no-op) loses them.
596        assert_eq!(seen.len(), 1000, "reclaimed work was lost");
597    }
598
599    /// The README work-stealing example is `no_run`, so its
600    /// doctest only compiles and never actually executes a steal. This mirrors
601    /// that example (one big task + crumb workers, a genuine `is_self`) and
602    /// asserts that steals *did* happen -- not merely that the result is correct,
603    /// which static pre-slicing would also satisfy.
604    #[tokio::test(flavor = "multi_thread")]
605    async fn test_readme_steal_actually_happens() {
606        let (tx, mut rx) = mpsc::unbounded_channel();
607        let released = Arc::new(AtomicBool::new(false));
608        let steals = Arc::new(Mutex::new(0));
609        let executor = StealExecutor {
610            tx,
611            speculative: 1,
612            next_id: Arc::new(Mutex::new(0)),
613            released: released.clone(),
614            steals: steals.clone(),
615        };
616        // README pattern: one big task plus seven single-number crumbs.
617        let pre_data = [0..1, 1..2, 2..3, 3..4, 4..5, 5..6, 6..7, 7..1000];
618        let task_queue = TaskQueue::new(pre_data.iter().cloned());
619        task_queue.set_threads(8, 1, Some(&executor)).unwrap();
620        drop(executor);
621        released.store(true, Ordering::Relaxed);
622        let mut seen = HashSet::new();
623        while let Some((_, i, res)) = rx.recv().await {
624            assert!(seen.insert(i), "number {i} computed twice");
625            assert_eq!(res, i);
626        }
627        assert_eq!(seen.len(), 1000, "not all numbers were computed");
628        // The discriminating check: steals must have actually occurred.
629        assert!(
630            *steals.lock().unwrap() > 0,
631            "no steal ever happened -- the README example would be silently broken"
632        );
633    }
634
635    /// A lightweight executor used only to exercise `set_threads` bookkeeping
636    /// without performing real work. It records how many workers it spawned and
637    /// keeps each handed [`Task`] alive in a never-ending background task so the
638    /// `running` deque stays populated and inspectable after the call returns.
639    struct HoldExecutor {
640        spawned: Arc<Mutex<usize>>,
641        tasks: Arc<Mutex<Vec<Task>>>,
642    }
643    struct HoldHandle;
644    impl Handle for HoldHandle {
645        type Id = ();
646        fn abort(&mut self) {}
647        fn is_self(&self, (): &()) -> bool {
648            false
649        }
650    }
651    impl Executor for HoldExecutor {
652        type Handle = HoldHandle;
653        fn execute(&self, task: Task, _q: TaskQueue<Self::Handle>) -> Self::Handle {
654            *self.spawned.lock().unwrap() += 1;
655            // Keep the `Task` alive with a strong reference so `Weak::upgrade`
656            // during `set_threads` reclaim always succeeds. No real work is done
657            // and no background task is spawned, so the test runtime exits cleanly.
658            self.tasks.lock().unwrap().push(task);
659            HoldHandle
660        }
661    }
662
663    #[tokio::test(flavor = "multi_thread")]
664    async fn test_set_threads_increase_spawns_exact_count() {
665        let ex = HoldExecutor {
666            spawned: Arc::new(Mutex::new(0)),
667            tasks: Arc::new(Mutex::new(Vec::new())),
668        };
669        let q = TaskQueue::new(std::iter::once(0..100));
670        q.set_threads(4, 1, Some(&ex)).unwrap();
671        assert_eq!(q.inner.lock().running.len(), 4);
672        assert_eq!(*ex.spawned.lock().unwrap(), 4);
673    }
674
675    #[tokio::test(flavor = "multi_thread")]
676    async fn test_set_threads_decrease_aborts_and_reclaims() {
677        let ex = HoldExecutor {
678            spawned: Arc::new(Mutex::new(0)),
679            tasks: Arc::new(Mutex::new(Vec::new())),
680        };
681        let q = TaskQueue::new(std::iter::once(0..100));
682        q.set_threads(4, 1, Some(&ex)).unwrap();
683        assert_eq!(q.inner.lock().running.len(), 4);
684        // Shrink to a single worker: the other 3 must be aborted and reclaimed.
685        q.set_threads(1, 1, Some(&ex)).unwrap();
686        assert_eq!(q.inner.lock().running.len(), 1);
687        // Tasks are reclaimed, not re-spawned: spawn count is unchanged.
688        assert_eq!(*ex.spawned.lock().unwrap(), 4);
689    }
690
691    #[tokio::test(flavor = "multi_thread")]
692    async fn test_set_threads_noop_keeps_worker_count() {
693        let ex = HoldExecutor {
694            spawned: Arc::new(Mutex::new(0)),
695            tasks: Arc::new(Mutex::new(Vec::new())),
696        };
697        let q = TaskQueue::new(std::iter::once(0..100));
698        q.set_threads(2, 1, Some(&ex)).unwrap();
699        assert_eq!(q.inner.lock().running.len(), 2);
700        // Calling again with the same count must be a no-op.
701        q.set_threads(2, 1, Some(&ex)).unwrap();
702        assert_eq!(q.inner.lock().running.len(), 2);
703    }
704
705    #[tokio::test(flavor = "multi_thread")]
706    async fn test_set_threads_none_executor_early_return() {
707        let ex = HoldExecutor {
708            spawned: Arc::new(Mutex::new(0)),
709            tasks: Arc::new(Mutex::new(Vec::new())),
710        };
711        let q = TaskQueue::new(std::iter::once(0..100));
712        // Need more workers but no executor available -> early return `None`.
713        assert!(q.set_threads::<HoldExecutor>(4, 1, None).is_none());
714        assert_eq!(q.inner.lock().running.len(), 0);
715        // Bring up 2 workers with a real executor.
716        q.set_threads(2, 1, Some(&ex)).unwrap();
717        assert_eq!(q.inner.lock().running.len(), 2);
718        // len == threads -> no-op branch, returns `Some(())` even without executor.
719        assert!(q.set_threads::<HoldExecutor>(2, 1, None).is_some());
720        assert_eq!(q.inner.lock().running.len(), 2);
721    }
722
723    /// Repeated increase/decrease must never lose or duplicate a task.
724    ///
725    /// Fifty independent single-element ranges are used so `remain == 1` and
726    /// `split_two` never fires; every move is then a pure `waiting` <-> `running`
727    /// transfer. The core invariant `waiting.len() + running.len() == total` must
728    /// hold after every resize, and `running` must land exactly at the requested
729    /// size (clamped to what `waiting` can supply). A broken hand-off (e.g. a
730    /// failed `Weak::upgrade` during reclaim, or a double-drain on increase) would
731    /// break one of these assertions immediately.
732    #[tokio::test(flavor = "multi_thread")]
733    async fn test_set_threads_oscillate_keeps_invariant() {
734        let ex = HoldExecutor {
735            spawned: Arc::new(Mutex::new(0)),
736            tasks: Arc::new(Mutex::new(Vec::new())),
737        };
738        // 50 independent single-element tasks: remain == 1, so `split_two` is dead
739        // code here and each resize is a deterministic transfer.
740        let q = TaskQueue::new((0..50).map(|i| i..i + 1));
741        let total = {
742            let g = q.inner.lock();
743            g.waiting.len() + g.running.len()
744        };
745        assert_eq!(total, 50);
746        // Oscillate the pool size many times across increase and decrease.
747        let pattern = [8usize, 2, 8, 3, 8, 1, 8, 4, 8, 2, 5, 8, 1];
748        for &threads in &pattern {
749            q.set_threads(threads, 1, Some(&ex)).unwrap();
750            let guard = q.inner.lock();
751            let running = guard.running.len();
752            let waiting = guard.waiting.len();
753            // No task may vanish or be double-claimed across a resize.
754            assert_eq!(
755                running + waiting,
756                total,
757                "task lost/duplicated at threads={threads}"
758            );
759            // `running` must match the request, clamped to the available pool.
760            assert_eq!(
761                running,
762                threads.min(total),
763                "running {running} != min({threads}, {total}) at threads={threads}"
764            );
765            drop(guard);
766        }
767        // `threads.max(1)` clamps zero to one; with `running == 1` after the
768        // oscillation pattern this is a no-op, leaving the single worker alive.
769        q.set_threads(0, 1, Some(&ex)).unwrap();
770        let guard = q.inner.lock();
771        assert_eq!(guard.running.len(), 1);
772        assert_eq!(guard.waiting.len(), total - 1);
773        drop(guard);
774    }
775
776    /// `add` pushes onto the waiting queue (`task_queue.rs` 50-53) and a live worker's
777    /// `steal` then pulls that freshly-added task off `waiting` (the `found = true;
778    /// break` branch, `task_queue.rs` line 91) before falling back to stealing from a
779    /// busy peer.
780    #[tokio::test(flavor = "multi_thread")]
781    async fn test_add_then_steal_pulls_from_waiting() {
782        let (tx, mut rx) = mpsc::unbounded_channel();
783        let released = Arc::new(AtomicBool::new(false));
784        let executor = StealExecutor {
785            tx,
786            speculative: 1,
787            next_id: Arc::new(Mutex::new(0)),
788            released: released.clone(),
789            steals: Arc::new(Mutex::new(0)),
790        };
791        // A single tiny initial task so exactly one worker is registered and the
792        // waiting queue starts empty.
793        let q = TaskQueue::new(std::iter::once(0..1));
794        q.set_threads(1, 1, Some(&executor)).unwrap();
795        // `add` lands a brand-new task in `waiting` (covers 50-53). A live worker
796        // exists, so `add` reports `true`.
797        assert!(
798            q.add(Task::new(100..110)),
799            "a live worker exists to pick up the added task"
800        );
801        assert_eq!(q.inner.lock().waiting.len(), 1);
802
803        drop(executor);
804        released.store(true, Ordering::Relaxed);
805
806        let mut seen = HashSet::new();
807        while let Some((_id, i, _res)) = rx.recv().await {
808            seen.insert(i);
809        }
810        // The initial task ran...
811        assert!(seen.contains(&0), "initial task was not executed");
812        // ...and the `add`ed task was pulled from `waiting` via `steal` (line 91).
813        for i in 100..110 {
814            assert!(
815                seen.contains(&i),
816                "added task range element {i} was never stolen"
817            );
818        }
819    }
820
821    // ---------------------------------------------------------------------
822    // Deterministic, runtime-free queue-level tests.
823    //
824    // The tokio executors above drive the queue end-to-end but can only observe
825    // *outcomes*: which branch of `steal` runs is decided by scheduling luck.
826    // `SyncExecutor` spawns nothing at all -- it hands each worker a real id and
827    // parks its `Task` in a slot -- so a test can call `steal` / `cancel_task` /
828    // `set_threads` directly from the test thread and assert on the queue's
829    // internal bookkeeping with zero races.
830    // ---------------------------------------------------------------------
831
832    struct SyncExecutor {
833        /// One slot per spawned worker. `None` means that worker has exited and
834        /// released the strong reference it held on its task.
835        slots: Arc<Mutex<Vec<Option<Task>>>>,
836        aborted: Arc<Mutex<Vec<usize>>>,
837    }
838    struct SyncHandle {
839        id: usize,
840        aborted: Arc<Mutex<Vec<usize>>>,
841    }
842    impl Handle for SyncHandle {
843        type Id = usize;
844        fn abort(&mut self) {
845            self.aborted.lock().unwrap().push(self.id);
846        }
847        fn is_self(&self, id: &usize) -> bool {
848            self.id == *id
849        }
850    }
851    impl Executor for SyncExecutor {
852        type Handle = SyncHandle;
853        fn execute(&self, task: Task, _q: TaskQueue<Self::Handle>) -> Self::Handle {
854            // Locks a mutex that is *not* the queue's, honouring the "never
855            // re-enter TaskQueue from execute" contract documented in executor.rs.
856            let id = {
857                let mut slots = self.slots.lock().unwrap();
858                let id = slots.len();
859                slots.push(Some(task));
860                id
861            };
862            SyncHandle {
863                id,
864                aborted: self.aborted.clone(),
865            }
866        }
867    }
868    impl SyncExecutor {
869        fn new() -> Self {
870            Self {
871                slots: Arc::new(Mutex::new(Vec::new())),
872                aborted: Arc::new(Mutex::new(Vec::new())),
873            }
874        }
875        /// The task worker `id` currently holds, as a state-sharing clone.
876        fn task_of(&self, id: usize) -> Task {
877            self.slots.lock().unwrap()[id].clone().unwrap()
878        }
879        /// Mirror a real worker's local `task` variable being replaced by `steal`.
880        fn rebind(&self, id: usize, task: &Task) {
881            self.slots.lock().unwrap()[id] = Some(task.clone());
882        }
883        /// Simulate worker `id` exiting: it drops its strong reference.
884        fn kill(&self, id: usize) {
885            self.slots.lock().unwrap()[id] = None;
886        }
887        fn live_workers(&self) -> usize {
888            self.slots.lock().unwrap().iter().flatten().count()
889        }
890        fn aborted(&self) -> Vec<usize> {
891            self.aborted.lock().unwrap().clone()
892        }
893    }
894
895    /// A task whose range invariant is broken (`start > end`), built through the
896    /// raw state field because `Task::new` would (correctly) refuse to make one.
897    fn corrupted_task() -> Task {
898        Task::from_raw_state(Arc::new(portable_atomic::AtomicU128::new(
899            (20u128 << 64) | 0xA,
900        )))
901    }
902
903    /// `Clone` is hand-written rather than derived (a derive would wrongly demand
904    /// `H: Clone`). It must alias the shared inner state, not copy it.
905    #[test]
906    fn clone_shares_the_same_inner_queue() {
907        let q: TaskQueue<SyncHandle> = TaskQueue::new(core::iter::empty());
908        let q2 = q.clone();
909        let _ = q2.add(Task::new(0..5));
910        assert_eq!(q.inner.lock().waiting.len(), 1);
911        assert!(Arc::ptr_eq(&q.inner, &q2.inner));
912    }
913
914    /// `TaskQueue::new` funnels every range through `Task::from`, so it inherits
915    /// that function's panic contract -- yet its own doc comment never mentions
916    /// it, and clippy's `missing_panics_doc` cannot see across the call.
917    #[test]
918    #[should_panic(expected = "range.start <= range.end")]
919    fn new_inherits_the_reversed_range_panic() {
920        // Struct literal on purpose: an inline `10..5` trips
921        // `clippy::reversed_empty_ranges`.
922        let bad = core::ops::Range {
923            start: 10u64,
924            end: 5u64,
925        };
926        let _: TaskQueue<SyncHandle> = TaskQueue::new(core::iter::once(bad));
927    }
928
929    /// `add` used to only append to `waiting`. It woke nobody and
930    /// returns `()`, so once every worker has exited (each one's `steal` returned
931    /// `false` and it left its loop) an added task is stranded forever with no
932    /// signal to the caller.
933    ///
934    /// The same test pins a related surprise: `set_threads` reports success even
935    /// when it spawned nothing at all because there was no work to hand out.
936    #[test]
937    fn add_is_inert_without_a_live_worker() {
938        let ex = SyncExecutor::new();
939        let q: TaskQueue<SyncHandle> = TaskQueue::new(core::iter::empty());
940        assert!(q.set_threads(4, 1, Some(&ex)).is_some());
941        assert_eq!(q.inner.lock().running.len(), 0, "nothing was spawned");
942
943        assert!(
944            !q.add(Task::new(0..100)),
945            "no live worker: the task is stranded"
946        );
947        assert_eq!(q.inner.lock().waiting.len(), 1);
948        // No worker exists to call `steal`, so only an explicit `set_threads`
949        // ever rescues the task.
950        q.set_threads(1, 1, Some(&ex)).unwrap();
951        let guard = q.inner.lock();
952        assert_eq!(guard.running.len(), 1);
953        assert_eq!(guard.waiting.len(), 0);
954        drop(guard);
955    }
956
957    /// The `true` branch of `add`: when at least one live worker exists,
958    /// `add` reports `true` (a worker can pick the new task up via `steal`).
959    #[test]
960    fn add_returns_true_when_a_live_worker_exists() {
961        let ex = SyncExecutor::new();
962        let q = TaskQueue::new(core::iter::once(0..1));
963        q.set_threads(1, 1, Some(&ex)).unwrap();
964        assert_eq!(q.inner.lock().running.len(), 1, "one live worker spawned");
965        assert!(
966            q.add(Task::new(5..6)),
967            "a live worker exists to pick up the new task"
968        );
969    }
970
971    /// An unregistered caller gets exactly the same `false` a worker gets when the
972    /// queue is drained -- even with work sitting in `waiting`.
973    #[test]
974    fn steal_rejects_an_unregistered_worker() {
975        let ex = SyncExecutor::new();
976        let q = TaskQueue::new(core::iter::once(0..100));
977        q.set_threads(1, 1, Some(&ex)).unwrap();
978        let _ = q.add(Task::new(500..510));
979        let mut t = Task::new(0..0);
980        assert!(!q.steal(&999, &mut t, 1, 1));
981        assert_eq!(t.get(), 0..0);
982        assert_eq!(
983            q.inner.lock().waiting.len(),
984            1,
985            "waiting must not be disturbed by an unknown caller"
986        );
987    }
988
989    #[test]
990    fn steal_prefers_waiting_over_robbing_a_peer() {
991        let ex = SyncExecutor::new();
992        let q = TaskQueue::new([0..2, 100..200].into_iter());
993        q.set_threads(2, 1, Some(&ex)).unwrap();
994        let _ = q.add(Task::new(500..510));
995        let mut t = ex.task_of(0);
996        assert!(q.steal(&0, &mut t, 1, 1));
997        assert_eq!(t.get(), 500..510);
998        assert_eq!(
999            ex.task_of(1).get(),
1000            100..200,
1001            "the fat peer keeps its range: waiting wins over split_two"
1002        );
1003        // The worker is re-registered against its brand-new task.
1004        assert_eq!(q.inner.lock().running[0].0.upgrade().unwrap(), t);
1005    }
1006
1007    /// The waiting drain loop pops until it finds *usable* work: an exhausted task
1008    /// yields `Ok(None)` and a corrupted one yields `Err`; both are discarded
1009    /// rather than handed to a worker.
1010    ///
1011    /// Note this path only became safe with the `take` hardening: while `take`
1012    /// still returned `Some(20..10)` for a corrupted task, `Task::new(range)`
1013    /// below would have tripped its `start <= end` assertion instead.
1014    #[test]
1015    fn steal_skips_exhausted_and_corrupted_waiting_tasks() {
1016        let ex = SyncExecutor::new();
1017        let q = TaskQueue::new(core::iter::once(0..2));
1018        q.set_threads(1, 1, Some(&ex)).unwrap();
1019        let _ = q.add(Task::new(7..7));
1020        let _ = q.add(corrupted_task());
1021        let _ = q.add(Task::new(500..510));
1022        let mut t = ex.task_of(0);
1023        assert!(q.steal(&0, &mut t, 1, 1));
1024        assert_eq!(t.get(), 500..510);
1025        assert_eq!(
1026            q.inner.lock().waiting.len(),
1027            0,
1028            "all three were popped; the two unusable ones are dropped"
1029        );
1030    }
1031
1032    /// A lone worker cannot steal from itself: the victim scan filters out any
1033    /// running task pointer-equal to the caller's own.
1034    #[test]
1035    fn steal_excludes_the_caller_as_a_victim() {
1036        let ex = SyncExecutor::new();
1037        let q = TaskQueue::new(core::iter::once(0..1000));
1038        q.set_threads(1, 1, Some(&ex)).unwrap();
1039        let mut t = ex.task_of(0);
1040        assert!(!q.steal(&0, &mut t, 1, 2));
1041        assert_eq!(t.get(), 0..1000, "the caller's own range must be untouched");
1042    }
1043
1044    #[test]
1045    fn steal_splits_the_busiest_peer() {
1046        let ex = SyncExecutor::new();
1047        let q = TaskQueue::new([0..2, 10..14, 100..200].into_iter());
1048        q.set_threads(3, 1, Some(&ex)).unwrap();
1049        let mut t = ex.task_of(0);
1050        assert!(q.steal(&0, &mut t, 1, 1));
1051        // 100..200 has the most work left, so it -- not 10..14 -- is halved.
1052        assert_eq!(t.get(), 150..200);
1053        assert_eq!(ex.task_of(2).get(), 100..150);
1054        assert_eq!(
1055            ex.task_of(1).get(),
1056            10..14,
1057            "the smaller peer is left alone"
1058        );
1059    }
1060
1061    /// When the fattest peer is too small to halve (`remain < min_chunk_size * 2`)
1062    /// the queue falls back to *sharing* it -- but only with speculation enabled.
1063    ///
1064    /// The two cases are tested separately because a worker that exhausts its
1065    /// task and finds no stealable work is deregistered from `running` (so a
1066    /// concurrent shrink does not mistakenly wait for it). In production a
1067    /// worker exits after its first `steal` returns `false`; retrying with
1068    /// different parameters on a deregistered worker is a test-only scenario.
1069    #[test]
1070    fn steal_shares_speculatively_only_when_allowed() {
1071        // Case 1: speculation disabled -> the crumb is left alone.
1072        let ex = SyncExecutor::new();
1073        let q = TaskQueue::new([0..1, 100..101].into_iter());
1074        q.set_threads(2, 1, Some(&ex)).unwrap();
1075        let mut t = ex.task_of(0);
1076        assert!(!q.steal(&0, &mut t, 1, 1));
1077        assert_eq!(t.get(), 0..1);
1078
1079        // Case 2: speculation enabled -> the caller aliases the peer's state.
1080        let ex = SyncExecutor::new();
1081        let q = TaskQueue::new([0..1, 100..101].into_iter());
1082        q.set_threads(2, 1, Some(&ex)).unwrap();
1083        let mut t = ex.task_of(0);
1084        assert!(q.steal(&0, &mut t, 1, 2));
1085        assert_eq!(t.get(), 100..101);
1086        assert_eq!(t, ex.task_of(1), "speculation must alias, not copy");
1087    }
1088
1089    /// The speculation cap limits how many workers share one cursor. Each sharer
1090    /// holds its own strong ref to the cursor (via `share_state`), and the cap
1091    /// `sharer_count() < max_speculative` admits a new sharer only while fewer
1092    /// than `max_speculative` workers already alias that cursor — keeping the
1093    /// total at `max_speculative`.
1094    #[test]
1095    fn steal_caps_the_number_of_speculative_sharers() {
1096        let ex = SyncExecutor::new();
1097        let q = TaskQueue::new([0..1, 1..2, 2..3].into_iter());
1098        q.set_threads(3, 1, Some(&ex)).unwrap();
1099        // All crumbs tie on `remain`, and `max_by_key` documents that ties resolve
1100        // to the *last* element, so worker 0 joins worker 2.
1101        let mut t0 = ex.task_of(0);
1102        assert!(q.steal(&0, &mut t0, 1, 2));
1103        ex.rebind(0, &t0);
1104        assert_eq!(t0, ex.task_of(2));
1105        // Worker 1 tries to become the third sharer of that same crumb and is
1106        // refused: the sharer count now exceeds the cap.
1107        let mut t1 = ex.task_of(1);
1108        assert!(!q.steal(&1, &mut t1, 1, 2));
1109        assert_eq!(t1.get(), 1..2, "the refused worker keeps its own range");
1110    }
1111
1112    /// `min_chunk_size * 2` used to be an unchecked
1113    /// multiplication on a caller-supplied `u64` — `fast-pull` forwards a
1114    /// user-configurable `options.min_chunk_size` straight into it, so a large
1115    /// value panicked in debug and silently wrapped (disabling split) in release.
1116    /// It now uses `saturating_mul`, so an overflowing value simply caps at
1117    /// `u64::MAX` and the split branch is skipped gracefully instead of panicking.
1118    #[test]
1119    fn steal_skips_split_on_a_huge_min_chunk_size() {
1120        let ex = SyncExecutor::new();
1121        let q = TaskQueue::new([0..2, 100..200].into_iter());
1122        q.set_threads(2, 1, Some(&ex)).unwrap();
1123        let mut t = ex.task_of(0);
1124        // No panic, no silent disable: the caller keeps its own tiny range because
1125        // the split branch is simply never taken (`remain >= u64::MAX` is false).
1126        assert!(!q.steal(&0, &mut t, u64::MAX, 1));
1127        assert_eq!(
1128            t.get(),
1129            0..2,
1130            "the caller is untouched when split is skipped"
1131        );
1132    }
1133
1134    /// The identical `min_chunk_size * 2` in `set_threads`'s split-to-grow loop is
1135    /// now `saturating_mul` too: an overflowing value no longer
1136    /// panics, the split-to-grow branch is just skipped.
1137    #[test]
1138    fn set_threads_skips_split_on_a_huge_min_chunk_size() {
1139        let ex = SyncExecutor::new();
1140        let q = TaskQueue::new(core::iter::once(0..100));
1141        assert!(q.set_threads(2, u64::MAX, Some(&ex)).is_some());
1142        // The single waiting task is still spawned, but no extra split worker is
1143        // created because `remain >= u64::MAX` is always false.
1144        assert_eq!(q.inner.lock().running.len(), 1);
1145    }
1146
1147    /// Baseline for the next test: with no sharing, one worker owns one task, so
1148    /// the liveness sweep collects its slot as soon as it exits.
1149    #[test]
1150    fn running_sweep_collects_an_exited_worker() {
1151        let ex = SyncExecutor::new();
1152        let q = TaskQueue::new([0..1, 100..101].into_iter());
1153        q.set_threads(2, 1, Some(&ex)).unwrap();
1154        assert_eq!(q.inner.lock().running.len(), 2);
1155        ex.kill(0);
1156        let _ = q.set_threads(2, 1, Some(&ex));
1157        assert_eq!(q.inner.lock().running.len(), 1, "the dead slot is swept");
1158    }
1159
1160    /// The liveness sweep used to key off the *cursor*
1161    /// refcount, which speculative sharing defeats — a dead worker's slot stayed
1162    /// propped up by its surviving twin and `set_threads` never refilled the pool.
1163    /// `WeakTask` now points at the worker's own identity (`TaskInner`), so the
1164    /// sweep reclaims a dead worker's slot regardless of how many twins share its
1165    /// cursor.
1166    #[test]
1167    fn speculative_sharing_no_longer_defeats_liveness_sweep() {
1168        let ex = SyncExecutor::new();
1169        let q = TaskQueue::new([0..1, 100..101].into_iter());
1170        q.set_threads(2, 1, Some(&ex)).unwrap();
1171
1172        let mut t0 = ex.task_of(0);
1173        assert!(q.steal(&0, &mut t0, 1, 2));
1174        ex.rebind(0, &t0);
1175        // Both `running` entries now weak-point at one and the same cursor.
1176        let guard = q.inner.lock();
1177        assert_eq!(
1178            guard.running[0].0.upgrade().unwrap(),
1179            guard.running[1].0.upgrade().unwrap()
1180        );
1181        drop(guard);
1182
1183        // Worker 0 exits, leaving worker 1 as the only live worker.
1184        ex.kill(0);
1185        drop(t0);
1186        assert_eq!(ex.live_workers(), 1);
1187
1188        // The liveness sweep alone must reclaim the dead worker's slot, even
1189        // though its speculative twin still references the shared cursor. Passing
1190        // `None` for the executor runs the sweep without spawning replacements.
1191        let _ = q.set_threads::<SyncExecutor>(2, 1, None);
1192        assert_eq!(
1193            q.inner.lock().running.len(),
1194            1,
1195            "dead worker reclaimed despite its speculative twin"
1196        );
1197    }
1198
1199    /// `handles` had no coverage inside `fast-steal` at all -- its only consumer
1200    /// lives in `fast-pull`. It must expose every running handle by mutable
1201    /// reference and hand the closure's return value back out.
1202    #[test]
1203    fn handles_exposes_every_running_worker() {
1204        let ex = SyncExecutor::new();
1205        let q = TaskQueue::new([0..1, 1..2, 2..3].into_iter());
1206        q.set_threads(3, 1, Some(&ex)).unwrap();
1207        let ids = q.handles(|iter| iter.map(|h| h.id).collect::<Vec<_>>());
1208        assert_eq!(ids, [0, 1, 2]);
1209        q.handles(|iter| {
1210            for h in iter {
1211                h.abort();
1212            }
1213        });
1214        assert_eq!(ex.aborted(), [0, 1, 2]);
1215    }
1216
1217    /// `cancel_task` is the speculation cleanup path: when one sharer finishes the
1218    /// shared range it aborts the others. It had no coverage in `fast-steal`.
1219    #[test]
1220    fn cancel_task_aborts_twins_but_spares_the_caller() {
1221        let ex = SyncExecutor::new();
1222        let q = TaskQueue::new([0..1, 100..101].into_iter());
1223        q.set_threads(2, 1, Some(&ex)).unwrap();
1224        let mut t0 = ex.task_of(0);
1225        assert!(q.steal(&0, &mut t0, 1, 2));
1226        ex.rebind(0, &t0);
1227
1228        // Worker 1 finished the shared range and cancels its twins.
1229        q.cancel_task(&t0, &1);
1230        assert_eq!(ex.aborted(), [0], "only the peer sharer is aborted");
1231        // The aborted twin is now deregistered, leaving only the
1232        // caller's own entry. (Previously it lingered in `running`.)
1233        assert_eq!(
1234            q.inner.lock().running.len(),
1235            1,
1236            "aborted twin is deregistered"
1237        );
1238    }
1239
1240    /// `cancel_task` aborts and deregisters the twin, but
1241    /// does **not** reclaim the remaining range into `waiting`. That is
1242    /// deliberate: reclaiming would race with the caller's own still-live `task`
1243    /// over the same range and cause duplicate execution. The soundness therefore
1244    /// depends on the caller invoking this only after the shared range is
1245    /// finished (as `fast-pull` does). Aimed at live work it silently strands the
1246    /// remainder — by design, not by accident.
1247    #[test]
1248    fn cancel_task_does_not_reclaim_unfinished_work() {
1249        let ex = SyncExecutor::new();
1250        let q = TaskQueue::new([0..1, 100..200].into_iter());
1251        q.set_threads(2, 1, Some(&ex)).unwrap();
1252        let victim = ex.task_of(1);
1253        assert_eq!(victim.remain(), 100);
1254        q.cancel_task(&victim, &0);
1255        assert_eq!(ex.aborted(), [1]);
1256        assert_eq!(
1257            q.inner.lock().running.len(),
1258            1,
1259            "aborted victim is deregistered"
1260        );
1261        assert_eq!(
1262            q.inner.lock().waiting.len(),
1263            0,
1264            "100 units of unfinished work were intentionally NOT reclaimed (would double-execute)"
1265        );
1266    }
1267
1268    /// An unregistered caller must not abort anyone via `cancel_task`.
1269    ///
1270    /// Without the registration guard a caller that was deregistered by a
1271    /// shrink (but is still running cooperatively) makes `is_self` vacuous —
1272    /// `false` for every entry — and aborts *all* twins, including the only
1273    /// worker that was supposed to pick up work from `waiting`.
1274    #[test]
1275    fn cancel_task_unregistered_caller_is_a_noop() {
1276        let ex = SyncExecutor::new();
1277        let q = TaskQueue::new([0..1, 100..200].into_iter());
1278        q.set_threads(2, 1, Some(&ex)).unwrap();
1279        // id 999 is not in `running`, so the call must be a no-op.
1280        q.cancel_task(&ex.task_of(1), &999);
1281        assert!(
1282            ex.aborted().is_empty(),
1283            "unregistered caller must not abort anyone"
1284        );
1285        assert_eq!(q.inner.lock().running.len(), 2, "no worker deregistered");
1286    }
1287
1288    // -----------------------------------------------------------------
1289    // Audit verification tests.
1290    // -----------------------------------------------------------------
1291
1292    /// Verifies that `retain` at the top of `set_threads` removes ALL dead
1293    /// entries before the shrink path runs, so the survivors are always alive
1294    /// and the decrease branch can proceed to reclaim their overflow peers
1295    /// without a separate liveness guard.
1296    #[test]
1297    fn set_threads_shrink_guard_never_fires_after_retain() {
1298        let ex = SyncExecutor::new();
1299        let q = TaskQueue::new((0..5).map(|i| i * 10..(i + 1) * 10));
1300        q.set_threads(5, 1, Some(&ex)).unwrap();
1301        assert_eq!(q.inner.lock().running.len(), 5);
1302
1303        // Kill the first 3 workers. A naive reading of the guard suggests
1304        // shrink-to-2 might early-return (first 2 are dead). But `retain`
1305        // removes them first, leaving [alive3, alive4], so shrink proceeds.
1306        ex.kill(0);
1307        ex.kill(1);
1308        ex.kill(2);
1309
1310        q.set_threads(2, 1, Some(&ex)).unwrap();
1311        let guard = q.inner.lock();
1312        assert_eq!(
1313            guard.running.len(),
1314            2,
1315            "retain sweeps dead entries; shrink always proceeds"
1316        );
1317        // The 2 survivors' tasks were NOT aborted.
1318        assert!(ex.aborted().is_empty() || ex.aborted().iter().all(|&id| id < 3));
1319    }
1320
1321    /// A worker deregistered by a failed `steal` (no work found) is rejected
1322    /// on all subsequent steal attempts — even after new work is added to
1323    /// `waiting`. Only `set_threads` can re-register a worker.
1324    #[test]
1325    fn steal_after_deregistration_is_permanently_rejected() {
1326        let ex = SyncExecutor::new();
1327        let q = TaskQueue::new(core::iter::once(0..1));
1328        q.set_threads(1, 1, Some(&ex)).unwrap();
1329
1330        let mut t = ex.task_of(0);
1331        // No work anywhere: steal fails and deregisters worker 0.
1332        assert!(!q.steal(&0, &mut t, 1, 1));
1333        assert_eq!(q.inner.lock().running.len(), 0, "worker deregistered");
1334
1335        // Add fresh work to waiting.
1336        let _ = q.add(Task::new(100..200));
1337        assert_eq!(q.inner.lock().waiting.len(), 1);
1338
1339        // The deregistered worker still cannot steal.
1340        assert!(!q.steal(&0, &mut t, 1, 1));
1341        assert_eq!(
1342            q.inner.lock().waiting.len(),
1343            1,
1344            "waiting undisturbed by rejected caller"
1345        );
1346
1347        // Only set_threads can rescue the stranded work.
1348        q.set_threads(1, 1, Some(&ex)).unwrap();
1349        assert_eq!(q.inner.lock().running.len(), 1);
1350        assert_eq!(q.inner.lock().waiting.len(), 0);
1351    }
1352
1353    /// Multiple registered workers calling `steal` concurrently from
1354    /// different OS threads: the mutex serialises access, so every task from
1355    /// `waiting` is handed out exactly once with no loss or duplication.
1356    #[test]
1357    fn concurrent_steal_from_waiting_no_loss_or_dup() {
1358        use std::thread;
1359
1360        let ex = SyncExecutor::new();
1361        let q = TaskQueue::new((0..4).map(|i| i * 100..(i + 1) * 100));
1362        q.set_threads(4, 1, Some(&ex)).unwrap();
1363
1364        // Add 4 fresh tasks to waiting.
1365        for i in 0..4u64 {
1366            let _ = q.add(Task::new(1000 + i * 100..1100 + i * 100));
1367        }
1368
1369        let mut handles = Vec::new();
1370        for id in 0..4usize {
1371            let q = q.clone();
1372            let mut t = ex.task_of(id);
1373            handles.push(thread::spawn(move || {
1374                let ok = q.steal(&id, &mut t, 1, 1);
1375                (id, ok, t.get())
1376            }));
1377        }
1378        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1379
1380        // All 4 succeed (4 waiting tasks available).
1381        let stolen: Vec<_> = results
1382            .iter()
1383            .filter(|(_, ok, _)| *ok)
1384            .map(|(_, _, r)| r.clone())
1385            .collect();
1386        assert_eq!(stolen.len(), 4, "all 4 workers got work from waiting");
1387
1388        // Ranges are non-overlapping and cover exactly 1000..1400.
1389        let mut sorted = stolen;
1390        sorted.sort_by_key(|r| r.start);
1391        for w in sorted.windows(2) {
1392            assert!(
1393                w[0].end <= w[1].start,
1394                "overlap detected: {:?} and {:?}",
1395                w[0],
1396                w[1]
1397            );
1398        }
1399        assert_eq!(sorted[0].start, 1000);
1400        assert_eq!(sorted[3].end, 1400);
1401    }
1402
1403    /// Multiple workers concurrently splitting the same fat peer: the CAS in
1404    /// `split_two` serialises splits so the resulting sub-ranges form a
1405    /// non-overlapping partition of the original.
1406    #[test]
1407    fn concurrent_steal_split_no_overlap() {
1408        use std::thread;
1409
1410        let ex = SyncExecutor::new();
1411        // Worker 3 gets the fat task; workers 0-2 get crumbs.
1412        let q = TaskQueue::new([0..1, 1..2, 2..3, 0..1000].into_iter());
1413        q.set_threads(4, 1, Some(&ex)).unwrap();
1414
1415        let mut handles = Vec::new();
1416        for id in 0..3usize {
1417            let q = q.clone();
1418            let mut t = ex.task_of(id);
1419            handles.push(thread::spawn(move || {
1420                let ok = q.steal(&id, &mut t, 1, 1);
1421                (id, ok, t.get())
1422            }));
1423        }
1424        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1425
1426        // All 3 must succeed by splitting the fat peer.
1427        for (id, ok, range) in &results {
1428            assert!(ok, "worker {id} failed to steal via split");
1429            assert!(range.end - range.start > 0, "worker {id} got empty range");
1430        }
1431
1432        // Sub-ranges must not overlap (they are partitions of 0..1000).
1433        let mut ranges: Vec<_> = results.iter().map(|(_, _, r)| r.clone()).collect();
1434        ranges.sort_by_key(|r| r.start);
1435        for w in ranges.windows(2) {
1436            assert!(
1437                w[0].end <= w[1].start,
1438                "split ranges overlap: {:?} and {:?}",
1439                w[0],
1440                w[1]
1441            );
1442        }
1443    }
1444
1445    /// `cancel_task` is a no-op when the registered caller has no speculative
1446    /// twin: nothing is aborted and no worker is deregistered. This guards
1447    /// against the call spuriously dropping the caller from `running`.
1448    #[test]
1449    fn cancel_task_registered_caller_without_twin_is_noop() {
1450        let ex = SyncExecutor::new();
1451        let q = TaskQueue::new([0..1, 100..200].into_iter());
1452        q.set_threads(2, 1, Some(&ex)).unwrap();
1453        // Worker 0 owns its own task and shares no cursor with anyone.
1454        q.cancel_task(&ex.task_of(0), &0);
1455        assert!(ex.aborted().is_empty(), "nothing to abort without a twin");
1456        assert_eq!(
1457            q.inner.lock().running.len(),
1458            2,
1459            "both workers must stay registered"
1460        );
1461    }
1462
1463    /// Pins the *correct* handoff when a speculative sharer is reclaimed by a
1464    /// `set_threads` shrink.
1465    ///
1466    /// A speculative sharer aliases its twin's progress cursor (same `state`
1467    /// pointer, `sharer_count` bumped by `share_state`). When the shrink aborts
1468    /// the sharer and pushes its task into `waiting`, that reclaimed task still
1469    /// points at the *surviving* twin's cursor. The next `steal` to drain it
1470    /// calls `take()` on the shared cursor: `take()` atomically claims the
1471    /// remaining range for the new worker and empties the shared cursor, so the
1472    /// survivor twin exits cleanly on its next `safe_add_start` and the range is
1473    /// executed exactly once.
1474    ///
1475    /// This guards against a regression where the handoff would be "fixed" by
1476    /// dropping the reclaimed task — which would instead LOSE the remaining work
1477    /// when the survivor is the only live worker left.
1478    #[test]
1479    fn shrink_reclaims_speculative_sharer_task_aliasing_live_cursor() {
1480        let ex = SyncExecutor::new();
1481        // `min_chunk = 10` makes the 5-element peers too small to split, forcing
1482        // the *share* branch of `steal` instead of `split_two`.
1483        let q = TaskQueue::new([0..5, 5..10].into_iter());
1484        q.set_threads(2, 10, Some(&ex)).unwrap();
1485
1486        // Worker 0 speculatively aliases worker 1's cursor (5..10).
1487        let mut t0 = ex.task_of(0);
1488        assert!(q.steal(&0, &mut t0, 10, 2));
1489        ex.rebind(0, &t0);
1490        assert_eq!(t0, ex.task_of(1), "worker 0 now aliases worker 1's cursor");
1491
1492        // Shrink to 1: aborts worker 1, reclaims its task into `waiting`.
1493        q.set_threads(1, 10, Some(&ex)).unwrap();
1494        let guard = q.inner.lock();
1495        assert_eq!(guard.running.len(), 1);
1496        let survivor = guard.running[0].0.upgrade().unwrap();
1497
1498        // The reclaimed waiting task aliases the *still-running* survivor's
1499        // cursor. Whoever steals it next re-executes the survivor's range.
1500        assert_eq!(
1501            guard.waiting.len(),
1502            1,
1503            "the aborted sharer's task was reclaimed"
1504        );
1505        assert_eq!(
1506            guard.waiting[0], survivor,
1507            "reclaimed task shares the survivor's cursor; steal's take() handshake transfers it safely"
1508        );
1509    }
1510}