1 | #![stable (feature = "wake_trait" , since = "1.51.0" )] |
2 | |
3 | //! Types and Traits for working with asynchronous tasks. |
4 | //! |
5 | //! **Note**: Some of the types in this module are only available |
6 | //! on platforms that support atomic loads and stores of pointers. |
7 | //! This may be detected at compile time using |
8 | //! `#[cfg(target_has_atomic = "ptr")]`. |
9 | |
10 | use crate::rc::Rc; |
11 | use core::mem::ManuallyDrop; |
12 | use core::task::{LocalWaker, RawWaker, RawWakerVTable}; |
13 | |
14 | #[cfg (target_has_atomic = "ptr" )] |
15 | use crate::sync::Arc; |
16 | #[cfg (target_has_atomic = "ptr" )] |
17 | use core::task::Waker; |
18 | |
19 | /// The implementation of waking a task on an executor. |
20 | /// |
21 | /// This trait can be used to create a [`Waker`]. An executor can define an |
22 | /// implementation of this trait, and use that to construct a [`Waker`] to pass |
23 | /// to the tasks that are executed on that executor. |
24 | /// |
25 | /// This trait is a memory-safe and ergonomic alternative to constructing a |
26 | /// [`RawWaker`]. It supports the common executor design in which the data used |
27 | /// to wake up a task is stored in an [`Arc`]. Some executors (especially |
28 | /// those for embedded systems) cannot use this API, which is why [`RawWaker`] |
29 | /// exists as an alternative for those systems. |
30 | /// |
31 | /// To construct a [`Waker`] from some type `W` implementing this trait, |
32 | /// wrap it in an [`Arc<W>`](Arc) and call `Waker::from()` on that. |
33 | /// It is also possible to convert to [`RawWaker`] in the same way. |
34 | /// |
35 | /// <!-- Ideally we'd link to the `From` impl, but rustdoc doesn't generate any page for it within |
36 | /// `alloc` because `alloc` neither defines nor re-exports `From` or `Waker`, and we can't |
37 | /// link ../../std/task/struct.Waker.html#impl-From%3CArc%3CW,+Global%3E%3E-for-Waker |
38 | /// without getting a link-checking error in CI. --> |
39 | /// |
40 | /// # Examples |
41 | /// |
42 | /// A basic `block_on` function that takes a future and runs it to completion on |
43 | /// the current thread. |
44 | /// |
45 | /// **Note:** This example trades correctness for simplicity. In order to prevent |
46 | /// deadlocks, production-grade implementations will also need to handle |
47 | /// intermediate calls to `thread::unpark` as well as nested invocations. |
48 | /// |
49 | /// ```rust |
50 | /// use std::future::Future; |
51 | /// use std::sync::Arc; |
52 | /// use std::task::{Context, Poll, Wake}; |
53 | /// use std::thread::{self, Thread}; |
54 | /// use core::pin::pin; |
55 | /// |
56 | /// /// A waker that wakes up the current thread when called. |
57 | /// struct ThreadWaker(Thread); |
58 | /// |
59 | /// impl Wake for ThreadWaker { |
60 | /// fn wake(self: Arc<Self>) { |
61 | /// self.0.unpark(); |
62 | /// } |
63 | /// } |
64 | /// |
65 | /// /// Run a future to completion on the current thread. |
66 | /// fn block_on<T>(fut: impl Future<Output = T>) -> T { |
67 | /// // Pin the future so it can be polled. |
68 | /// let mut fut = pin!(fut); |
69 | /// |
70 | /// // Create a new context to be passed to the future. |
71 | /// let t = thread::current(); |
72 | /// let waker = Arc::new(ThreadWaker(t)).into(); |
73 | /// let mut cx = Context::from_waker(&waker); |
74 | /// |
75 | /// // Run the future to completion. |
76 | /// loop { |
77 | /// match fut.as_mut().poll(&mut cx) { |
78 | /// Poll::Ready(res) => return res, |
79 | /// Poll::Pending => thread::park(), |
80 | /// } |
81 | /// } |
82 | /// } |
83 | /// |
84 | /// block_on(async { |
85 | /// println!("Hi from inside a future!" ); |
86 | /// }); |
87 | /// ``` |
88 | #[cfg (target_has_atomic = "ptr" )] |
89 | #[stable (feature = "wake_trait" , since = "1.51.0" )] |
90 | pub trait Wake { |
91 | /// Wake this task. |
92 | #[stable (feature = "wake_trait" , since = "1.51.0" )] |
93 | fn wake(self: Arc<Self>); |
94 | |
95 | /// Wake this task without consuming the waker. |
96 | /// |
97 | /// If an executor supports a cheaper way to wake without consuming the |
98 | /// waker, it should override this method. By default, it clones the |
99 | /// [`Arc`] and calls [`wake`] on the clone. |
100 | /// |
101 | /// [`wake`]: Wake::wake |
102 | #[stable (feature = "wake_trait" , since = "1.51.0" )] |
103 | fn wake_by_ref(self: &Arc<Self>) { |
104 | self.clone().wake(); |
105 | } |
106 | } |
107 | #[cfg (target_has_atomic = "ptr" )] |
108 | #[stable (feature = "wake_trait" , since = "1.51.0" )] |
109 | impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker { |
110 | /// Use a [`Wake`]-able type as a `Waker`. |
111 | /// |
112 | /// No heap allocations or atomic operations are used for this conversion. |
113 | fn from(waker: Arc<W>) -> Waker { |
114 | // SAFETY: This is safe because raw_waker safely constructs |
115 | // a RawWaker from Arc<W>. |
116 | unsafe { Waker::from_raw(waker:raw_waker(waker)) } |
117 | } |
118 | } |
119 | #[cfg (target_has_atomic = "ptr" )] |
120 | #[stable (feature = "wake_trait" , since = "1.51.0" )] |
121 | impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker { |
122 | /// Use a `Wake`-able type as a `RawWaker`. |
123 | /// |
124 | /// No heap allocations or atomic operations are used for this conversion. |
125 | fn from(waker: Arc<W>) -> RawWaker { |
126 | raw_waker(waker) |
127 | } |
128 | } |
129 | |
130 | // NB: This private function for constructing a RawWaker is used, rather than |
131 | // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that |
132 | // the safety of `From<Arc<W>> for Waker` does not depend on the correct |
133 | // trait dispatch - instead both impls call this function directly and |
134 | // explicitly. |
135 | #[cfg (target_has_atomic = "ptr" )] |
136 | #[inline (always)] |
137 | fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker { |
138 | // Increment the reference count of the arc to clone it. |
139 | // |
140 | // The #[inline(always)] is to ensure that raw_waker and clone_waker are |
141 | // always generated in the same code generation unit as one another, and |
142 | // therefore that the structurally identical const-promoted RawWakerVTable |
143 | // within both functions is deduplicated at LLVM IR code generation time. |
144 | // This allows optimizing Waker::will_wake to a single pointer comparison of |
145 | // the vtable pointers, rather than comparing all four function pointers |
146 | // within the vtables. |
147 | #[inline (always)] |
148 | unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker { |
149 | unsafe { Arc::increment_strong_count(waker as *const W) }; |
150 | RawWaker::new( |
151 | waker, |
152 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), |
153 | ) |
154 | } |
155 | |
156 | // Wake by value, moving the Arc into the Wake::wake function |
157 | unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
158 | let waker = unsafe { Arc::from_raw(waker as *const W) }; |
159 | <W as Wake>::wake(waker); |
160 | } |
161 | |
162 | // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it |
163 | unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
164 | let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) }; |
165 | <W as Wake>::wake_by_ref(&waker); |
166 | } |
167 | |
168 | // Decrement the reference count of the Arc on drop |
169 | unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
170 | unsafe { Arc::decrement_strong_count(waker as *const W) }; |
171 | } |
172 | |
173 | RawWaker::new( |
174 | Arc::into_raw(waker) as *const (), |
175 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), |
176 | ) |
177 | } |
178 | |
179 | /// An analogous trait to `Wake` but used to construct a `LocalWaker`. This API |
180 | /// works in exactly the same way as `Wake`, except that it uses an `Rc` instead |
181 | /// of an `Arc`, and the result is a `LocalWaker` instead of a `Waker`. |
182 | /// |
183 | /// The benefits of using `LocalWaker` over `Waker` are that it allows the local waker |
184 | /// to hold data that does not implement `Send` and `Sync`. Additionally, it saves calls |
185 | /// to `Arc::clone`, which requires atomic synchronization. |
186 | /// |
187 | /// |
188 | /// # Examples |
189 | /// |
190 | /// This is a simplified example of a `spawn` and a `block_on` function. The `spawn` function |
191 | /// is used to push new tasks onto the run queue, while the block on function will remove them |
192 | /// and poll them. When a task is woken, it will put itself back on the run queue to be polled |
193 | /// by the executor. |
194 | /// |
195 | /// **Note:** This example trades correctness for simplicity. A real world example would interleave |
196 | /// poll calls with calls to an io reactor to wait for events instead of spinning on a loop. |
197 | /// |
198 | /// ```rust |
199 | /// #![feature(local_waker)] |
200 | /// #![feature(noop_waker)] |
201 | /// use std::task::{LocalWake, ContextBuilder, LocalWaker, Waker}; |
202 | /// use std::future::Future; |
203 | /// use std::pin::Pin; |
204 | /// use std::rc::Rc; |
205 | /// use std::cell::RefCell; |
206 | /// use std::collections::VecDeque; |
207 | /// |
208 | /// |
209 | /// thread_local! { |
210 | /// // A queue containing all tasks ready to do progress |
211 | /// static RUN_QUEUE: RefCell<VecDeque<Rc<Task>>> = RefCell::default(); |
212 | /// } |
213 | /// |
214 | /// type BoxedFuture = Pin<Box<dyn Future<Output = ()>>>; |
215 | /// |
216 | /// struct Task(RefCell<BoxedFuture>); |
217 | /// |
218 | /// impl LocalWake for Task { |
219 | /// fn wake(self: Rc<Self>) { |
220 | /// RUN_QUEUE.with_borrow_mut(|queue| { |
221 | /// queue.push_back(self) |
222 | /// }) |
223 | /// } |
224 | /// } |
225 | /// |
226 | /// fn spawn<F>(future: F) |
227 | /// where |
228 | /// F: Future<Output=()> + 'static + Send + Sync |
229 | /// { |
230 | /// let task = RefCell::new(Box::pin(future)); |
231 | /// RUN_QUEUE.with_borrow_mut(|queue| { |
232 | /// queue.push_back(Rc::new(Task(task))); |
233 | /// }); |
234 | /// } |
235 | /// |
236 | /// fn block_on<F>(future: F) |
237 | /// where |
238 | /// F: Future<Output=()> + 'static + Sync + Send |
239 | /// { |
240 | /// spawn(future); |
241 | /// loop { |
242 | /// let Some(task) = RUN_QUEUE.with_borrow_mut(|queue| queue.pop_front()) else { |
243 | /// // we exit, since there are no more tasks remaining on the queue |
244 | /// return; |
245 | /// }; |
246 | /// |
247 | /// // cast the Rc<Task> into a `LocalWaker` |
248 | /// let local_waker: LocalWaker = task.clone().into(); |
249 | /// // Build the context using `ContextBuilder` |
250 | /// let mut cx = ContextBuilder::from_waker(Waker::noop()) |
251 | /// .local_waker(&local_waker) |
252 | /// .build(); |
253 | /// |
254 | /// // Poll the task |
255 | /// let _ = task.0 |
256 | /// .borrow_mut() |
257 | /// .as_mut() |
258 | /// .poll(&mut cx); |
259 | /// } |
260 | /// } |
261 | /// |
262 | /// block_on(async { |
263 | /// println!("hello world" ); |
264 | /// }); |
265 | /// ``` |
266 | /// |
267 | #[unstable (feature = "local_waker" , issue = "118959" )] |
268 | pub trait LocalWake { |
269 | /// Wake this task. |
270 | #[unstable (feature = "local_waker" , issue = "118959" )] |
271 | fn wake(self: Rc<Self>); |
272 | |
273 | /// Wake this task without consuming the local waker. |
274 | /// |
275 | /// If an executor supports a cheaper way to wake without consuming the |
276 | /// waker, it should override this method. By default, it clones the |
277 | /// [`Rc`] and calls [`wake`] on the clone. |
278 | /// |
279 | /// [`wake`]: LocalWaker::wake |
280 | #[unstable (feature = "local_waker" , issue = "118959" )] |
281 | fn wake_by_ref(self: &Rc<Self>) { |
282 | self.clone().wake(); |
283 | } |
284 | } |
285 | |
286 | #[unstable (feature = "local_waker" , issue = "118959" )] |
287 | impl<W: LocalWake + 'static> From<Rc<W>> for LocalWaker { |
288 | /// Use a `Wake`-able type as a `LocalWaker`. |
289 | /// |
290 | /// No heap allocations or atomic operations are used for this conversion. |
291 | fn from(waker: Rc<W>) -> LocalWaker { |
292 | // SAFETY: This is safe because raw_waker safely constructs |
293 | // a RawWaker from Rc<W>. |
294 | unsafe { LocalWaker::from_raw(waker:local_raw_waker(waker)) } |
295 | } |
296 | } |
297 | #[allow (ineffective_unstable_trait_impl)] |
298 | #[unstable (feature = "local_waker" , issue = "118959" )] |
299 | impl<W: LocalWake + 'static> From<Rc<W>> for RawWaker { |
300 | /// Use a `Wake`-able type as a `RawWaker`. |
301 | /// |
302 | /// No heap allocations or atomic operations are used for this conversion. |
303 | fn from(waker: Rc<W>) -> RawWaker { |
304 | local_raw_waker(waker) |
305 | } |
306 | } |
307 | |
308 | // NB: This private function for constructing a RawWaker is used, rather than |
309 | // inlining this into the `From<Rc<W>> for RawWaker` impl, to ensure that |
310 | // the safety of `From<Rc<W>> for Waker` does not depend on the correct |
311 | // trait dispatch - instead both impls call this function directly and |
312 | // explicitly. |
313 | #[inline (always)] |
314 | fn local_raw_waker<W: LocalWake + 'static>(waker: Rc<W>) -> RawWaker { |
315 | // Increment the reference count of the Rc to clone it. |
316 | // |
317 | // Refer to the comment on raw_waker's clone_waker regarding why this is |
318 | // always inline. |
319 | #[inline (always)] |
320 | unsafe fn clone_waker<W: LocalWake + 'static>(waker: *const ()) -> RawWaker { |
321 | unsafe { Rc::increment_strong_count(waker as *const W) }; |
322 | RawWaker::new( |
323 | waker, |
324 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), |
325 | ) |
326 | } |
327 | |
328 | // Wake by value, moving the Rc into the LocalWake::wake function |
329 | unsafe fn wake<W: LocalWake + 'static>(waker: *const ()) { |
330 | let waker = unsafe { Rc::from_raw(waker as *const W) }; |
331 | <W as LocalWake>::wake(waker); |
332 | } |
333 | |
334 | // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it |
335 | unsafe fn wake_by_ref<W: LocalWake + 'static>(waker: *const ()) { |
336 | let waker = unsafe { ManuallyDrop::new(Rc::from_raw(waker as *const W)) }; |
337 | <W as LocalWake>::wake_by_ref(&waker); |
338 | } |
339 | |
340 | // Decrement the reference count of the Rc on drop |
341 | unsafe fn drop_waker<W: LocalWake + 'static>(waker: *const ()) { |
342 | unsafe { Rc::decrement_strong_count(waker as *const W) }; |
343 | } |
344 | |
345 | RawWaker::new( |
346 | Rc::into_raw(waker) as *const (), |
347 | &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), |
348 | ) |
349 | } |
350 | |