| // Copyright 2022 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| use { |
| futures::{ |
| future::Future, |
| task::{Context, Poll}, |
| }, |
| log::warn, |
| pin_project::pin_project, |
| std::pin::Pin, |
| }; |
| |
| /// An extension trait that allows futures to be cancelled. Cancellation is signalled |
| /// by a second future. |
| pub trait OrCancel<F: Future, C: Future> { |
| /// Create a Future, which polls both the original future and |cancel_fut|. |
| /// If |cancel_fut| resolves before the original future, Err(Cancelled) is returned. |
| /// Otherwise, the result of the original future is returned. |
| fn or_cancelled(self, cancel_fut: C) -> OrCancelledFuture<F, C>; |
| } |
| |
| /// An error indicating that a future was cancelled. |
| #[derive(PartialEq, Debug)] |
| pub struct Cancelled; |
| |
| impl<C: Future, F: Future> OrCancel<F, C> for F { |
| fn or_cancelled(self, cancel_fut: C) -> OrCancelledFuture<F, C> { |
| OrCancelledFuture { fut: self, cancel_fut } |
| } |
| } |
| |
| /// Implementation for the future returned by OrCancel::or_cancelled. |
| #[pin_project] |
| pub struct OrCancelledFuture<F: Future, C: Future> { |
| #[pin] |
| fut: F, |
| #[pin] |
| cancel_fut: C, |
| } |
| |
| impl<F: Future, C: Future> Future for OrCancelledFuture<F, C> { |
| type Output = Result<<F as Future>::Output, Cancelled>; |
| fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| let this = self.project(); |
| match this.cancel_fut.poll(cx) { |
| Poll::Ready(_) => return Poll::Ready(Err(Cancelled)), |
| Poll::Pending => (), |
| } |
| match this.fut.poll(cx) { |
| Poll::Ready(ready_result) => Poll::Ready(Ok(ready_result)), |
| Poll::Pending => Poll::Pending, |
| } |
| } |
| } |
| |
| // An extension trait that allows naming futures, and exports diagnostics when |
| // the future is dropped without polling to completion. |
| pub trait NamedFutureExt<F: Future> { |
| /// Produces a named Future. When the Future is dropped, logs a warning if |
| /// the Future was not polled to completion. |
| fn named(self, name: &'static str) -> WarnOnIncompleteFuture<F>; |
| } |
| |
| impl<F: Future> NamedFutureExt<F> for F { |
| fn named(self, name: &'static str) -> WarnOnIncompleteFuture<F> { |
| WarnOnIncompleteFuture::new(self, name, LogWarn) |
| } |
| } |
| |
| /// Implementation of the Future returned from |
| /// |NamedFutureExt::named|. |
| pub type WarnOnIncompleteFuture<F> = HookOnIncompleteFuture<F, LogWarn>; |
| |
| /// A future that calls some hook when it is dropped before polling to |
| /// completion. |
| #[pin_project] |
| pub struct HookOnIncompleteFuture<F: Future, LF: OnIncompleteHook> { |
| #[pin] |
| fut: F, |
| inner: HookOnIncompleteFutureInner<LF>, |
| } |
| |
| impl<F: Future, LF: OnIncompleteHook> HookOnIncompleteFuture<F, LF> { |
| fn new(fut: F, name: &'static str, on_incomplete: LF) -> Self { |
| HookOnIncompleteFuture { |
| fut, |
| inner: HookOnIncompleteFutureInner { name, polled_to_completion: false, on_incomplete }, |
| } |
| } |
| } |
| |
| impl<F: Future, LF: OnIncompleteHook> Future for HookOnIncompleteFuture<F, LF> { |
| type Output = <F as Future>::Output; |
| fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| let this = self.project(); |
| match this.fut.poll(cx) { |
| Poll::Ready(ready_result) => { |
| this.inner.polled_to_completion = true; |
| Poll::Ready(ready_result) |
| } |
| Poll::Pending => Poll::Pending, |
| } |
| } |
| } |
| |
| /// Inner state for |HookOnIncompleteFuture|. |
| struct HookOnIncompleteFutureInner<LF: OnIncompleteHook> { |
| name: &'static str, |
| polled_to_completion: bool, |
| on_incomplete: LF, |
| } |
| |
| impl<LF: OnIncompleteHook> std::ops::Drop for HookOnIncompleteFutureInner<LF> { |
| fn drop(&mut self) { |
| if !self.polled_to_completion { |
| self.on_incomplete.on_incomplete(self.name); |
| } |
| } |
| } |
| |
| /// Trait providing implementations for functions called when |
| /// |HookOnIncompleteFutureInner| is dropped. This exists purely to enable |
| /// testing |HookOnIncompleteFuture|. |
| pub trait OnIncompleteHook { |
| fn on_incomplete(&self, name: &str); |
| } |
| |
| /// |OnIncompleteHook| implementation that logs a warning when a future is |
| /// incomplete. |
| pub struct LogWarn; |
| |
| impl OnIncompleteHook for LogWarn { |
| fn on_incomplete(&self, name: &str) { |
| warn!("Future {} dropped before polling to completion", name); |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use super::*; |
| use futures::{ |
| channel::oneshot, |
| future::{pending, ready, FutureExt}, |
| }; |
| use std::cell::RefCell; |
| |
| #[fuchsia::test] |
| async fn stop_if_cancelled() { |
| assert!(pending::<()>().or_cancelled(ready(())).await.is_err()); |
| |
| let (_send, recv) = oneshot::channel::<()>(); |
| |
| assert!(recv.or_cancelled(ready(())).await.is_err()); |
| } |
| |
| #[fuchsia::test] |
| async fn resolve_if_not_cancelled() { |
| assert!(ready(()).or_cancelled(pending::<()>()).await.is_ok()); |
| |
| let (send, recv) = oneshot::channel::<()>(); |
| |
| futures::future::join( |
| async move { |
| send.send(()).unwrap(); |
| }, |
| async move { |
| assert_eq!(recv.or_cancelled(pending::<()>()).await.unwrap(), Ok(())); |
| }, |
| ) |
| .await; |
| } |
| |
| struct RecordOnDrop<'a>(&'a RefCell<bool>); |
| |
| impl OnIncompleteHook for RecordOnDrop<'_> { |
| fn on_incomplete(&self, _: &str) { |
| self.0.replace(true); |
| } |
| } |
| |
| #[fuchsia::test] |
| async fn no_record_when_polled_to_completion() { |
| let on_drop_called = RefCell::new(false); |
| let record_on_drop = RecordOnDrop(&on_drop_called); |
| |
| let fut = HookOnIncompleteFuture::new(ready(()), "my_fut", record_on_drop); |
| fut.await; |
| assert_eq!(*on_drop_called.borrow(), false); |
| } |
| |
| #[fuchsia::test] |
| async fn record_when_poll_incomplete() { |
| let on_drop_called = RefCell::new(false); |
| let record_on_drop = RecordOnDrop(&on_drop_called); |
| |
| let fut = HookOnIncompleteFuture::new(pending::<()>(), "my_fut", record_on_drop); |
| assert!(fut.now_or_never().is_none()); |
| assert_eq!(*on_drop_called.borrow(), true); |
| } |
| } |