A future adapter that retries the future if it panics

retry_on_panic.rs

use std::{
    future::Future,
    panic::{catch_unwind, UnwindSafe},
    pin::Pin,
    task::{Context, Poll},
};

use pin_project::pin_project;

#[pin_project]
#[derive(Debug, Clone)]
pub struct RetryOnPanic<F> {
    original: F,

    #[pin]
    attempt: F,
}

pub fn retry_on_panic<F>(fut: F) -> RetryOnPanic<F>
where
    F: Future + Clone,
    for<'a> &'a mut F: UnwindSafe,
{
    RetryOnPanic {
        attempt: fut.clone(),
        original: fut,
    }
}
impl<F: Future + Clone> Future for RetryOnPanic<F>
where
    for<'a> &'a mut F: UnwindSafe,
{
    type Output = F::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let waker = cx.waker();

        let attempt = self.as_mut().project().attempt;

        match catch_unwind(|| {
            let mut cx = Context::from_waker(waker);
            attempt.poll(&mut cx)
        }) {
            Ok(polled) => polled,
            Err(..) => {
                let retry = self.as_mut().project().original.clone();
                self.as_mut().project().attempt.set(retry);
                waker.wake_by_ref();
                Poll::Pending
            }
        }
    }
}