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 crate::rc::Rc;
11use core::mem::ManuallyDrop;
12use core::task::{LocalWaker, RawWaker, RawWakerVTable};
13
14#[cfg(target_has_atomic = "ptr")]
15use crate::sync::Arc;
16#[cfg(target_has_atomic = "ptr")]
17use 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")]
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(waker: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`. 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")]
268pub 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")]
287impl<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")]
299impl<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)]
314fn 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