run_test_suite_lib/
cancel.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use futures::future::Future;
6use futures::task::{Context, Poll};
7use log::warn;
8use pin_project::pin_project;
9use std::pin::Pin;
10
11/// An extension trait that allows futures to be cancelled. Cancellation is signalled
12/// by a second future.
13pub trait OrCancel<F: Future, C: Future> {
14    /// Create a Future, which polls both the original future and |cancel_fut|.
15    /// If |cancel_fut| resolves before the original future, Err(Cancelled) is returned.
16    /// Otherwise, the result of the original future is returned.
17    fn or_cancelled(self, cancel_fut: C) -> OrCancelledFuture<F, C>;
18}
19
20/// An error indicating that a future was cancelled.
21#[derive(PartialEq, Debug)]
22pub struct Cancelled<T>(pub T);
23
24impl<C: Future, F: Future> OrCancel<F, C> for F {
25    fn or_cancelled(self, cancel_fut: C) -> OrCancelledFuture<F, C> {
26        OrCancelledFuture { fut: self, cancel_fut }
27    }
28}
29
30/// Implementation for the future returned by OrCancel::or_cancelled.
31#[pin_project]
32pub struct OrCancelledFuture<F: Future, C: Future> {
33    #[pin]
34    fut: F,
35    #[pin]
36    cancel_fut: C,
37}
38
39impl<F: Future, C: Future> Future for OrCancelledFuture<F, C> {
40    type Output = Result<<F as Future>::Output, Cancelled<<C as Future>::Output>>;
41    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
42        let this = self.project();
43        match this.cancel_fut.poll(cx) {
44            Poll::Ready(ready) => return Poll::Ready(Err(Cancelled(ready))),
45            Poll::Pending => (),
46        }
47        match this.fut.poll(cx) {
48            Poll::Ready(ready_result) => Poll::Ready(Ok(ready_result)),
49            Poll::Pending => Poll::Pending,
50        }
51    }
52}
53
54// An extension trait that allows naming futures, and exports diagnostics when
55// the future is dropped without polling to completion.
56pub trait NamedFutureExt<F: Future> {
57    /// Produces a named Future. When the Future is dropped, logs a warning if
58    /// the Future was not polled to completion.
59    fn named(self, name: &'static str) -> WarnOnIncompleteFuture<F>;
60}
61
62impl<F: Future> NamedFutureExt<F> for F {
63    fn named(self, name: &'static str) -> WarnOnIncompleteFuture<F> {
64        WarnOnIncompleteFuture::new(self, name, LogWarn)
65    }
66}
67
68/// Implementation of the Future returned from
69/// |NamedFutureExt::named|.
70pub type WarnOnIncompleteFuture<F> = HookOnIncompleteFuture<F, LogWarn>;
71
72/// A future that calls some hook when it is dropped before polling to
73/// completion.
74#[pin_project]
75pub struct HookOnIncompleteFuture<F: Future, LF: OnIncompleteHook> {
76    #[pin]
77    fut: F,
78    inner: HookOnIncompleteFutureInner<LF>,
79}
80
81impl<F: Future, LF: OnIncompleteHook> HookOnIncompleteFuture<F, LF> {
82    fn new(fut: F, name: &'static str, on_incomplete: LF) -> Self {
83        HookOnIncompleteFuture {
84            fut,
85            inner: HookOnIncompleteFutureInner { name, polled_to_completion: false, on_incomplete },
86        }
87    }
88}
89
90impl<F: Future, LF: OnIncompleteHook> Future for HookOnIncompleteFuture<F, LF> {
91    type Output = <F as Future>::Output;
92    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
93        let this = self.project();
94        match this.fut.poll(cx) {
95            Poll::Ready(ready_result) => {
96                this.inner.polled_to_completion = true;
97                Poll::Ready(ready_result)
98            }
99            Poll::Pending => Poll::Pending,
100        }
101    }
102}
103
104/// Inner state for |HookOnIncompleteFuture|.
105struct HookOnIncompleteFutureInner<LF: OnIncompleteHook> {
106    name: &'static str,
107    polled_to_completion: bool,
108    on_incomplete: LF,
109}
110
111impl<LF: OnIncompleteHook> std::ops::Drop for HookOnIncompleteFutureInner<LF> {
112    fn drop(&mut self) {
113        if !self.polled_to_completion {
114            self.on_incomplete.on_incomplete(self.name);
115        }
116    }
117}
118
119/// Trait providing implementations for functions called when
120/// |HookOnIncompleteFutureInner| is dropped. This exists purely to enable
121/// testing |HookOnIncompleteFuture|.
122pub trait OnIncompleteHook {
123    fn on_incomplete(&self, name: &str);
124}
125
126/// |OnIncompleteHook| implementation that logs a warning when a future is
127/// incomplete.
128pub struct LogWarn;
129
130impl OnIncompleteHook for LogWarn {
131    fn on_incomplete(&self, name: &str) {
132        warn!("Future {} dropped before polling to completion", name);
133    }
134}
135
136#[cfg(test)]
137mod test {
138    use super::*;
139    use futures::channel::oneshot;
140    use futures::future::{pending, ready, FutureExt};
141    use std::cell::RefCell;
142
143    #[fuchsia::test]
144    async fn stop_if_cancelled() {
145        assert!(pending::<()>().or_cancelled(ready(())).await.is_err());
146
147        let (_send, recv) = oneshot::channel::<()>();
148
149        assert!(recv.or_cancelled(ready(())).await.is_err());
150    }
151
152    #[fuchsia::test]
153    async fn resolve_if_not_cancelled() {
154        assert!(ready(()).or_cancelled(pending::<()>()).await.is_ok());
155
156        let (send, recv) = oneshot::channel::<()>();
157
158        futures::future::join(
159            async move {
160                send.send(()).unwrap();
161            },
162            async move {
163                assert_eq!(recv.or_cancelled(pending::<()>()).await.unwrap(), Ok(()));
164            },
165        )
166        .await;
167    }
168
169    struct RecordOnDrop<'a>(&'a RefCell<bool>);
170
171    impl OnIncompleteHook for RecordOnDrop<'_> {
172        fn on_incomplete(&self, _: &str) {
173            self.0.replace(true);
174        }
175    }
176
177    #[fuchsia::test]
178    async fn no_record_when_polled_to_completion() {
179        let on_drop_called = RefCell::new(false);
180        let record_on_drop = RecordOnDrop(&on_drop_called);
181
182        let fut = HookOnIncompleteFuture::new(ready(()), "my_fut", record_on_drop);
183        fut.await;
184        assert_eq!(*on_drop_called.borrow(), false);
185    }
186
187    #[fuchsia::test]
188    async fn record_when_poll_incomplete() {
189        let on_drop_called = RefCell::new(false);
190        let record_on_drop = RecordOnDrop(&on_drop_called);
191
192        let fut = HookOnIncompleteFuture::new(pending::<()>(), "my_fut", record_on_drop);
193        assert!(fut.now_or_never().is_none());
194        assert_eq!(*on_drop_called.borrow(), true);
195    }
196}
run_test_suite_lib/cancel.rs

run_test_suite_lib/
cancel.rs
