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