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