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