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