alloc/
task.rs

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}