almost_enough/
tree.rs

1//! Hierarchical cancellation with tree structure.
2//!
3//! [`ChildStopper`] provides cancellation that can form parent-child relationships.
4//! When a parent is cancelled, all children are also cancelled. Children can be
5//! cancelled independently without affecting siblings or parents.
6//!
7//! # Overview
8//!
9//! - [`ChildStopper::new()`] - Create a root (no parent)
10//! - [`ChildStopper::with_parent()`] - Create a child of any `Stop` implementation
11//! - [`tree.child()`](ChildStopper::child) - Create a child of this tree node
12//!
13//! # Example
14//!
15//! ```rust
16//! use almost_enough::{ChildStopper, Stop};
17//!
18//! let parent = ChildStopper::new();
19//! let child_a = parent.child();
20//! let child_b = parent.child();
21//!
22//! // Children can be cancelled independently
23//! child_a.cancel();
24//! assert!(child_a.should_stop());
25//! assert!(!child_b.should_stop());
26//!
27//! // Parent cancellation propagates to all children
28//! parent.cancel();
29//! assert!(child_b.should_stop());
30//! ```
31//!
32//! # Grandchildren
33//!
34//! Children can have their own children, creating a cancellation tree:
35//!
36//! ```rust
37//! use almost_enough::{ChildStopper, Stop};
38//!
39//! let grandparent = ChildStopper::new();
40//! let parent = grandparent.child();
41//! let child = parent.child();
42//!
43//! // Grandparent cancellation propagates through the tree
44//! grandparent.cancel();
45//! assert!(parent.should_stop());
46//! assert!(child.should_stop());
47//! ```
48//!
49//! # With Other Stop Types
50//!
51//! You can create a `ChildStopper` as a child of any `Stop` implementation:
52//!
53//! ```rust
54//! use almost_enough::{Stopper, ChildStopper, Stop};
55//!
56//! let root = Stopper::new();
57//! let child = ChildStopper::with_parent(root.clone());
58//!
59//! root.cancel();
60//! assert!(child.should_stop());
61//! ```
62
63use alloc::sync::Arc;
64use core::any::{Any, TypeId};
65use core::sync::atomic::{AtomicU8, Ordering};
66
67use crate::{Stop, StopReason, StopToken, Unstoppable};
68
69/// `TreeInner::state`: not cancelled, with nothing above but parent nodes.
70const RUNNING: u8 = 0;
71/// `TreeInner::state`: this node was cancelled.
72const CANCELLED: u8 = 1;
73/// `TreeInner::state`: not cancelled, and `above` must be checked.
74const ABOVE: u8 = 2;
75
76// The checks test `state != RUNNING`, then `state & CANCELLED`; a `match` on
77// the state compiles longer on x86-64.
78const _: () = assert!(RUNNING == 0 && CANCELLED != 0 && ABOVE != 0 && ABOVE & CANCELLED == 0);
79
80/// Inner state for a tree node.
81///
82/// A check walks the chain with no vtable call: each level loads and tests
83/// the state byte, then loads and tests the parent pointer.
84struct TreeInner {
85    /// `RUNNING`, `CANCELLED`, or `ABOVE`, which `cancel()` also overwrites.
86    state: AtomicU8,
87    /// The parent node, if it is another `ChildStopper`.
88    parent: Option<Arc<TreeInner>>,
89    /// Any other parent, checked only in state `ABOVE`; `Unstoppable`
90    /// (stored as nothing) otherwise.
91    above: StopToken,
92}
93
94impl TreeInner {
95    fn new(parent: Option<Arc<TreeInner>>, above: StopToken) -> Self {
96        let state = if above.may_stop() { ABOVE } else { RUNNING };
97        Self {
98            state: AtomicU8::new(state),
99            parent,
100            above,
101        }
102    }
103
104    /// Check this node, then walk its ancestors. The first level is peeled
105    /// out of the loop so a root returns without a jump.
106    #[inline]
107    #[track_caller]
108    fn check(&self) -> Result<(), StopReason> {
109        let state = self.state.load(Ordering::Relaxed);
110        if state != RUNNING {
111            if state & CANCELLED != 0 {
112                return Err(StopReason::Cancelled);
113            }
114            return self.above.check();
115        }
116        match &self.parent {
117            None => Ok(()),
118            Some(parent) => parent.check_ancestors(),
119        }
120    }
121
122    #[inline]
123    #[track_caller]
124    fn check_ancestors(&self) -> Result<(), StopReason> {
125        let mut node = self;
126        loop {
127            if node.state.load(Ordering::Relaxed) != RUNNING {
128                return node.check_stopped_ancestor();
129            }
130            match &node.parent {
131                Some(parent) => node = parent,
132                None => return Ok(()),
133            }
134        }
135    }
136
137    #[inline]
138    #[track_caller]
139    fn should_stop(&self) -> bool {
140        let state = self.state.load(Ordering::Relaxed);
141        if state != RUNNING {
142            return state & CANCELLED != 0 || self.above.should_stop();
143        }
144        match &self.parent {
145            None => false,
146            Some(parent) => parent.should_stop_ancestors(),
147        }
148    }
149
150    /// An ancestor that is cancelled or has a stop above: out of line, so the
151    /// inlined check holds one `StopToken` dispatch, not two.
152    #[cold]
153    #[inline(never)]
154    #[track_caller]
155    fn check_stopped_ancestor(&self) -> Result<(), StopReason> {
156        if self.state.load(Ordering::Relaxed) & CANCELLED != 0 {
157            return Err(StopReason::Cancelled);
158        }
159        self.above.check()
160    }
161
162    #[cold]
163    #[inline(never)]
164    #[track_caller]
165    fn should_stop_stopped_ancestor(&self) -> bool {
166        self.state.load(Ordering::Relaxed) & CANCELLED != 0 || self.above.should_stop()
167    }
168
169    #[inline]
170    #[track_caller]
171    fn should_stop_ancestors(&self) -> bool {
172        let mut node = self;
173        loop {
174            if node.state.load(Ordering::Relaxed) != RUNNING {
175                return node.should_stop_stopped_ancestor();
176            }
177            match &node.parent {
178                Some(parent) => node = parent,
179                None => return false,
180            }
181        }
182    }
183}
184
185impl core::fmt::Debug for TreeInner {
186    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
187        let parent = if self.parent.is_some() {
188            Some("<ChildStopper>")
189        } else if self.above.may_stop() {
190            Some("<StopToken>")
191        } else {
192            None
193        };
194        f.debug_struct("TreeInner")
195            .field(
196                "self_cancelled",
197                &(self.state.load(Ordering::Relaxed) == CANCELLED),
198            )
199            .field("parent", &parent)
200            .finish()
201    }
202}
203
204/// A cancellation primitive with tree-structured parent-child relationships.
205///
206/// `ChildStopper` uses a unified clone model: clone to share, any clone can cancel.
207/// When cancelled, it does NOT affect its parent or siblings - only this node
208/// and any of its children.
209///
210/// # Example
211///
212/// ```rust
213/// use almost_enough::{ChildStopper, Stop};
214///
215/// let parent = ChildStopper::new();
216/// let child = parent.child();
217///
218/// // Clone to share across threads
219/// let child_clone = child.clone();
220///
221/// // Any clone can cancel
222/// child_clone.cancel();
223/// assert!(child.should_stop());
224///
225/// // Parent is not affected
226/// assert!(!parent.should_stop());
227/// ```
228///
229/// # Performance
230///
231/// - Size: 8 bytes (one pointer)
232/// - `check()`: per level of the tree, a state byte and a parent pointer,
233///   with no vtable call. A parent that isn't a `ChildStopper` is checked
234///   once, at the top of the chain, as a [`StopToken`] would check it.
235#[derive(Debug, Clone)]
236pub struct ChildStopper {
237    inner: Arc<TreeInner>,
238}
239
240impl ChildStopper {
241    /// Create a new root tree node (no parent).
242    ///
243    /// This creates a tree root that can have children added via [`child()`](Self::child).
244    ///
245    /// # Example
246    ///
247    /// ```rust
248    /// use almost_enough::{ChildStopper, Stop};
249    ///
250    /// let root = ChildStopper::new();
251    /// let child = root.child();
252    ///
253    /// root.cancel();
254    /// assert!(child.should_stop());
255    /// ```
256    #[inline]
257    pub fn new() -> Self {
258        Self {
259            inner: Arc::new(TreeInner::new(None, StopToken::new(Unstoppable))),
260        }
261    }
262
263    /// Create a new tree node with a parent.
264    ///
265    /// The child will stop if either:
266    /// - [`cancel()`](Self::cancel) is called on this node (or any clone)
267    /// - Any ancestor in the parent chain is cancelled
268    ///
269    /// # Example
270    ///
271    /// ```rust
272    /// use almost_enough::{Stopper, ChildStopper, Stop};
273    ///
274    /// let root = Stopper::new();
275    /// let child = ChildStopper::with_parent(root.clone());
276    ///
277    /// root.cancel();
278    /// assert!(child.should_stop());
279    /// ```
280    #[inline]
281    pub fn with_parent<T: Stop + 'static>(parent: T) -> Self {
282        if TypeId::of::<T>() == TypeId::of::<ChildStopper>() {
283            let any_ref: &dyn Any = &parent;
284            return any_ref.downcast_ref::<ChildStopper>().unwrap().child();
285        }
286        Self {
287            inner: Arc::new(TreeInner::new(None, StopToken::new(parent))),
288        }
289    }
290
291    /// Create a child of this tree node.
292    ///
293    /// The child will stop if either this node or any ancestor is cancelled.
294    /// Cancelling the child does NOT affect this node.
295    ///
296    /// # Example
297    ///
298    /// ```rust
299    /// use almost_enough::{ChildStopper, Stop};
300    ///
301    /// let parent = ChildStopper::new();
302    /// let child = parent.child();
303    /// let grandchild = child.child();
304    ///
305    /// child.cancel();
306    /// assert!(!parent.should_stop());  // Parent unaffected
307    /// assert!(child.should_stop());
308    /// assert!(grandchild.should_stop());  // Inherits from parent
309    /// ```
310    #[inline]
311    pub fn child(&self) -> ChildStopper {
312        Self {
313            inner: Arc::new(TreeInner::new(
314                Some(Arc::clone(&self.inner)),
315                StopToken::new(Unstoppable),
316            )),
317        }
318    }
319
320    /// Cancel this node (and all its children).
321    ///
322    /// This does NOT affect the parent or siblings. A Release swap: writes
323    /// made before a `cancel()` of this node are visible to a thread whose
324    /// [`is_cancelled`](Self::is_cancelled) returns `true` after it.
325    #[inline]
326    pub fn cancel(&self) {
327        // A swap: see `Stopper::cancel`.
328        self.inner.state.swap(CANCELLED, Ordering::Release);
329    }
330
331    /// Check if this node is cancelled (either directly or via ancestor).
332    ///
333    /// When it returns `true`, an Acquire fence follows, so this thread also
334    /// sees the writes made before whichever `cancel()` it observed (for a
335    /// parent that isn't a `ChildStopper`, if that parent's cancel is a
336    /// Release swap or store, as every stop in this crate's is). Checks are
337    /// Relaxed.
338    #[inline]
339    pub fn is_cancelled(&self) -> bool {
340        let cancelled = self.inner.should_stop();
341        if cancelled {
342            core::sync::atomic::fence(Ordering::Acquire);
343        }
344        cancelled
345    }
346}
347
348impl Default for ChildStopper {
349    fn default() -> Self {
350        Self::new()
351    }
352}
353
354impl Stop for ChildStopper {
355    #[inline]
356    #[track_caller]
357    fn check(&self) -> Result<(), StopReason> {
358        self.inner.check()
359    }
360
361    #[inline]
362    #[track_caller]
363    fn should_stop(&self) -> bool {
364        self.inner.should_stop()
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371    use crate::Stopper;
372
373    /// See `stopper::tests`: the child's checks are Relaxed; its
374    /// `is_cancelled()` fences, whichever ancestor was cancelled.
375    #[cfg(feature = "std")]
376    #[test]
377    fn is_cancelled_after_the_stop_sees_writes_made_before_cancel() {
378        use core::sync::atomic::{AtomicU32, Ordering::Relaxed};
379        for parent_is_stopper in [false, true] {
380            for _ in 0..if cfg!(miri) { 15 } else { 500 } {
381                let stopper = Stopper::new();
382                let root = if parent_is_stopper {
383                    ChildStopper::with_parent(stopper.clone())
384                } else {
385                    ChildStopper::new()
386                };
387                let leaf = root.child();
388                let data = Arc::new(AtomicU32::new(0));
389                let reader = {
390                    let (leaf, data) = (leaf.clone(), Arc::clone(&data));
391                    std::thread::spawn(move || {
392                        while !leaf.should_stop() {
393                            core::hint::spin_loop();
394                        }
395                        assert!(leaf.is_cancelled());
396                        data.load(Relaxed)
397                    })
398                };
399                data.store(42, Relaxed);
400                if parent_is_stopper {
401                    stopper.cancel();
402                } else {
403                    root.cancel();
404                }
405                assert_eq!(reader.join().unwrap(), 42);
406            }
407        }
408    }
409
410    #[cfg(feature = "std")]
411    #[test]
412    fn is_cancelled_after_a_second_cancel_sees_the_first_cancellers_writes() {
413        use core::sync::atomic::{AtomicU32, Ordering::Relaxed};
414        for _ in 0..if cfg!(miri) { 30 } else { 1_000 } {
415            let root = ChildStopper::new();
416            let leaf = root.child();
417            let data = AtomicU32::new(0);
418            std::thread::scope(|scope| {
419                let reader = scope.spawn(|| {
420                    while !leaf.should_stop() {
421                        core::hint::spin_loop();
422                    }
423                    assert!(leaf.is_cancelled());
424                    data.load(Relaxed)
425                });
426                scope.spawn(|| {
427                    while !root.should_stop() {
428                        core::hint::spin_loop();
429                    }
430                    root.cancel();
431                });
432                data.store(42, Relaxed);
433                root.cancel();
434                assert_eq!(reader.join().unwrap(), 42);
435            });
436        }
437    }
438
439    #[test]
440    fn tree_root_basic() {
441        let root = ChildStopper::new();
442        assert!(!root.is_cancelled());
443        assert!(!root.should_stop());
444        assert!(root.check().is_ok());
445
446        root.cancel();
447
448        assert!(root.is_cancelled());
449        assert!(root.should_stop());
450        assert_eq!(root.check(), Err(StopReason::Cancelled));
451    }
452
453    #[test]
454    fn tree_child_inherits_parent() {
455        let parent = ChildStopper::new();
456        let child = parent.child();
457
458        assert!(!child.is_cancelled());
459
460        parent.cancel();
461
462        assert!(child.is_cancelled());
463    }
464
465    #[test]
466    fn tree_child_cancel_independent() {
467        let parent = ChildStopper::new();
468        let child = parent.child();
469
470        child.cancel();
471
472        assert!(child.is_cancelled());
473        assert!(!parent.is_cancelled());
474    }
475
476    #[test]
477    fn tree_siblings_independent() {
478        let parent = ChildStopper::new();
479        let child_a = parent.child();
480        let child_b = parent.child();
481
482        child_a.cancel();
483
484        assert!(child_a.is_cancelled());
485        assert!(!child_b.is_cancelled());
486
487        parent.cancel();
488        assert!(child_b.is_cancelled());
489    }
490
491    #[test]
492    fn tree_grandchild() {
493        let grandparent = ChildStopper::new();
494        let parent = grandparent.child();
495        let child = parent.child();
496
497        assert!(!child.is_cancelled());
498
499        grandparent.cancel();
500        assert!(child.is_cancelled());
501    }
502
503    #[test]
504    fn tree_three_generations() {
505        let g1 = ChildStopper::new();
506        let g2 = g1.child();
507        let g3 = g2.child();
508
509        assert!(!g3.is_cancelled());
510
511        // Cancel middle generation
512        g2.cancel();
513
514        assert!(!g1.is_cancelled());
515        assert!(g2.is_cancelled());
516        assert!(g3.is_cancelled());
517    }
518
519    #[test]
520    fn tree_with_stopper_parent() {
521        let root = Stopper::new();
522        let child = ChildStopper::with_parent(root.clone());
523
524        assert!(!child.is_cancelled());
525
526        root.cancel();
527
528        assert!(child.is_cancelled());
529    }
530
531    #[test]
532    fn tree_clone_shares_state() {
533        let t1 = ChildStopper::new();
534        let t2 = t1.clone();
535
536        t2.cancel();
537
538        assert!(t1.is_cancelled());
539        assert!(t2.is_cancelled());
540    }
541
542    #[test]
543    fn tree_is_send_sync() {
544        fn assert_send_sync<T: Send + Sync>() {}
545        assert_send_sync::<ChildStopper>();
546    }
547
548    #[test]
549    fn tree_is_default() {
550        let t: ChildStopper = Default::default();
551        assert!(!t.is_cancelled());
552    }
553
554    #[test]
555    fn child_stopper_parents_are_walked_directly() {
556        let root = ChildStopper::new();
557        assert!(root.inner.parent.is_none());
558        assert_eq!(root.inner.state.load(Ordering::Relaxed), RUNNING);
559        let child = root.child();
560        assert!(Arc::ptr_eq(
561            child.inner.parent.as_ref().unwrap(),
562            &root.inner
563        ));
564        let adopted = ChildStopper::with_parent(child.clone());
565        assert!(Arc::ptr_eq(
566            adopted.inner.parent.as_ref().unwrap(),
567            &child.inner
568        ));
569        let other = ChildStopper::with_parent(Stopper::new());
570        assert!(other.inner.parent.is_none());
571        assert_eq!(other.inner.state.load(Ordering::Relaxed), ABOVE);
572        let never = ChildStopper::with_parent(crate::Unstoppable);
573        assert_eq!(never.inner.state.load(Ordering::Relaxed), RUNNING);
574    }
575
576    #[test]
577    fn cancelling_a_node_under_another_stop_overrides_it() {
578        let stopper = Stopper::new();
579        let node = ChildStopper::with_parent(stopper.clone());
580        let leaf = node.child();
581        assert!(!leaf.should_stop());
582        node.cancel();
583        assert_eq!(node.inner.state.load(Ordering::Relaxed), CANCELLED);
584        assert_eq!(leaf.check(), Err(StopReason::Cancelled));
585        assert!(!stopper.should_stop());
586    }
587
588    #[test]
589    fn a_deep_chain_stops_below_the_cancelled_node_only() {
590        let root = ChildStopper::with_parent(Stopper::new());
591        let mut chain = alloc::vec![root];
592        for _ in 0..100 {
593            let next = chain.last().unwrap().child();
594            chain.push(next);
595        }
596        assert!(chain.iter().all(|node| node.check().is_ok()));
597        chain[40].cancel();
598        for (depth, node) in chain.iter().enumerate() {
599            assert_eq!(node.should_stop(), depth >= 40, "depth {depth}");
600            assert_eq!(node.is_cancelled(), depth >= 40, "depth {depth}");
601            let expected = if depth >= 40 {
602                Err(StopReason::Cancelled)
603            } else {
604                Ok(())
605            };
606            assert_eq!(node.check(), expected, "depth {depth}");
607        }
608    }
609
610    #[test]
611    fn a_stopper_at_the_top_stops_the_whole_chain() {
612        let stopper = Stopper::new();
613        let leaf = ChildStopper::with_parent(stopper.clone()).child().child();
614        assert!(!leaf.should_stop());
615        stopper.cancel();
616        assert!(leaf.should_stop());
617        assert_eq!(leaf.check(), Err(StopReason::Cancelled));
618    }
619
620    #[test]
621    fn the_top_stops_reason_reaches_the_leaf() {
622        struct TimedOut;
623        impl Stop for TimedOut {
624            fn check(&self) -> Result<(), StopReason> {
625                Err(StopReason::TimedOut)
626            }
627        }
628        let leaf = ChildStopper::with_parent(TimedOut).child().child();
629        assert_eq!(leaf.check(), Err(StopReason::TimedOut));
630        assert!(leaf.should_stop());
631    }
632
633    #[test]
634    fn a_child_stopper_inside_a_stop_token_is_checked_above() {
635        let root = ChildStopper::new();
636        let middle = ChildStopper::with_parent(StopToken::new(root.clone()));
637        assert!(middle.inner.parent.is_none());
638        assert_eq!(middle.inner.state.load(Ordering::Relaxed), ABOVE);
639        let leaf = middle.child();
640        assert!(leaf.check().is_ok());
641        root.cancel();
642        assert_eq!(leaf.check(), Err(StopReason::Cancelled));
643        assert!(leaf.should_stop());
644    }
645
646    #[test]
647    fn a_cancelled_node_reports_cancelled_over_a_timed_out_parent() {
648        struct TimedOut;
649        impl Stop for TimedOut {
650            fn check(&self) -> Result<(), StopReason> {
651                Err(StopReason::TimedOut)
652            }
653        }
654        let node = ChildStopper::with_parent(TimedOut);
655        node.cancel();
656        assert_eq!(node.check(), Err(StopReason::Cancelled));
657        assert_eq!(node.child().check(), Err(StopReason::Cancelled));
658    }
659
660    #[test]
661    fn an_unstoppable_parent_makes_a_root() {
662        let node = ChildStopper::with_parent(crate::Unstoppable);
663        assert!(node.inner.parent.is_none());
664        assert!(!node.inner.above.may_stop());
665        assert!(alloc::format!("{node:?}").contains("parent: None"));
666        assert!(node.check().is_ok());
667        node.cancel();
668        assert_eq!(node.check(), Err(StopReason::Cancelled));
669    }
670
671    #[test]
672    fn debug_names_the_parent_kind() {
673        let root = ChildStopper::new();
674        assert!(alloc::format!("{root:?}").contains("parent: None"));
675        assert!(alloc::format!("{:?}", root.child()).contains("<ChildStopper>"));
676        let adopted = ChildStopper::with_parent(Stopper::new());
677        assert!(alloc::format!("{adopted:?}").contains("<StopToken>"));
678    }
679}
almost_enough/tree.rs

almost_enough/
tree.rs
