dialcache/
cancel.rs

1//! Cooperative cancellation for remote reads.
2
3use std::future::Future;
4use std::panic::{catch_unwind, AssertUnwindSafe};
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll, Waker};
8
9use parking_lot::Mutex;
10
11#[derive(Default)]
12struct Inner {
13    cancelled: bool,
14    wakers: Vec<Waker>,
15    callbacks: Vec<Box<dyn FnOnce() + Send>>,
16}
17
18/// A cancellation request handed to remote adapters with each read.
19///
20/// The cache cancels the token once the read deadline passes. Cancellation is
21/// a request; it does not prove that a dispatched command stopped.
22#[derive(Clone, Default)]
23pub struct CancelToken {
24    inner: Arc<Mutex<Inner>>,
25}
26
27impl std::fmt::Debug for CancelToken {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        f.debug_struct("CancelToken")
30            .field("cancelled", &self.is_cancelled())
31            .finish()
32    }
33}
34
35impl CancelToken {
36    /// A token that has not been cancelled; clones share its state.
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Whether cancellation was requested.
42    pub fn is_cancelled(&self) -> bool {
43        self.inner.lock().cancelled
44    }
45
46    /// Request cancellation once. Later calls do nothing.
47    pub fn cancel(&self) {
48        let (wakers, callbacks) = {
49            let mut inner = self.inner.lock();
50            if inner.cancelled {
51                return;
52            }
53            inner.cancelled = true;
54            (
55                std::mem::take(&mut inner.wakers),
56                std::mem::take(&mut inner.callbacks),
57            )
58        };
59        for waker in wakers {
60            waker.wake();
61        }
62        // A callback belongs to an adapter; its panic must not become the
63        // outcome of the read that is being abandoned.
64        for callback in callbacks {
65            let _ = catch_unwind(AssertUnwindSafe(callback));
66        }
67    }
68
69    /// Run `callback` when cancellation is requested, or immediately if it already was.
70    pub fn on_cancel(&self, callback: impl FnOnce() + Send + 'static) {
71        let mut callback = Some(Box::new(callback) as Box<dyn FnOnce() + Send>);
72        {
73            let mut inner = self.inner.lock();
74            if !inner.cancelled {
75                inner
76                    .callbacks
77                    .push(callback.take().expect("callback present"));
78            }
79        }
80        if let Some(callback) = callback {
81            let _ = catch_unwind(AssertUnwindSafe(callback));
82        }
83    }
84
85    /// Completes once cancellation is requested.
86    pub fn cancelled(&self) -> Cancelled {
87        Cancelled {
88            token: self.clone(),
89        }
90    }
91}
92
93/// Future returned by [`CancelToken::cancelled`].
94#[derive(Debug)]
95pub struct Cancelled {
96    token: CancelToken,
97}
98
99impl Future for Cancelled {
100    type Output = ();
101
102    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
103        let mut inner = self.token.inner.lock();
104        if inner.cancelled {
105            return Poll::Ready(());
106        }
107        if !inner.wakers.iter().any(|w| w.will_wake(cx.waker())) {
108            inner.wakers.push(cx.waker().clone());
109        }
110        Poll::Pending
111    }
112}
dialcache/cancel.rs

dialcache/
cancel.rs
