66 lines
2.1 KiB
Rust
66 lines
2.1 KiB
Rust
use std::future::Future;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
|
|
use pin_project::{pin_project, pinned_drop};
|
|
|
|
/// Trait allowing you to attach a function to a [`Future`] that will be called if the future is
|
|
/// cancelled.
|
|
pub trait OnCancelFutureExt
|
|
where
|
|
Self: Future + Sized,
|
|
{
|
|
/// Wraps the future with an [`OnCancelFutureExt`] that will execute the given function
|
|
/// when the future is cancelled.
|
|
fn on_cancel<D: FnMut()>(self, on_drop: D) -> OnCancelFuture<Self, D>;
|
|
}
|
|
impl<F: Future> OnCancelFutureExt for F {
|
|
fn on_cancel<D: FnMut()>(self, on_cancel: D) -> OnCancelFuture<Self, D> {
|
|
OnCancelFuture {
|
|
inner: self,
|
|
on_cancel,
|
|
is_ready: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Wrapper around a [`Future`] that calls an `on_cancel` callback if the future is cancelled
|
|
/// before it resolved to ready. See [`OnCancelFuture::on_cancel`] for more details. A future is
|
|
/// considered cancelled if it is dropped before resolving to [`Poll::Ready`], see <https://google.github.io/comprehensive-rust/async/pitfalls/cancellation.html#cancellation>.
|
|
#[pin_project(PinnedDrop)]
|
|
pub struct OnCancelFuture<F: Future, D: FnMut()> {
|
|
#[pin]
|
|
inner: F,
|
|
/// Function that is called when the future is cancelled.
|
|
on_cancel: D,
|
|
/// Whether the inner future is has returned [`Poll::Ready`] (indicating the future is
|
|
/// complete).
|
|
is_ready: bool,
|
|
}
|
|
|
|
impl<F: Future, D: FnMut()> Future for OnCancelFuture<F, D> {
|
|
type Output = F::Output;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
|
|
let this = self.project();
|
|
let output = this.inner.poll(cx);
|
|
*this.is_ready = output.is_ready();
|
|
output
|
|
}
|
|
}
|
|
|
|
#[pinned_drop]
|
|
impl<F: Future, D: FnMut()> PinnedDrop for OnCancelFuture<F, D> {
|
|
fn drop(self: Pin<&mut Self>) {
|
|
// If the future was dropped before it was resolved to ready, the future was cancelled.
|
|
let this = self.project();
|
|
if !*this.is_ready {
|
|
(this.on_cancel)();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "on_cancel_tests.rs"]
|
|
mod tests;
|