JoinSet Task Supervisor

Spawn a batch of async tasks, collect results as they finish, and abort everything still running when a deadline expires.

RustResilience

Code

snippet
use std::time::Duration;
use tokio::task::JoinSet;
use tokio::time::{sleep, timeout};

async fn run_task(id: u32) -> u32 {
    sleep(Duration::from_millis(20 * id as u64)).await;
    id * id
}

async fn supervise(ids: &[u32], deadline: Duration) -> Vec<u32> {
    let mut set = JoinSet::new();
    for &id in ids {
        set.spawn(run_task(id));
    }

    let mut results = Vec::new();
    let outcome = {
        let collect = async {
            while let Some(joined) = set.join_next().await {
                match joined {
                    Ok(value) => results.push(value),
                    Err(error) => eprintln!("task failed: {error}"),
                }
            }
        };
        timeout(deadline, collect).await
    };

    if outcome.is_err() {
        eprintln!("deadline reached, aborting {} tasks", set.len());
        set.abort_all();
    }

    results.sort_unstable();
    results
}

#[tokio::main]
async fn main() {
    let finished = supervise(&[1, 2, 3], Duration::from_secs(1)).await;
    println!("finished: {finished:?}");

    let partial = supervise(&[1, 2, 30], Duration::from_millis(200)).await;
    println!("partial: {partial:?}");
}