Cancellation Token Budget

use std::time::Duration;
use tokio::time::sleep;
use tokio_util::sync::CancellationToken;

async fn worker(id: u32, token: CancellationToken) {
    tokio::select! {
        _ = token.cancelled() => {
            println!("worker {id} cancelled");
        }
        _ = sleep(Duration::from_millis(40 * id as u64)) => {
            println!("worker {id} finished");
        }
    }
}

async fn run_with_budget(budget: Duration) -> u32 {
    let token = CancellationToken::new();
    let mut handles = Vec::new();
    for id in 1..=4 {
        let child = token.child_token();
        handles.push(tokio::spawn(worker(id, child)));
    }

    let timer = {
        let token = token.clone();
        tokio::spawn(async move {
            sleep(budget).await;
            token.cancel();
        })
    };

    let mut returned = 0;
    for handle in handles {
        if handle.await.is_ok() {
            returned += 1;
        }
    }
    let _ = timer.await;
    returned
}

#[tokio::main]
async fn main() {
    let returned = run_with_budget(Duration::from_millis(100)).await;
    println!("workers returned: {returned}");
}