1use core::cell::UnsafeCell;
2use core::fmt;
3use core::task::Waker;
4
5use atomic::AtomicUsize;
6use atomic::Ordering::{AcqRel, Acquire, Release};
7
8#[cfg(feature = "portable-atomic")]
9use portable_atomic as atomic;
10
11#[cfg(not(feature = "portable-atomic"))]
12use core::sync::atomic;
13
14/// A synchronization primitive for task wakeup.
15///
16/// Sometimes the task interested in a given event will change over time.
17/// An `AtomicWaker` can coordinate concurrent notifications with the consumer
18/// potentially "updating" the underlying task to wake up. This is useful in
19/// scenarios where a computation completes in another thread and wants to
20/// notify the consumer, but the consumer is in the process of being migrated to
21/// a new logical task.
22///
23/// Consumers should call `register` before checking the result of a computation
24/// and producers should call `wake` after producing the computation (this
25/// differs from the usual `thread::park` pattern). It is also permitted for
26/// `wake` to be called **before** `register`. This results in a no-op.
27///
28/// A single `AtomicWaker` may be reused for any number of calls to `register` or
29/// `wake`.
30///
31/// # Memory ordering
32///
33/// Calling `register` "acquires" all memory "released" by calls to `wake`
34/// before the call to `register`. Later calls to `wake` will wake the
35/// registered waker (on contention this wake might be triggered in `register`).
36///
37/// For concurrent calls to `register` (should be avoided) the ordering is only
38/// guaranteed for the winning call.
39///
40/// # Examples
41///
42/// Here is a simple example providing a `Flag` that can be signalled manually
43/// when it is ready.
44///
45/// ```
46/// use futures::future::Future;
47/// use futures::task::{Context, Poll, AtomicWaker};
48/// use std::sync::Arc;
49/// use std::sync::atomic::AtomicBool;
50/// use std::sync::atomic::Ordering::Relaxed;
51/// use std::pin::Pin;
52///
53/// struct Inner {
54/// waker: AtomicWaker,
55/// set: AtomicBool,
56/// }
57///
58/// #[derive(Clone)]
59/// struct Flag(Arc<Inner>);
60///
61/// impl Flag {
62/// pub fn new() -> Self {
63/// Self(Arc::new(Inner {
64/// waker: AtomicWaker::new(),
65/// set: AtomicBool::new(false),
66/// }))
67/// }
68///
69/// pub fn signal(&self) {
70/// self.0.set.store(true, Relaxed);
71/// self.0.waker.wake();
72/// }
73/// }
74///
75/// impl Future for Flag {
76/// type Output = ();
77///
78/// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
79/// // quick check to avoid registration if already done.
80/// if self.0.set.load(Relaxed) {
81/// return Poll::Ready(());
82/// }
83///
84/// self.0.waker.register(cx.waker());
85///
86/// // Need to check condition **after** `register` to avoid a race
87/// // condition that would result in lost notifications.
88/// if self.0.set.load(Relaxed) {
89/// Poll::Ready(())
90/// } else {
91/// Poll::Pending
92/// }
93/// }
94/// }
95/// ```
96pub struct AtomicWaker {
97 state: AtomicUsize,
98 waker: UnsafeCell<Option<Waker>>,
99}
100
101// `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell
102// stores a `Waker` value produced by calls to `register` and many threads can
103// race to take the waker (to wake it) by calling `wake`.
104//
105// If a new `Waker` instance is produced by calling `register` before an
106// existing one is consumed, then the existing one is overwritten.
107//
108// While `AtomicWaker` is single-producer, the implementation ensures memory
109// safety. In the event of concurrent calls to `register`, there will be a
110// single winner whose waker will get stored in the cell. The losers will not
111// have their tasks woken. As such, callers should ensure to add synchronization
112// to calls to `register`.
113//
114// The implementation uses a single `AtomicUsize` value to coordinate access to
115// the `Waker` cell. There are two bits that are operated on independently.
116// These are represented by `REGISTERING` and `WAKING`.
117//
118// The `REGISTERING` bit is set when a producer enters the critical section. The
119// `WAKING` bit is set when a consumer enters the critical section. Neither bit
120// being set is represented by `WAITING`.
121//
122// A thread obtains an exclusive lock on the waker cell by transitioning the
123// state from `WAITING` to `REGISTERING` or `WAKING`, depending on the operation
124// the thread wishes to perform. When this transition is made, it is guaranteed
125// that no other thread will access the waker cell.
126//
127// # Registering
128//
129// On a call to `register`, an attempt to transition the state from WAITING to
130// REGISTERING is made. On success, the caller obtains a lock on the waker cell.
131//
132// If the lock is obtained, then the thread sets the waker cell to the waker
133// provided as an argument. Then it attempts to transition the state back from
134// `REGISTERING` -> `WAITING`.
135//
136// If this transition is successful, then the registering process is complete
137// and the next call to `wake` will observe the waker.
138//
139// If the transition fails, then there was a concurrent call to `wake` that was
140// unable to access the waker cell (due to the registering thread holding the
141// lock). To handle this, the registering thread removes the waker it just set
142// from the cell and calls `wake` on it. This call to wake represents the
143// attempt to wake by the other thread (that set the `WAKING` bit). The state is
144// then transitioned from `REGISTERING | WAKING` back to `WAITING`. This
145// transition must succeed because, at this point, the state cannot be
146// transitioned by another thread.
147//
148// # Waking
149//
150// On a call to `wake`, an attempt to transition the state from `WAITING` to
151// `WAKING` is made. On success, the caller obtains a lock on the waker cell.
152//
153// If the lock is obtained, then the thread takes ownership of the current value
154// in the waker cell, and calls `wake` on it. The state is then transitioned
155// back to `WAITING`. This transition must succeed as, at this point, the state
156// cannot be transitioned by another thread.
157//
158// If the thread is unable to obtain the lock, the `WAKING` bit is still. This
159// is because it has either been set by the current thread but the previous
160// value included the `REGISTERING` bit **or** a concurrent thread is in the
161// `WAKING` critical section. Either way, no action must be taken.
162//
163// If the current thread is the only concurrent call to `wake` and another
164// thread is in the `register` critical section, when the other thread **exits**
165// the `register` critical section, it will observe the `WAKING` bit and handle
166// the wake itself.
167//
168// If another thread is in the `wake` critical section, then it will handle
169// waking the task.
170//
171// # A potential race (is safely handled).
172//
173// Imagine the following situation:
174//
175// * Thread A obtains the `wake` lock and wakes a task.
176//
177// * Before thread A releases the `wake` lock, the woken task is scheduled.
178//
179// * Thread B attempts to wake the task. In theory this should result in the
180// task being woken, but it cannot because thread A still holds the wake lock.
181//
182// This case is handled by requiring users of `AtomicWaker` to call `register`
183// **before** attempting to observe the application state change that resulted
184// in the task being awoken. The wakers also change the application state before
185// calling wake.
186//
187// Because of this, the waker will do one of two things.
188//
189// 1) Observe the application state change that Thread B is woken for. In this
190// case, it is OK for Thread B's wake to be lost.
191//
192// 2) Call register before attempting to observe the application state. Since
193// Thread A still holds the `wake` lock, the call to `register` will result
194// in the task waking itself and get scheduled again.
195
196/// Idle state
197const WAITING: usize = 0;
198
199/// A new waker value is being registered with the `AtomicWaker` cell.
200const REGISTERING: usize = 0b01;
201
202/// The waker currently registered with the `AtomicWaker` cell is being woken.
203const WAKING: usize = 0b10;
204
205impl AtomicWaker {
206 /// Create an `AtomicWaker`.
207 pub const fn new() -> Self {
208 // Make sure that task is Sync
209 trait AssertSync: Sync {}
210 impl AssertSync for Waker {}
211
212 Self { state: AtomicUsize::new(WAITING), waker: UnsafeCell::new(None) }
213 }
214
215 /// Registers the waker to be notified on calls to `wake`.
216 ///
217 /// The new task will take place of any previous tasks that were registered
218 /// by previous calls to `register`. Any calls to `wake` that happen after
219 /// a call to `register` (as defined by the memory ordering rules), will
220 /// notify the `register` caller's task and deregister the waker from future
221 /// notifications. Because of this, callers should ensure `register` gets
222 /// invoked with a new `Waker` **each** time they require a wakeup.
223 ///
224 /// It is safe to call `register` with multiple other threads concurrently
225 /// calling `wake`. This will result in the `register` caller's current
226 /// task being notified once.
227 ///
228 /// This function is safe to call concurrently, but this is generally a bad
229 /// idea. Concurrent calls to `register` will attempt to register different
230 /// tasks to be notified. One of the callers will win and have its task set,
231 /// but there is no guarantee as to which caller will succeed.
232 ///
233 /// # Examples
234 ///
235 /// Here is how `register` is used when implementing a flag.
236 ///
237 /// ```
238 /// use futures::future::Future;
239 /// use futures::task::{Context, Poll, AtomicWaker};
240 /// use std::sync::atomic::AtomicBool;
241 /// use std::sync::atomic::Ordering::Relaxed;
242 /// use std::pin::Pin;
243 ///
244 /// struct Flag {
245 /// waker: AtomicWaker,
246 /// set: AtomicBool,
247 /// }
248 ///
249 /// impl Future for Flag {
250 /// type Output = ();
251 ///
252 /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
253 /// // Register **before** checking `set` to avoid a race condition
254 /// // that would result in lost notifications.
255 /// self.waker.register(cx.waker());
256 ///
257 /// if self.set.load(Relaxed) {
258 /// Poll::Ready(())
259 /// } else {
260 /// Poll::Pending
261 /// }
262 /// }
263 /// }
264 /// ```
265 pub fn register(&self, waker: &Waker) {
266 match self
267 .state
268 .compare_exchange(WAITING, REGISTERING, Acquire, Acquire)
269 .unwrap_or_else(|x| x)
270 {
271 WAITING => {
272 unsafe {
273 // Locked acquired, update the waker cell
274
275 // Avoid cloning the waker if the old waker will awaken the same task.
276 match &*self.waker.get() {
277 Some(old_waker) if old_waker.will_wake(waker) => (),
278 _ => *self.waker.get() = Some(waker.clone()),
279 }
280
281 // Release the lock. If the state transitioned to include
282 // the `WAKING` bit, this means that at least one wake has
283 // been called concurrently.
284 //
285 // Start by assuming that the state is `REGISTERING` as this
286 // is what we just set it to. If this holds, we know that no
287 // other writes were performed in the meantime, so there is
288 // nothing to acquire, only release. In case of concurrent
289 // wakers, we need to acquire their releases, so success needs
290 // to do both.
291 let res = self.state.compare_exchange(REGISTERING, WAITING, AcqRel, Acquire);
292
293 match res {
294 Ok(_) => {
295 // memory ordering: acquired self.state during CAS
296 // - if previous wakes went through it syncs with
297 // their final release (`fetch_and`)
298 // - if there was no previous wake the next wake
299 // will wake us, no sync needed.
300 }
301 Err(actual) => {
302 // This branch can only be reached if at least one
303 // concurrent thread called `wake`. In this
304 // case, `actual` **must** be `REGISTERING |
305 // `WAKING`.
306 debug_assert_eq!(actual, REGISTERING | WAKING);
307
308 // Take the waker to wake once the atomic operation has
309 // completed.
310 let waker = (*self.waker.get()).take().unwrap();
311
312 // We need to return to WAITING state (clear our lock and
313 // concurrent WAKING flag). This needs to acquire all
314 // WAKING fetch_or releases and it needs to release our
315 // update to self.waker, so we need a `swap` operation.
316 self.state.swap(WAITING, AcqRel);
317
318 // memory ordering: we acquired the state for all
319 // concurrent wakes, but future wakes might still
320 // need to wake us in case we can't make progress
321 // from the pending wakes.
322 //
323 // So we simply schedule to come back later (we could
324 // also simply leave the registration in place above).
325 waker.wake();
326 }
327 }
328 }
329 }
330 WAKING => {
331 // Currently in the process of waking the task, i.e.,
332 // `wake` is currently being called on the old task handle.
333 //
334 // memory ordering: we acquired the state for all
335 // concurrent wakes, but future wakes might still
336 // need to wake us in case we can't make progress
337 // from the pending wakes.
338 //
339 // So we simply schedule to come back later (we
340 // could also spin here trying to acquire the lock
341 // to register).
342 waker.wake_by_ref();
343 }
344 state => {
345 // In this case, a concurrent thread is holding the
346 // "registering" lock. This probably indicates a bug in the
347 // caller's code as racing to call `register` doesn't make much
348 // sense.
349 //
350 // memory ordering: don't care. a concurrent register() is going
351 // to succeed and provide proper memory ordering.
352 //
353 // We just want to maintain memory safety. It is ok to drop the
354 // call to `register`.
355 debug_assert!(state == REGISTERING || state == REGISTERING | WAKING);
356 }
357 }
358 }
359
360 /// Calls `wake` on the last `Waker` passed to `register`.
361 ///
362 /// If `register` has not been called yet, then this does nothing.
363 pub fn wake(&self) {
364 if let Some(waker) = self.take() {
365 waker.wake();
366 }
367 }
368
369 /// Returns the last `Waker` passed to `register`, so that the user can wake it.
370 ///
371 ///
372 /// Sometimes, just waking the AtomicWaker is not fine grained enough. This allows the user
373 /// to take the waker and then wake it separately, rather than performing both steps in one
374 /// atomic action.
375 ///
376 /// If a waker has not been registered, this returns `None`.
377 pub fn take(&self) -> Option<Waker> {
378 // AcqRel ordering is used in order to acquire the value of the `task`
379 // cell as well as to establish a `release` ordering with whatever
380 // memory the `AtomicWaker` is associated with.
381 match self.state.fetch_or(WAKING, AcqRel) {
382 WAITING => {
383 // The waking lock has been acquired.
384 let waker = unsafe { (*self.waker.get()).take() };
385
386 // Release the lock
387 self.state.fetch_and(!WAKING, Release);
388
389 waker
390 }
391 state => {
392 // There is a concurrent thread currently updating the
393 // associated task.
394 //
395 // Nothing more to do as the `WAKING` bit has been set. It
396 // doesn't matter if there are concurrent registering threads or
397 // not.
398 //
399 debug_assert!(
400 state == REGISTERING || state == REGISTERING | WAKING || state == WAKING
401 );
402 None
403 }
404 }
405 }
406}
407
408impl Default for AtomicWaker {
409 fn default() -> Self {
410 Self::new()
411 }
412}
413
414impl fmt::Debug for AtomicWaker {
415 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
416 write!(f, "AtomicWaker")
417 }
418}
419
420unsafe impl Send for AtomicWaker {}
421unsafe impl Sync for AtomicWaker {}
422