Tokio Task Join Panic Recovery

use tokio::task::JoinHandle;

async fn risky_compute(index: usize) -> usize {
    if index == 2 {
        panic!("worker {index} hit unrecoverable state");
    }
    index * 100
}

async fn supervised_batch() -> (Vec<usize>, Vec<String>) {
    let handles: Vec<JoinHandle<usize>> = (0..5)
        .map(|i| tokio::spawn(risky_compute(i)))
        .collect();

    let mut results = Vec::new();
    let mut failures = Vec::new();

    for (position, handle) in handles.into_iter().enumerate() {
        match handle.await {
            Ok(value) => results.push(value),
            Err(join_err) if join_err.is_panic() => {
                let message = join_err
                    .into_panic()
                    .downcast_ref::<String>()
                    .cloned()
                    .unwrap_or_else(|| "unknown panic payload".into());
                failures.push(format!("slot {position}: {message}"));
            }
            Err(join_err) => {
                failures.push(format!("slot {position}: {join_err}"));
            }
        }
    }
    (results, failures)
}

struct RestartingWorker {
    max_restarts: usize,
}

impl RestartingWorker {
    async fn run_with_restart<F, Fut>(&self, mut job: F)
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = usize>,
    {
        let mut restarts = 0;
        loop {
            let handle = tokio::spawn(job());
            match handle.await {
                Ok(count) => {
                    println!("worker completed {count} items");
                    return;
                }
                Err(err) if err.is_panic() => {
                    restarts += 1;
                    if restarts > self.max_restarts {
                        println!("giving up after {restarts} restarts");
                        return;
                    }
                    println!("worker panicked, restart {restarts}");
                }
                Err(_) => return,
            }
        }
    }
}

#[tokio::main]
async fn main() {
    let (results, failures) = supervised_batch().await;
    println!("results: {results:?}");
    println!("failures: {failures:?}");

    let worker = RestartingWorker { max_restarts: 2 };
    let mut invocations = 0;
    worker
        .run_with_restart(move || {
            invocations += 1;
            let attempt = invocations;
            async move {
                if attempt < 3 {
                    panic!("transient failure {attempt}");
                }
                42
            }
        })
        .await;
}