1 #![cfg_attr(any(loom, not(feature = "sync")), allow(dead_code, unreachable_pub))] 2 3 use crate::loom::cell::UnsafeCell; 4 use crate::loom::sync::atomic::{self, AtomicUsize}; 5 6 use std::fmt; 7 use std::panic::{resume_unwind, AssertUnwindSafe, RefUnwindSafe, UnwindSafe}; 8 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; 9 use std::task::Waker; 10 11 /// A synchronization primitive for task waking. 12 /// 13 /// `AtomicWaker` will coordinate concurrent wakes with the consumer 14 /// potentially "waking" the underlying task. This is useful in scenarios 15 /// where a computation completes in another thread and wants to wake the 16 /// consumer, but the consumer is in the process of being migrated to a new 17 /// logical task. 18 /// 19 /// Consumers should call `register` before checking the result of a computation 20 /// and producers should call `wake` after producing the computation (this 21 /// differs from the usual `thread::park` pattern). It is also permitted for 22 /// `wake` to be called **before** `register`. This results in a no-op. 23 /// 24 /// A single `AtomicWaker` may be reused for any number of calls to `register` or 25 /// `wake`. 26 pub(crate) struct AtomicWaker { 27 state: AtomicUsize, 28 waker: UnsafeCell<Option<Waker>>, 29 } 30 31 impl RefUnwindSafe for AtomicWaker {} 32 impl UnwindSafe for AtomicWaker {} 33 34 // `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell 35 // stores a `Waker` value produced by calls to `register` and many threads can 36 // race to take the waker by calling `wake`. 37 // 38 // If a new `Waker` instance is produced by calling `register` before an existing 39 // one is consumed, then the existing one is overwritten. 40 // 41 // While `AtomicWaker` is single-producer, the implementation ensures memory 42 // safety. In the event of concurrent calls to `register`, there will be a 43 // single winner whose waker will get stored in the cell. The losers will not 44 // have their tasks woken. As such, callers should ensure to add synchronization 45 // to calls to `register`. 46 // 47 // The implementation uses a single `AtomicUsize` value to coordinate access to 48 // the `Waker` cell. There are two bits that are operated on independently. These 49 // are represented by `REGISTERING` and `WAKING`. 50 // 51 // The `REGISTERING` bit is set when a producer enters the critical section. The 52 // `WAKING` bit is set when a consumer enters the critical section. Neither 53 // bit being set is represented by `WAITING`. 54 // 55 // A thread obtains an exclusive lock on the waker cell by transitioning the 56 // state from `WAITING` to `REGISTERING` or `WAKING`, depending on the 57 // operation the thread wishes to perform. When this transition is made, it is 58 // guaranteed that no other thread will access the waker cell. 59 // 60 // # Registering 61 // 62 // On a call to `register`, an attempt to transition the state from WAITING to 63 // REGISTERING is made. On success, the caller obtains a lock on the waker cell. 64 // 65 // If the lock is obtained, then the thread sets the waker cell to the waker 66 // provided as an argument. Then it attempts to transition the state back from 67 // `REGISTERING` -> `WAITING`. 68 // 69 // If this transition is successful, then the registering process is complete 70 // and the next call to `wake` will observe the waker. 71 // 72 // If the transition fails, then there was a concurrent call to `wake` that 73 // was unable to access the waker cell (due to the registering thread holding the 74 // lock). To handle this, the registering thread removes the waker it just set 75 // from the cell and calls `wake` on it. This call to wake represents the 76 // attempt to wake by the other thread (that set the `WAKING` bit). The 77 // state is then transitioned from `REGISTERING | WAKING` back to `WAITING`. 78 // This transition must succeed because, at this point, the state cannot be 79 // transitioned by another thread. 80 // 81 // # Waking 82 // 83 // On a call to `wake`, an attempt to transition the state from `WAITING` to 84 // `WAKING` is made. On success, the caller obtains a lock on the waker cell. 85 // 86 // If the lock is obtained, then the thread takes ownership of the current value 87 // in the waker cell, and calls `wake` on it. The state is then transitioned 88 // back to `WAITING`. This transition must succeed as, at this point, the state 89 // cannot be transitioned by another thread. 90 // 91 // If the thread is unable to obtain the lock, the `WAKING` bit is still set. 92 // This is because it has either been set by the current thread but the previous 93 // value included the `REGISTERING` bit **or** a concurrent thread is in the 94 // `WAKING` critical section. Either way, no action must be taken. 95 // 96 // If the current thread is the only concurrent call to `wake` and another 97 // thread is in the `register` critical section, when the other thread **exits** 98 // the `register` critical section, it will observe the `WAKING` bit and 99 // handle the waker itself. 100 // 101 // If another thread is in the `waker` critical section, then it will handle 102 // waking the caller task. 103 // 104 // # A potential race (is safely handled). 105 // 106 // Imagine the following situation: 107 // 108 // * Thread A obtains the `wake` lock and wakes a task. 109 // 110 // * Before thread A releases the `wake` lock, the woken task is scheduled. 111 // 112 // * Thread B attempts to wake the task. In theory this should result in the 113 // task being woken, but it cannot because thread A still holds the wake 114 // lock. 115 // 116 // This case is handled by requiring users of `AtomicWaker` to call `register` 117 // **before** attempting to observe the application state change that resulted 118 // in the task being woken. The wakers also change the application state 119 // before calling wake. 120 // 121 // Because of this, the task will do one of two things. 122 // 123 // 1) Observe the application state change that Thread B is waking on. In 124 // this case, it is OK for Thread B's wake to be lost. 125 // 126 // 2) Call register before attempting to observe the application state. Since 127 // Thread A still holds the `wake` lock, the call to `register` will result 128 // in the task waking itself and get scheduled again. 129 130 /// Idle state. 131 const WAITING: usize = 0; 132 133 /// A new waker value is being registered with the `AtomicWaker` cell. 134 const REGISTERING: usize = 0b01; 135 136 /// The task currently registered with the `AtomicWaker` cell is being woken. 137 const WAKING: usize = 0b10; 138 139 impl AtomicWaker { 140 /// Create an `AtomicWaker` new() -> AtomicWaker141 pub(crate) fn new() -> AtomicWaker { 142 AtomicWaker { 143 state: AtomicUsize::new(WAITING), 144 waker: UnsafeCell::new(None), 145 } 146 } 147 148 /* 149 /// Registers the current waker to be notified on calls to `wake`. 150 pub(crate) fn register(&self, waker: Waker) { 151 self.do_register(waker); 152 } 153 */ 154 155 /// Registers the provided waker to be notified on calls to `wake`. 156 /// 157 /// The new waker will take place of any previous wakers that were registered 158 /// by previous calls to `register`. Any calls to `wake` that happen after 159 /// a call to `register` (as defined by the memory ordering rules), will 160 /// wake the `register` caller's task. 161 /// 162 /// It is safe to call `register` with multiple other threads concurrently 163 /// calling `wake`. This will result in the `register` caller's current 164 /// task being woken once. 165 /// 166 /// This function is safe to call concurrently, but this is generally a bad 167 /// idea. Concurrent calls to `register` will attempt to register different 168 /// tasks to be woken. One of the callers will win and have its task set, 169 /// but there is no guarantee as to which caller will succeed. register_by_ref(&self, waker: &Waker)170 pub(crate) fn register_by_ref(&self, waker: &Waker) { 171 self.do_register(waker); 172 } 173 do_register<W>(&self, waker: W) where W: WakerRef,174 fn do_register<W>(&self, waker: W) 175 where 176 W: WakerRef, 177 { 178 fn catch_unwind<F: FnOnce() -> R, R>(f: F) -> std::thread::Result<R> { 179 std::panic::catch_unwind(AssertUnwindSafe(f)) 180 } 181 182 match self 183 .state 184 .compare_exchange(WAITING, REGISTERING, Acquire, Acquire) 185 .unwrap_or_else(|x| x) 186 { 187 WAITING => { 188 unsafe { 189 // If `into_waker` panics (because it's code outside of 190 // AtomicWaker) we need to prime a guard that is called on 191 // unwind to restore the waker to a WAITING state. Otherwise 192 // any future calls to register will incorrectly be stuck 193 // believing it's being updated by someone else. 194 let new_waker_or_panic = catch_unwind(move || waker.into_waker()); 195 196 // Set the field to contain the new waker, or if 197 // `into_waker` panicked, leave the old value. 198 let mut maybe_panic = None; 199 let mut old_waker = None; 200 match new_waker_or_panic { 201 Ok(new_waker) => { 202 old_waker = self.waker.with_mut(|t| (*t).take()); 203 self.waker.with_mut(|t| *t = Some(new_waker)); 204 } 205 Err(panic) => maybe_panic = Some(panic), 206 } 207 208 // Release the lock. If the state transitioned to include 209 // the `WAKING` bit, this means that a wake has been 210 // called concurrently, so we have to remove the waker and 211 // wake it.` 212 // 213 // Start by assuming that the state is `REGISTERING` as this 214 // is what we jut set it to. 215 let res = self 216 .state 217 .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire); 218 219 match res { 220 Ok(_) => { 221 // We don't want to give the caller the panic if it 222 // was someone else who put in that waker. 223 let _ = catch_unwind(move || { 224 drop(old_waker); 225 }); 226 } 227 Err(actual) => { 228 // This branch can only be reached if a 229 // concurrent thread called `wake`. In this 230 // case, `actual` **must** be `REGISTERING | 231 // WAKING`. 232 debug_assert_eq!(actual, REGISTERING | WAKING); 233 234 // Take the waker to wake once the atomic operation has 235 // completed. 236 let mut waker = self.waker.with_mut(|t| (*t).take()); 237 238 // Just swap, because no one could change state 239 // while state == `Registering | `Waking` 240 self.state.swap(WAITING, AcqRel); 241 242 // If `into_waker` panicked, then the waker in the 243 // waker slot is actually the old waker. 244 if maybe_panic.is_some() { 245 old_waker = waker.take(); 246 } 247 248 // We don't want to give the caller the panic if it 249 // was someone else who put in that waker. 250 if let Some(old_waker) = old_waker { 251 let _ = catch_unwind(move || { 252 old_waker.wake(); 253 }); 254 } 255 256 // The atomic swap was complete, now wake the waker 257 // and return. 258 // 259 // If this panics, we end up in a consumed state and 260 // return the panic to the caller. 261 if let Some(waker) = waker { 262 debug_assert!(maybe_panic.is_none()); 263 waker.wake(); 264 } 265 } 266 } 267 268 if let Some(panic) = maybe_panic { 269 // If `into_waker` panicked, return the panic to the caller. 270 resume_unwind(panic); 271 } 272 } 273 } 274 WAKING => { 275 // Currently in the process of waking the task, i.e., 276 // `wake` is currently being called on the old waker. 277 // So, we call wake on the new waker. 278 // 279 // If this panics, someone else is responsible for restoring the 280 // state of the waker. 281 waker.wake(); 282 283 // This is equivalent to a spin lock, so use a spin hint. 284 // TODO: once we bump MSRV to 1.49+, use `hint::spin_loop` instead. 285 #[allow(deprecated)] 286 atomic::spin_loop_hint(); 287 } 288 state => { 289 // In this case, a concurrent thread is holding the 290 // "registering" lock. This probably indicates a bug in the 291 // caller's code as racing to call `register` doesn't make much 292 // sense. 293 // 294 // We just want to maintain memory safety. It is ok to drop the 295 // call to `register`. 296 debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); 297 } 298 } 299 } 300 301 /// Wakes the task that last called `register`. 302 /// 303 /// If `register` has not been called yet, then this does nothing. wake(&self)304 pub(crate) fn wake(&self) { 305 if let Some(waker) = self.take_waker() { 306 // If wake panics, we've consumed the waker which is a legitimate 307 // outcome. 308 waker.wake(); 309 } 310 } 311 312 /// Attempts to take the `Waker` value out of the `AtomicWaker` with the 313 /// intention that the caller will wake the task later. take_waker(&self) -> Option<Waker>314 pub(crate) fn take_waker(&self) -> Option<Waker> { 315 // AcqRel ordering is used in order to acquire the value of the `waker` 316 // cell as well as to establish a `release` ordering with whatever 317 // memory the `AtomicWaker` is associated with. 318 match self.state.fetch_or(WAKING, AcqRel) { 319 WAITING => { 320 // The waking lock has been acquired. 321 let waker = unsafe { self.waker.with_mut(|t| (*t).take()) }; 322 323 // Release the lock 324 self.state.fetch_and(!WAKING, Release); 325 326 waker 327 } 328 state => { 329 // There is a concurrent thread currently updating the 330 // associated waker. 331 // 332 // Nothing more to do as the `WAKING` bit has been set. It 333 // doesn't matter if there are concurrent registering threads or 334 // not. 335 // 336 debug_assert!( 337 state == REGISTERING || state == REGISTERING | WAKING || state == WAKING 338 ); 339 None 340 } 341 } 342 } 343 } 344 345 impl Default for AtomicWaker { default() -> Self346 fn default() -> Self { 347 AtomicWaker::new() 348 } 349 } 350 351 impl fmt::Debug for AtomicWaker { fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result352 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { 353 write!(fmt, "AtomicWaker") 354 } 355 } 356 357 unsafe impl Send for AtomicWaker {} 358 unsafe impl Sync for AtomicWaker {} 359 360 trait WakerRef { wake(self)361 fn wake(self); into_waker(self) -> Waker362 fn into_waker(self) -> Waker; 363 } 364 365 impl WakerRef for Waker { wake(self)366 fn wake(self) { 367 self.wake() 368 } 369 into_waker(self) -> Waker370 fn into_waker(self) -> Waker { 371 self 372 } 373 } 374 375 impl WakerRef for &Waker { wake(self)376 fn wake(self) { 377 self.wake_by_ref() 378 } 379 into_waker(self) -> Waker380 fn into_waker(self) -> Waker { 381 self.clone() 382 } 383 } 384