1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
|
//! Cancellable futures.
use std::{
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll, Waker},
};
use futures::{future::FusedFuture, Future};
use pin_project::pin_project;
/// A cancellable future type, loosely influenced by `RemoteHandle`.
///
/// This type is useful for cases when we can't cancel a future simply by
/// dropping it, because the future is owned by some other object (like a
/// `FuturesUnordered`) that won't give it up.
//
// We could use `tokio_util`'s cancellable futures instead here, but I don't
// think we want an unconditional tokio_util dependency.
#[pin_project]
pub(crate) struct Cancel<F> {
/// Shared state between the `Cancel` and the `CancelHandle`.
//
// It would be nice not to have to stick this behind a mutex, but that would
// make it a bit tricky to manage the Waker.
inner: Arc<Mutex<Inner>>,
/// The inner future.
#[pin]
fut: F,
}
/// Inner state shared between `Cancel` and the `CancelHandle.
struct Inner {
/// True if this future has been cancelled.
cancelled: bool,
/// A waker to use in telling this future that it's cancelled.
waker: Option<Waker>,
}
/// An object that can be used to cancel a future.
#[derive(Clone)]
pub(crate) struct CancelHandle {
/// The shared state for the cancellable future between `Cancel` and
/// `CancelHandle`.
inner: Arc<Mutex<Inner>>,
}
impl<F> Cancel<F> {
/// Wrap `fut` in a new future that can be cancelled.
///
/// Returns a handle to cancel the future, and the cancellable future.
pub(crate) fn new(fut: F) -> (CancelHandle, Cancel<F>) {
let inner = Arc::new(Mutex::new(Inner {
cancelled: false,
waker: None,
}));
let handle = CancelHandle {
inner: inner.clone(),
};
let future = Cancel { inner, fut };
(handle, future)
}
}
impl CancelHandle {
/// Cancel the associated future, if it has not already finished.
#[allow(dead_code)] // TODO RPC
pub(crate) fn cancel(&self) {
let mut inner = self.inner.lock().expect("poisoned lock");
inner.cancelled = true;
if let Some(waker) = inner.waker.take() {
waker.wake();
}
}
}
/// An error returned from a `Cancel` future if it is cancelled.
#[derive(thiserror::Error, Clone, Debug)]
#[error("Future was cancelled")]
pub(crate) struct Cancelled;
impl<F: Future> Future for Cancel<F> {
type Output = Result<F::Output, Cancelled>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
{
let mut inner = self.inner.lock().expect("lock poisoned");
if inner.cancelled {
return Poll::Ready(Err(Cancelled));
}
inner.waker = Some(cx.waker().clone());
}
let this = self.project();
this.fut.poll(cx).map(Ok)
}
}
impl<F: FusedFuture> FusedFuture for Cancel<F> {
fn is_terminated(&self) -> bool {
{
let inner = self.inner.lock().expect("lock poisoned");
if inner.cancelled {
return true;
}
}
self.fut.is_terminated()
}
}
#[cfg(test)]
mod test {
// @@ begin test lint list maintained by maint/add_warning @@
#![allow(clippy::bool_assert_comparison)]
#![allow(clippy::clone_on_copy)]
#![allow(clippy::dbg_macro)]
#![allow(clippy::print_stderr)]
#![allow(clippy::print_stdout)]
#![allow(clippy::single_char_pattern)]
#![allow(clippy::unwrap_used)]
#![allow(clippy::unchecked_duration_subtraction)]
#![allow(clippy::useless_vec)]
#![allow(clippy::needless_pass_by_value)]
//! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
use super::*;
use futures_await_test::async_test;
#[async_test]
async fn not_cancelled() {
let f = futures::future::ready("hello");
let (_h, f) = Cancel::new(f);
assert_eq!(f.await.unwrap(), "hello");
}
#[async_test]
async fn cancelled() {
let f = futures::future::pending::<()>();
let (h, f) = Cancel::new(f);
let (r, ()) = futures::join!(f, async {
h.cancel();
});
assert!(matches!(r, Err(Cancelled)));
let (_tx, rx) = futures::channel::oneshot::channel::<()>();
let (h, f) = Cancel::new(rx);
let (r, ()) = futures::join!(f, async {
h.cancel();
});
assert!(matches!(r, Err(Cancelled)));
}
}
|
