Structured Concurrency — The Pattern That Defines Senior Tokio

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

pub struct CancelToken {
    // Flag for this node.
    flag: Arc<AtomicBool>,
    // Parent chain — checked on is_cancelled() to support "cancel
    // bubbles down from ancestor". Children hold references to all
    // ancestor flags (cheap: just Arc clones).
    ancestors: Vec<Arc<AtomicBool>>,
}

impl CancelToken {
    pub fn new() -> Self {
        Self {
            flag: Arc::new(AtomicBool::new(false)),
            ancestors: Vec::new(),
        }
    }

    /// Shares the *same* node — not a child. Use this to hand a copy
    /// of this token to another thread without splitting the subtree.
    pub fn clone_handle(&self) -> Self {
        // TODO: hint — share `flag` and `ancestors` via Arc::clone
        todo!()
    }

    /// Create a child token. Cancelling *this* child does not cancel
    /// the parent. Cancelling the parent cancels this child.
    pub fn child(&self) -> Self {
        // TODO: new flag for the child; ancestors = self.ancestors ++ [self.flag]
        todo!()
    }

    pub fn cancel(&self) {
        self.flag.store(true, Ordering::SeqCst);
    }

    pub fn is_cancelled(&self) -> bool {
        // TODO: true iff self.flag OR any ancestor flag is set
        todo!()
    }
}

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

// Re-declare the pieces of CancelToken we need. (In the capstone you'll
// import from the drill above; here we duplicate so this drill compiles
// standalone.)
use std::sync::atomic::AtomicBool;

#[derive(Clone)]
pub struct Cancel {
    flag: Arc<AtomicBool>,
}

impl Cancel {
    pub fn new() -> Self { Self { flag: Arc::new(AtomicBool::new(false)) } }
    pub fn cancel(&self) { self.flag.store(true, Ordering::SeqCst); }
    pub fn is_cancelled(&self) -> bool { self.flag.load(Ordering::SeqCst) }
}

pub struct TaskSet {
    handles: Vec<JoinHandle<()>>,
    cancel: Cancel,
    /// Incremented every time a spawned task returns. Lets tests assert
    /// that all tasks actually ran (and didn't panic silently).
    completed: Arc<AtomicUsize>,
}

impl TaskSet {
    pub fn new(cancel: Cancel) -> Self {
        Self {
            handles: Vec::new(),
            cancel,
            completed: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Spawn a worker. The closure receives the shared Cancel handle so
    /// it can observe shutdown.
    pub fn spawn<F>(&mut self, f: F)
    where
        F: FnOnce(Cancel) + Send + 'static,
    {
        // TODO: wrap f in a closure that runs it, then increments
        // `completed`. Push the handle onto self.handles.
        let _ = f;
        todo!()
    }

    pub fn completed_count(&self) -> usize {
        self.completed.load(Ordering::SeqCst)
    }

    /// Cancel + join every task. Returns the time the drain took.
    pub fn shutdown(self) -> Duration {
        // TODO:
        //   1. Record a start Instant
        //   2. self.cancel.cancel()
        //   3. For each handle, handle.join()
        //   4. Return elapsed
        todo!()
    }
}