Skip to main content

dope_runtime/executor/
join_handle.rs

1use std::{
2    cell::{Cell, RefCell},
3    future::Future,
4    pin::Pin,
5    rc::Rc,
6    task::{Context, Poll, Waker},
7};
8
9pub struct JoinHandle<T> {
10    kind: JoinHandleKind<T>,
11}
12
13enum JoinHandleKind<T> {
14    Local(Rc<LocalJoinState<T>>),
15}
16
17struct LocalJoinState<T> {
18    ready: Cell<bool>,
19    value: RefCell<Option<T>>,
20    waker: RefCell<Option<Waker>>,
21}
22
23impl<T: 'static> JoinHandle<T> {
24    pub(crate) fn new_local() -> (Self, impl FnOnce(T) + 'static) {
25        let state = Rc::new(LocalJoinState {
26            ready: Cell::new(false),
27            value: RefCell::new(None),
28            waker: RefCell::new(None),
29        });
30        let state2 = state.clone();
31        let complete = move |v: T| {
32            *state2.value.borrow_mut() = Some(v);
33            state2.ready.set(true);
34            if let Some(w) = state2.waker.borrow_mut().take() {
35                w.wake();
36            }
37        };
38
39        (
40            Self {
41                kind: JoinHandleKind::Local(state),
42            },
43            complete,
44        )
45    }
46}
47
48impl<T: 'static> Future for JoinHandle<T> {
49    type Output = T;
50
51    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
52        let me = self.get_mut();
53        match &mut me.kind {
54            JoinHandleKind::Local(state) => {
55                if let Some(v) = state.value.borrow_mut().take() {
56                    return Poll::Ready(v);
57                }
58
59                let mut slot = state.waker.borrow_mut();
60                let replace = match slot.as_ref() {
61                    Some(w) => !w.will_wake(cx.waker()),
62                    None => true,
63                };
64                if replace {
65                    *slot = Some(cx.waker().clone());
66                }
67                drop(slot);
68
69                if let Some(v) = state.value.borrow_mut().take() {
70                    return Poll::Ready(v);
71                }
72                Poll::Pending
73            }
74        }
75    }
76}