39 lines
892 B
Rust
39 lines
892 B
Rust
use std::{
|
|
future::Future,
|
|
panic::AssertUnwindSafe,
|
|
pin::Pin,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
pub(crate) struct CatchUnwindFuture<F> {
|
|
future: F,
|
|
}
|
|
|
|
pub(crate) fn catch_unwind<F>(future: F) -> CatchUnwindFuture<F>
|
|
where
|
|
F: Future + Unpin,
|
|
{
|
|
CatchUnwindFuture { future }
|
|
}
|
|
|
|
impl<F> Future for CatchUnwindFuture<F>
|
|
where
|
|
F: Future + Unpin,
|
|
{
|
|
type Output = std::thread::Result<F::Output>;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
let future = &mut self.get_mut().future;
|
|
let waker = cx.waker().clone();
|
|
let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
|
|
let mut context = Context::from_waker(&waker);
|
|
Pin::new(future).poll(&mut context)
|
|
}));
|
|
|
|
match res {
|
|
Ok(poll) => poll.map(Ok),
|
|
Err(e) => Poll::Ready(Err(e)),
|
|
}
|
|
}
|
|
}
|