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 /// Registers the current waker to be notified on calls to `wake`. 145 /// 146 /// This is the same as calling `register_task` with `task::current()`. 147 #[cfg(feature = "io-driver")] register(&self, waker: Waker)148 pub(crate) fn register(&self, waker: Waker) { 149 self.do_register(waker); 150 } 151 152 /// Registers the provided waker to be notified on calls to `wake`. 153 /// 154 /// The new waker will take place of any previous wakers that were registered 155 /// by previous calls to `register`. Any calls to `wake` that happen after 156 /// a call to `register` (as defined by the memory ordering rules), will 157 /// wake the `register` caller's task. 158 /// 159 /// It is safe to call `register` with multiple other threads concurrently 160 /// calling `wake`. This will result in the `register` caller's current 161 /// task being woken once. 162 /// 163 /// This function is safe to call concurrently, but this is generally a bad 164 /// idea. Concurrent calls to `register` will attempt to register different 165 /// tasks to be woken. One of the callers will win and have its task set, 166 /// but there is no guarantee as to which caller will succeed. register_by_ref(&self, waker: &Waker)167 pub(crate) fn register_by_ref(&self, waker: &Waker) { 168 self.do_register(waker); 169 } 170 do_register<W>(&self, waker: W) where W: WakerRef,171 fn do_register<W>(&self, waker: W) 172 where 173 W: WakerRef, 174 { 175 match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) { 176 WAITING => { 177 unsafe { 178 // Locked acquired, update the waker cell 179 self.waker.with_mut(|t| *t = Some(waker.into_waker())); 180 181 // Release the lock. If the state transitioned to include 182 // the `WAKING` bit, this means that a wake has been 183 // called concurrently, so we have to remove the waker and 184 // wake it.` 185 // 186 // Start by assuming that the state is `REGISTERING` as this 187 // is what we jut set it to. 188 let res = self 189 .state 190 .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire); 191 192 match res { 193 Ok(_) => {} 194 Err(actual) => { 195 // This branch can only be reached if a 196 // concurrent thread called `wake`. In this 197 // case, `actual` **must** be `REGISTERING | 198 // `WAKING`. 199 debug_assert_eq!(actual, REGISTERING | WAKING); 200 201 // Take the waker to wake once the atomic operation has 202 // completed. 203 let waker = self.waker.with_mut(|t| (*t).take()).unwrap(); 204 205 // Just swap, because no one could change state 206 // while state == `Registering | `Waking` 207 self.state.swap(WAITING, AcqRel); 208 209 // The atomic swap was complete, now 210 // wake the waker and return. 211 waker.wake(); 212 } 213 } 214 } 215 } 216 WAKING => { 217 // Currently in the process of waking the task, i.e., 218 // `wake` is currently being called on the old waker. 219 // So, we call wake on the new waker. 220 waker.wake(); 221 222 // This is equivalent to a spin lock, so use a spin hint. 223 atomic::spin_loop_hint(); 224 } 225 state => { 226 // In this case, a concurrent thread is holding the 227 // "registering" lock. This probably indicates a bug in the 228 // caller's code as racing to call `register` doesn't make much 229 // sense. 230 // 231 // We just want to maintain memory safety. It is ok to drop the 232 // call to `register`. 233 debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); 234 } 235 } 236 } 237 238 /// Wakes the task that last called `register`. 239 /// 240 /// If `register` has not been called yet, then this does nothing. wake(&self)241 pub(crate) fn wake(&self) { 242 if let Some(waker) = self.take_waker() { 243 waker.wake(); 244 } 245 } 246 247 /// Attempts to take the `Waker` value out of the `AtomicWaker` with the 248 /// intention that the caller will wake the task later. take_waker(&self) -> Option<Waker>249 pub(crate) fn take_waker(&self) -> Option<Waker> { 250 // AcqRel ordering is used in order to acquire the value of the `waker` 251 // cell as well as to establish a `release` ordering with whatever 252 // memory the `AtomicWaker` is associated with. 253 match self.state.fetch_or(WAKING, AcqRel) { 254 WAITING => { 255 // The waking lock has been acquired. 256 let waker = unsafe { self.waker.with_mut(|t| (*t).take()) }; 257 258 // Release the lock 259 self.state.fetch_and(!WAKING, Release); 260 261 waker 262 } 263 state => { 264 // There is a concurrent thread currently updating the 265 // associated waker. 266 // 267 // Nothing more to do as the `WAKING` bit has been set. It 268 // doesn't matter if there are concurrent registering threads or 269 // not. 270 // 271 debug_assert!( 272 state == REGISTERING || state == REGISTERING | WAKING || state == WAKING 273 ); 274 None 275 } 276 } 277 } 278 } 279 280 impl Default for AtomicWaker { default() -> Self281 fn default() -> Self { 282 AtomicWaker::new() 283 } 284 } 285 286 impl fmt::Debug for AtomicWaker { fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result287 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { 288 write!(fmt, "AtomicWaker") 289 } 290 } 291 292 unsafe impl Send for AtomicWaker {} 293 unsafe impl Sync for AtomicWaker {} 294 295 trait WakerRef { wake(self)296 fn wake(self); into_waker(self) -> Waker297 fn into_waker(self) -> Waker; 298 } 299 300 impl WakerRef for Waker { wake(self)301 fn wake(self) { 302 self.wake() 303 } 304 into_waker(self) -> Waker305 fn into_waker(self) -> Waker { 306 self 307 } 308 } 309 310 impl WakerRef for &Waker { wake(self)311 fn wake(self) { 312 self.wake_by_ref() 313 } 314 into_waker(self) -> Waker315 fn into_waker(self) -> Waker { 316 self.clone() 317 } 318 } 319