use std::future::{Future, poll_fn};
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::task::Poll;
pub(crate) async fn guarded<F: Future>(fut: F) -> Result<F::Output, String> {
let mut fut: Pin<Box<F>> = Box::pin(fut);
poll_fn(
move |cx| match std::panic::catch_unwind(AssertUnwindSafe(|| fut.as_mut().poll(cx))) {
Ok(Poll::Ready(out)) => Poll::Ready(Ok(out)),
Ok(Poll::Pending) => Poll::Pending,
Err(payload) => Poll::Ready(Err(panic_message(&*payload))),
},
)
.await
}
fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
payload
.downcast_ref::<&'static str>()
.map(|s| (*s).to_owned())
.or_else(|| payload.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "non-string panic payload".to_owned())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn ok_output_passes_through() {
assert_eq!(guarded(async { 42 }).await, Ok(42));
}
#[tokio::test]
async fn completes_normally_after_await() {
let r = guarded(async {
tokio::task::yield_now().await;
"done"
})
.await;
assert_eq!(r, Ok("done"));
}
#[tokio::test]
async fn panic_before_await_becomes_err() {
let r: Result<(), String> = guarded(async { panic!("boom") }).await;
assert!(r.clone().is_err(), "panic must become Err, got {r:?}");
assert!(r.unwrap_err().contains("boom"));
}
#[tokio::test]
async fn panic_after_await_becomes_err() {
let r: Result<(), String> = guarded(async {
tokio::task::yield_now().await;
panic!("late {}", 7);
})
.await;
assert!(r.unwrap_err().contains("late 7"));
}
}