futures_signals/
future.rs

1use std::pin::Pin;
2// TODO use parking_lot ?
3use std::sync::{Arc, Weak, Mutex};
4use std::future::Future;
5use std::task::{Poll, Waker, Context};
6// TODO use parking_lot ?
7use std::sync::atomic::{AtomicBool, Ordering};
8use discard::{Discard, DiscardOnDrop};
9use pin_project::pin_project;
10
11
12#[derive(Debug)]
13struct CancelableFutureState {
14    is_cancelled: AtomicBool,
15    waker: Mutex<Option<Waker>>,
16}
17
18
19#[derive(Debug)]
20pub struct CancelableFutureHandle {
21    state: Weak<CancelableFutureState>,
22}
23
24impl Discard for CancelableFutureHandle {
25    fn discard(self) {
26        if let Some(state) = self.state.upgrade() {
27            let mut lock = state.waker.lock().unwrap();
28
29            // TODO verify that this is correct
30            state.is_cancelled.store(true, Ordering::SeqCst);
31
32            if let Some(waker) = lock.take() {
33                drop(lock);
34                waker.wake();
35            }
36        }
37    }
38}
39
40
41#[pin_project(project = CancelableFutureProj)]
42#[derive(Debug)]
43#[must_use = "Futures do nothing unless polled"]
44pub struct CancelableFuture<A, B> {
45    state: Arc<CancelableFutureState>,
46    #[pin]
47    future: Option<A>,
48    when_cancelled: Option<B>,
49}
50
51impl<A, B> Future for CancelableFuture<A, B>
52    where A: Future,
53          B: FnOnce() -> A::Output {
54
55    type Output = A::Output;
56
57    // TODO should this inline ?
58    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
59        let CancelableFutureProj { state, mut future, when_cancelled } = self.project();
60
61        // TODO is this correct ?
62        if state.is_cancelled.load(Ordering::SeqCst) {
63            // This is necessary in order to prevent the future from calling `waker.wake()` later
64            future.set(None);
65            let callback = when_cancelled.take().unwrap();
66            // TODO figure out how to call the callback immediately when discard is called, e.g. using two Arc<Mutex<>>
67            Poll::Ready(callback())
68
69        } else {
70            match future.as_pin_mut().unwrap().poll(cx) {
71                Poll::Pending => {
72                    // TODO is this correct ?
73                    *state.waker.lock().unwrap() = Some(cx.waker().clone());
74                    Poll::Pending
75                },
76                a => a,
77            }
78        }
79    }
80}
81
82
83// TODO figure out a more efficient way to implement this
84// TODO replace with futures_util::abortable ?
85pub fn cancelable_future<A, B>(future: A, when_cancelled: B) -> (DiscardOnDrop<CancelableFutureHandle>, CancelableFuture<A, B>)
86    where A: Future,
87          B: FnOnce() -> A::Output {
88
89    let state = Arc::new(CancelableFutureState {
90        is_cancelled: AtomicBool::new(false),
91        waker: Mutex::new(None),
92    });
93
94    let cancel_handle = DiscardOnDrop::new(CancelableFutureHandle {
95        state: Arc::downgrade(&state),
96    });
97
98    let cancel_future = CancelableFuture {
99        state,
100        future: Some(future),
101        when_cancelled: Some(when_cancelled),
102    };
103
104    (cancel_handle, cancel_future)
105}
futures_signals/future.rs

futures_signals/
future.rs
