]>
Commit | Line | Data |
---|---|---|
5869c6ff | 1 | #![stable(feature = "wake_trait", since = "1.51.0")] |
ba9703b0 | 2 | //! Types and Traits for working with asynchronous tasks. |
f9f354fc | 3 | use core::mem::ManuallyDrop; |
ba9703b0 XL |
4 | use core::task::{RawWaker, RawWakerVTable, Waker}; |
5 | ||
6 | use crate::sync::Arc; | |
7 | ||
8 | /// The implementation of waking a task on an executor. | |
9 | /// | |
10 | /// This trait can be used to create a [`Waker`]. An executor can define an | |
11 | /// implementation of this trait, and use that to construct a Waker to pass | |
12 | /// to the tasks that are executed on that executor. | |
13 | /// | |
14 | /// This trait is a memory-safe and ergonomic alternative to constructing a | |
15 | /// [`RawWaker`]. It supports the common executor design in which the data used | |
3dfed10e | 16 | /// to wake up a task is stored in an [`Arc`]. Some executors (especially |
ba9703b0 XL |
17 | /// those for embedded systems) cannot use this API, which is why [`RawWaker`] |
18 | /// exists as an alternative for those systems. | |
5869c6ff XL |
19 | /// |
20 | /// [arc]: ../../std/sync/struct.Arc.html | |
21 | /// | |
22 | /// # Examples | |
23 | /// | |
24 | /// A basic `block_on` function that takes a future and runs it to completion on | |
25 | /// the current thread. | |
26 | /// | |
27 | /// **Note:** This example trades correctness for simplicity. In order to prevent | |
28 | /// deadlocks, production-grade implementations will also need to handle | |
29 | /// intermediate calls to `thread::unpark` as well as nested invocations. | |
30 | /// | |
31 | /// ```rust | |
32 | /// use std::future::Future; | |
33 | /// use std::sync::Arc; | |
34 | /// use std::task::{Context, Poll, Wake}; | |
35 | /// use std::thread::{self, Thread}; | |
36 | /// | |
37 | /// /// A waker that wakes up the current thread when called. | |
38 | /// struct ThreadWaker(Thread); | |
39 | /// | |
40 | /// impl Wake for ThreadWaker { | |
41 | /// fn wake(self: Arc<Self>) { | |
42 | /// self.0.unpark(); | |
43 | /// } | |
44 | /// } | |
45 | /// | |
46 | /// /// Run a future to completion on the current thread. | |
47 | /// fn block_on<T>(fut: impl Future<Output = T>) -> T { | |
48 | /// // Pin the future so it can be polled. | |
49 | /// let mut fut = Box::pin(fut); | |
50 | /// | |
51 | /// // Create a new context to be passed to the future. | |
52 | /// let t = thread::current(); | |
53 | /// let waker = Arc::new(ThreadWaker(t)).into(); | |
54 | /// let mut cx = Context::from_waker(&waker); | |
55 | /// | |
56 | /// // Run the future to completion. | |
57 | /// loop { | |
58 | /// match fut.as_mut().poll(&mut cx) { | |
59 | /// Poll::Ready(res) => return res, | |
60 | /// Poll::Pending => thread::park(), | |
61 | /// } | |
62 | /// } | |
63 | /// } | |
64 | /// | |
65 | /// block_on(async { | |
66 | /// println!("Hi from inside a future!"); | |
67 | /// }); | |
68 | /// ``` | |
69 | #[stable(feature = "wake_trait", since = "1.51.0")] | |
ba9703b0 XL |
70 | pub trait Wake { |
71 | /// Wake this task. | |
5869c6ff | 72 | #[stable(feature = "wake_trait", since = "1.51.0")] |
ba9703b0 XL |
73 | fn wake(self: Arc<Self>); |
74 | ||
75 | /// Wake this task without consuming the waker. | |
76 | /// | |
77 | /// If an executor supports a cheaper way to wake without consuming the | |
78 | /// waker, it should override this method. By default, it clones the | |
5869c6ff XL |
79 | /// [`Arc`] and calls [`wake`] on the clone. |
80 | /// | |
81 | /// [`wake`]: Wake::wake | |
82 | #[stable(feature = "wake_trait", since = "1.51.0")] | |
ba9703b0 XL |
83 | fn wake_by_ref(self: &Arc<Self>) { |
84 | self.clone().wake(); | |
85 | } | |
86 | } | |
87 | ||
5869c6ff | 88 | #[stable(feature = "wake_trait", since = "1.51.0")] |
ba9703b0 XL |
89 | impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker { |
90 | fn from(waker: Arc<W>) -> Waker { | |
91 | // SAFETY: This is safe because raw_waker safely constructs | |
92 | // a RawWaker from Arc<W>. | |
93 | unsafe { Waker::from_raw(raw_waker(waker)) } | |
94 | } | |
95 | } | |
96 | ||
5869c6ff | 97 | #[stable(feature = "wake_trait", since = "1.51.0")] |
ba9703b0 XL |
98 | impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker { |
99 | fn from(waker: Arc<W>) -> RawWaker { | |
100 | raw_waker(waker) | |
101 | } | |
102 | } | |
103 | ||
104 | // NB: This private function for constructing a RawWaker is used, rather than | |
105 | // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that | |
106 | // the safety of `From<Arc<W>> for Waker` does not depend on the correct | |
107 | // trait dispatch - instead both impls call this function directly and | |
108 | // explicitly. | |
109 | #[inline(always)] | |
110 | fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker { | |
111 | // Increment the reference count of the arc to clone it. | |
112 | unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker { | |
5869c6ff | 113 | unsafe { Arc::increment_strong_count(waker as *const W) }; |
f9f354fc XL |
114 | RawWaker::new( |
115 | waker as *const (), | |
116 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), | |
117 | ) | |
ba9703b0 XL |
118 | } |
119 | ||
120 | // Wake by value, moving the Arc into the Wake::wake function | |
121 | unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) { | |
3dfed10e | 122 | let waker = unsafe { Arc::from_raw(waker as *const W) }; |
ba9703b0 XL |
123 | <W as Wake>::wake(waker); |
124 | } | |
125 | ||
126 | // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it | |
127 | unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) { | |
3dfed10e | 128 | let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) }; |
ba9703b0 XL |
129 | <W as Wake>::wake_by_ref(&waker); |
130 | } | |
131 | ||
132 | // Decrement the reference count of the Arc on drop | |
133 | unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) { | |
5869c6ff | 134 | unsafe { Arc::decrement_strong_count(waker as *const W) }; |
ba9703b0 XL |
135 | } |
136 | ||
137 | RawWaker::new( | |
138 | Arc::into_raw(waker) as *const (), | |
139 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), | |
140 | ) | |
141 | } |