portable_atomic_util/
task.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3//! Types and Traits for working with asynchronous tasks.
4
5// This module is based on alloc::task::Wake.
6//
7// The code has been adjusted to work with stable Rust.
8//
9// Source: https://github.com/rust-lang/rust/blob/1.80.0/library/alloc/src/task.rs.
10//
11// Copyright & License of the original code:
12// - https://github.com/rust-lang/rust/blob/1.80.0/COPYRIGHT
13// - https://github.com/rust-lang/rust/blob/1.80.0/LICENSE-APACHE
14// - https://github.com/rust-lang/rust/blob/1.80.0/LICENSE-MIT
15
16use core::{
17    mem::ManuallyDrop,
18    task::{RawWaker, RawWakerVTable, Waker},
19};
20
21use crate::Arc;
22
23/// The implementation of waking a task on an executor.
24///
25/// This is an equivalent to [`std::task::Wake`], but using [`portable_atomic_util::Arc`](crate::Arc)
26/// as a reference-counted pointer. See the documentation for [`std::task::Wake`] for more details.
27///
28/// **Note:** Unlike `std::task::Wake`, all methods take `this:` instead of `self:`.
29/// This is because using `portable_atomic_util::Arc` as a receiver requires the
30/// [unstable `arbitrary_self_types` feature](https://github.com/rust-lang/rust/issues/44874).
31///
32/// # Examples
33///
34/// A basic `block_on` function that takes a future and runs it to completion on
35/// the current thread.
36///
37/// **Note:** This example trades correctness for simplicity. In order to prevent
38/// deadlocks, production-grade implementations will also need to handle
39/// intermediate calls to `thread::unpark` as well as nested invocations.
40///
41/// ```
42/// use portable_atomic_util::{task::Wake, Arc};
43/// use std::{
44///     future::Future,
45///     task::{Context, Poll},
46///     thread::{self, Thread},
47/// };
48///
49/// /// A waker that wakes up the current thread when called.
50/// struct ThreadWaker(Thread);
51///
52/// impl Wake for ThreadWaker {
53///     fn wake(this: Arc<Self>) {
54///         this.0.unpark();
55///     }
56/// }
57///
58/// /// Run a future to completion on the current thread.
59/// fn block_on<T>(fut: impl Future<Output = T>) -> T {
60///     // Pin the future so it can be polled.
61///     let mut fut = Box::pin(fut);
62///
63///     // Create a new context to be passed to the future.
64///     let t = thread::current();
65///     let waker = Arc::new(ThreadWaker(t)).into();
66///     let mut cx = Context::from_waker(&waker);
67///
68///     // Run the future to completion.
69///     loop {
70///         match fut.as_mut().poll(&mut cx) {
71///             Poll::Ready(res) => return res,
72///             Poll::Pending => thread::park(),
73///         }
74///     }
75/// }
76///
77/// block_on(async {
78///     println!("Hi from inside a future!");
79/// });
80/// ```
81pub trait Wake {
82    /// Wake this task.
83    fn wake(this: Arc<Self>);
84
85    /// Wake this task without consuming the waker.
86    ///
87    /// If an executor supports a cheaper way to wake without consuming the
88    /// waker, it should override this method. By default, it clones the
89    /// [`Arc`] and calls [`wake`] on the clone.
90    ///
91    /// [`wake`]: Wake::wake
92    fn wake_by_ref(this: &Arc<Self>) {
93        Self::wake(this.clone());
94    }
95}
96
97impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
98    /// Use a `Wake`-able type as a `Waker`.
99    ///
100    /// No heap allocations or atomic operations are used for this conversion.
101    fn from(waker: Arc<W>) -> Self {
102        // SAFETY: This is safe because raw_waker safely constructs
103        // a RawWaker from Arc<W>.
104        unsafe { Self::from_raw(raw_waker(waker)) }
105    }
106}
107
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>) -> Self {
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    //
126    // The #[inline(always)] is to ensure that raw_waker and clone_waker are
127    // always generated in the same code generation unit as one another, and
128    // therefore that the structurally identical const-promoted RawWakerVTable
129    // within both functions is deduplicated at LLVM IR code generation time.
130    // This allows optimizing Waker::will_wake to a single pointer comparison of
131    // the vtable pointers, rather than comparing all four function pointers
132    // within the vtables.
133    #[inline(always)]
134    unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
135        // SAFETY: the caller must uphold the safety contract.
136        unsafe { Arc::increment_strong_count(waker as *const W) };
137        RawWaker::new(
138            waker,
139            &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
140        )
141    }
142
143    // Wake by value, moving the Arc into the Wake::wake function
144    unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
145        // SAFETY: the caller must uphold the safety contract.
146        let waker = unsafe { Arc::from_raw(waker as *const W) };
147        <W as Wake>::wake(waker);
148    }
149
150    // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
151    unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
152        // SAFETY: the caller must uphold the safety contract.
153        let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
154        <W as Wake>::wake_by_ref(&waker);
155    }
156
157    // Decrement the reference count of the Arc on drop
158    unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
159        // SAFETY: the caller must uphold the safety contract.
160        unsafe { Arc::decrement_strong_count(waker as *const W) };
161    }
162
163    RawWaker::new(
164        Arc::into_raw(waker) as *const (),
165        &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
166    )
167}
portable_atomic_util/task.rs

portable_atomic_util/
task.rs
