tds_driver/client/
cancel.rs

1use futures_util::task::AtomicWaker;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::Arc;
4use std::task::Context;
5
6/// Cancels the request in flight on a [`Client`]'s connection from outside the
7/// future awaiting it: another task, or a `select!` branch beside it.
8///
9/// Unlike [`Client::cancel_query`], which can only run once the response
10/// stream has been dropped, this also stops a request that is still being sent
11/// or that the server hasn't answered yet, such as a statement still sorting,
12/// aggregating or waiting on a lock.
13///
14/// [`Client`]: crate::Client
15/// [`Client::cancel_query`]: crate::Client::cancel_query
16#[derive(Clone, Debug, Default)]
17pub struct CancelHandle(Arc<Signal>);
18
19#[derive(Debug, Default)]
20struct Signal {
21    requested: AtomicBool,
22    waker: AtomicWaker,
23}
24
25impl CancelHandle {
26    /// Asks the server to stop the request in flight on the connection.
27    ///
28    /// The Attention goes out the next time the response is read, so keep
29    /// polling the future awaiting it. It then finishes early, with whatever
30    /// the server sent before it stopped, or with server error 3980 if the
31    /// request was aborted before it ran.
32    ///
33    /// A cancel is forgotten if it comes before a request is sent, so it never
34    /// stops a later one, or if the response is dropped before the Attention
35    /// goes out; call [`Client::cancel_query`] then instead.
36    ///
37    /// [`Client::cancel_query`]: crate::Client::cancel_query
38    pub fn cancel(&self) {
39        self.0.requested.store(true, Ordering::Release);
40        self.0.waker.wake();
41    }
42
43    /// Whether a cancel is pending. Registers `cx` to be woken by the next one.
44    pub(crate) fn poll_requested(&self, cx: &mut Context<'_>) -> bool {
45        self.0.waker.register(cx.waker());
46        self.0.requested.load(Ordering::Acquire)
47    }
48
49    /// Forgets a pending cancel.
50    pub(crate) fn clear(&self) {
51        self.0.requested.store(false, Ordering::Release);
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use std::sync::atomic::AtomicUsize;
59    use std::task::{Wake, Waker};
60
61    struct CountWakes(AtomicUsize);
62
63    impl Wake for CountWakes {
64        fn wake(self: Arc<Self>) {
65            self.0.fetch_add(1, Ordering::SeqCst);
66        }
67    }
68
69    #[test]
70    fn cancel_wakes_the_reader_and_stays_pending_until_cleared() {
71        let wakes = Arc::new(CountWakes(AtomicUsize::new(0)));
72        let waker = Waker::from(Arc::clone(&wakes));
73        let mut cx = Context::from_waker(&waker);
74        let handle = CancelHandle::default();
75
76        assert!(!handle.poll_requested(&mut cx));
77        handle.clone().cancel();
78        assert_eq!(1, wakes.0.load(Ordering::SeqCst));
79        assert!(handle.poll_requested(&mut cx));
80
81        handle.clear();
82        assert!(!handle.poll_requested(&mut cx));
83    }
84}
tds_driver/client/cancel.rs

tds_driver/client/
cancel.rs
