1 use std::cell::{Cell, RefCell};
2 use std::future::Future;
3 use std::mem::ManuallyDrop;
4 use std::pin::Pin;
5 use std::rc::Rc;
6 use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
7 
8 struct Inner {
9     future: Pin<Box<dyn Future<Output = ()> + 'static>>,
10     waker: Waker,
11 }
12 
13 pub(crate) struct Task {
14     // The actual Future that we're executing as part of this task.
15     //
16     // This is an Option so that the Future can be immediately dropped when it's
17     // finished
18     inner: RefCell<Option<Inner>>,
19 
20     // This is used to ensure that the Task will only be queued once
21     is_queued: Cell<bool>,
22 }
23 
24 impl Task {
spawn(future: Pin<Box<dyn Future<Output = ()> + 'static>>)25     pub(crate) fn spawn(future: Pin<Box<dyn Future<Output = ()> + 'static>>) {
26         let this = Rc::new(Self {
27             inner: RefCell::new(None),
28             is_queued: Cell::new(false),
29         });
30 
31         let waker = unsafe { Waker::from_raw(Task::into_raw_waker(Rc::clone(&this))) };
32 
33         *this.inner.borrow_mut() = Some(Inner { future, waker });
34 
35         Task::wake_by_ref(&this);
36     }
37 
wake_by_ref(this: &Rc<Self>)38     fn wake_by_ref(this: &Rc<Self>) {
39         // If we've already been placed on the run queue then there's no need to
40         // requeue ourselves since we're going to run at some point in the
41         // future anyway.
42         if this.is_queued.replace(true) {
43             return;
44         }
45 
46         crate::queue::QUEUE.with(|queue| {
47             queue.push_task(Rc::clone(this));
48         });
49     }
50 
51     /// Creates a standard library `RawWaker` from an `Rc` of ourselves.
52     ///
53     /// Note that in general this is wildly unsafe because everything with
54     /// Futures requires `Sync` + `Send` with regard to Wakers. For wasm,
55     /// however, everything is guaranteed to be singlethreaded (since we're
56     /// compiled without the `atomics` feature) so we "safely lie" and say our
57     /// `Rc` pointer is good enough.
into_raw_waker(this: Rc<Self>) -> RawWaker58     unsafe fn into_raw_waker(this: Rc<Self>) -> RawWaker {
59         unsafe fn raw_clone(ptr: *const ()) -> RawWaker {
60             let ptr = ManuallyDrop::new(Rc::from_raw(ptr as *const Task));
61             Task::into_raw_waker((*ptr).clone())
62         }
63 
64         unsafe fn raw_wake(ptr: *const ()) {
65             let ptr = Rc::from_raw(ptr as *const Task);
66             Task::wake_by_ref(&ptr);
67         }
68 
69         unsafe fn raw_wake_by_ref(ptr: *const ()) {
70             let ptr = ManuallyDrop::new(Rc::from_raw(ptr as *const Task));
71             Task::wake_by_ref(&ptr);
72         }
73 
74         unsafe fn raw_drop(ptr: *const ()) {
75             drop(Rc::from_raw(ptr as *const Task));
76         }
77 
78         const VTABLE: RawWakerVTable =
79             RawWakerVTable::new(raw_clone, raw_wake, raw_wake_by_ref, raw_drop);
80 
81         RawWaker::new(Rc::into_raw(this) as *const (), &VTABLE)
82     }
83 
run(&self)84     pub(crate) fn run(&self) {
85         let mut borrow = self.inner.borrow_mut();
86 
87         // Wakeups can come in after a Future has finished and been destroyed,
88         // so handle this gracefully by just ignoring the request to run.
89         let inner = match borrow.as_mut() {
90             Some(inner) => inner,
91             None => return,
92         };
93 
94         // Ensure that if poll calls `waker.wake()` we can get enqueued back on
95         // the run queue.
96         self.is_queued.set(false);
97 
98         let poll = {
99             let mut cx = Context::from_waker(&inner.waker);
100             inner.future.as_mut().poll(&mut cx)
101         };
102 
103         // If a future has finished (`Ready`) then clean up resources associated
104         // with the future ASAP. This ensures that we don't keep anything extra
105         // alive in-memory by accident. Our own struct, `Rc<Task>` won't
106         // actually go away until all wakers referencing us go away, which may
107         // take quite some time, so ensure that the heaviest of resources are
108         // released early.
109         if let Poll::Ready(_) = poll {
110             *borrow = None;
111         }
112     }
113 }
114