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
10use core::mem::ManuallyDrop;
11#[cfg(target_has_atomic = "ptr")]
12use core::task::Waker;
13use core::task::{LocalWaker, RawWaker, RawWakerVTable};
14
15use crate::rc::Rc;
16#[cfg(target_has_atomic = "ptr")]
17use crate::sync::Arc;
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")]
90pub 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")]
109impl<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(raw_waker(waker)) }
117 }
118}
119#[cfg(target_has_atomic = "ptr")]
120#[stable(feature = "wake_trait", since = "1.51.0")]
121impl<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)]
137fn 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`.
180///
181/// This API works in exactly the same way as `Wake`,
182/// except that it uses an `Rc` instead of an `Arc`,
183/// and the result is a `LocalWaker` instead of a `Waker`.
184///
185/// The benefits of using `LocalWaker` over `Waker` are that it allows the local waker
186/// to hold data that does not implement `Send` and `Sync`. Additionally, it saves calls
187/// to `Arc::clone`, which requires atomic synchronization.
188///
189///
190/// # Examples
191///
192/// This is a simplified example of a `spawn` and a `block_on` function. The `spawn` function
193/// is used to push new tasks onto the run queue, while the block on function will remove them
194/// and poll them. When a task is woken, it will put itself back on the run queue to be polled
195/// by the executor.
196///
197/// **Note:** This example trades correctness for simplicity. A real world example would interleave
198/// poll calls with calls to an io reactor to wait for events instead of spinning on a loop.
199///
200/// ```rust
201/// #![feature(local_waker)]
202/// use std::task::{LocalWake, ContextBuilder, LocalWaker, Waker};
203/// use std::future::Future;
204/// use std::pin::Pin;
205/// use std::rc::Rc;
206/// use std::cell::RefCell;
207/// use std::collections::VecDeque;
208///
209///
210/// thread_local! {
211/// // A queue containing all tasks ready to do progress
212/// static RUN_QUEUE: RefCell<VecDeque<Rc<Task>>> = RefCell::default();
213/// }
214///
215/// type BoxedFuture = Pin<Box<dyn Future<Output = ()>>>;
216///
217/// struct Task(RefCell<BoxedFuture>);
218///
219/// impl LocalWake for Task {
220/// fn wake(self: Rc<Self>) {
221/// RUN_QUEUE.with_borrow_mut(|queue| {
222/// queue.push_back(self)
223/// })
224/// }
225/// }
226///
227/// fn spawn<F>(future: F)
228/// where
229/// F: Future<Output=()> + 'static + Send + Sync
230/// {
231/// let task = RefCell::new(Box::pin(future));
232/// RUN_QUEUE.with_borrow_mut(|queue| {
233/// queue.push_back(Rc::new(Task(task)));
234/// });
235/// }
236///
237/// fn block_on<F>(future: F)
238/// where
239/// F: Future<Output=()> + 'static + Sync + Send
240/// {
241/// spawn(future);
242/// loop {
243/// let Some(task) = RUN_QUEUE.with_borrow_mut(|queue| queue.pop_front()) else {
244/// // we exit, since there are no more tasks remaining on the queue
245/// return;
246/// };
247///
248/// // cast the Rc<Task> into a `LocalWaker`
249/// let local_waker: LocalWaker = task.clone().into();
250/// // Build the context using `ContextBuilder`
251/// let mut cx = ContextBuilder::from_waker(Waker::noop())
252/// .local_waker(&local_waker)
253/// .build();
254///
255/// // Poll the task
256/// let _ = task.0
257/// .borrow_mut()
258/// .as_mut()
259/// .poll(&mut cx);
260/// }
261/// }
262///
263/// block_on(async {
264/// println!("hello world");
265/// });
266/// ```
267///
268#[unstable(feature = "local_waker", issue = "118959")]
269pub trait LocalWake {
270 /// Wake this task.
271 #[unstable(feature = "local_waker", issue = "118959")]
272 fn wake(self: Rc<Self>);
273
274 /// Wake this task without consuming the local waker.
275 ///
276 /// If an executor supports a cheaper way to wake without consuming the
277 /// waker, it should override this method. By default, it clones the
278 /// [`Rc`] and calls [`wake`] on the clone.
279 ///
280 /// [`wake`]: LocalWaker::wake
281 #[unstable(feature = "local_waker", issue = "118959")]
282 fn wake_by_ref(self: &Rc<Self>) {
283 self.clone().wake();
284 }
285}
286
287#[unstable(feature = "local_waker", issue = "118959")]
288impl<W: LocalWake + 'static> From<Rc<W>> for LocalWaker {
289 /// Use a `Wake`-able type as a `LocalWaker`.
290 ///
291 /// No heap allocations or atomic operations are used for this conversion.
292 fn from(waker: Rc<W>) -> LocalWaker {
293 // SAFETY: This is safe because raw_waker safely constructs
294 // a RawWaker from Rc<W>.
295 unsafe { LocalWaker::from_raw(local_raw_waker(waker)) }
296 }
297}
298#[allow(ineffective_unstable_trait_impl)]
299#[unstable(feature = "local_waker", issue = "118959")]
300impl<W: LocalWake + 'static> From<Rc<W>> for RawWaker {
301 /// Use a `Wake`-able type as a `RawWaker`.
302 ///
303 /// No heap allocations or atomic operations are used for this conversion.
304 fn from(waker: Rc<W>) -> RawWaker {
305 local_raw_waker(waker)
306 }
307}
308
309// NB: This private function for constructing a RawWaker is used, rather than
310// inlining this into the `From<Rc<W>> for RawWaker` impl, to ensure that
311// the safety of `From<Rc<W>> for Waker` does not depend on the correct
312// trait dispatch - instead both impls call this function directly and
313// explicitly.
314#[inline(always)]
315fn local_raw_waker<W: LocalWake + 'static>(waker: Rc<W>) -> RawWaker {
316 // Increment the reference count of the Rc to clone it.
317 //
318 // Refer to the comment on raw_waker's clone_waker regarding why this is
319 // always inline.
320 #[inline(always)]
321 unsafe fn clone_waker<W: LocalWake + 'static>(waker: *const ()) -> RawWaker {
322 unsafe { Rc::increment_strong_count(waker as *const W) };
323 RawWaker::new(
324 waker,
325 &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
326 )
327 }
328
329 // Wake by value, moving the Rc into the LocalWake::wake function
330 unsafe fn wake<W: LocalWake + 'static>(waker: *const ()) {
331 let waker = unsafe { Rc::from_raw(waker as *const W) };
332 <W as LocalWake>::wake(waker);
333 }
334
335 // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
336 unsafe fn wake_by_ref<W: LocalWake + 'static>(waker: *const ()) {
337 let waker = unsafe { ManuallyDrop::new(Rc::from_raw(waker as *const W)) };
338 <W as LocalWake>::wake_by_ref(&waker);
339 }
340
341 // Decrement the reference count of the Rc on drop
342 unsafe fn drop_waker<W: LocalWake + 'static>(waker: *const ()) {
343 unsafe { Rc::decrement_strong_count(waker as *const W) };
344 }
345
346 RawWaker::new(
347 Rc::into_raw(waker) as *const (),
348 &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
349 )
350}
351