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}