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