42 lines
945 B
Rust
42 lines
945 B
Rust
use std::{
|
|
future::Future,
|
|
pin::Pin,
|
|
sync::Mutex,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
pub(crate) struct CatchUnwindFuture<F> {
|
|
future: Mutex<F>,
|
|
}
|
|
|
|
pub(crate) fn catch_unwind<F>(future: F) -> CatchUnwindFuture<F>
|
|
where
|
|
F: Future + Unpin,
|
|
{
|
|
CatchUnwindFuture {
|
|
future: Mutex::new(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 = &self.future;
|
|
let waker = cx.waker().clone();
|
|
let res = std::panic::catch_unwind(|| {
|
|
let mut context = Context::from_waker(&waker);
|
|
let mut guard = future.lock().unwrap();
|
|
Pin::new(&mut *guard).poll(&mut context)
|
|
});
|
|
|
|
match res {
|
|
Ok(poll) => poll.map(Ok),
|
|
Err(e) => Poll::Ready(Err(e)),
|
|
}
|
|
}
|
|
}
|