1 use crate::loom::sync::atomic::AtomicUsize; 2 3 use std::fmt; 4 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; 5 use std::usize; 6 7 pub(super) struct State { 8 val: AtomicUsize, 9 } 10 11 /// Current state value 12 #[derive(Copy, Clone)] 13 pub(super) struct Snapshot(usize); 14 15 type UpdateResult = Result<Snapshot, Snapshot>; 16 17 /// The task is currently being run. 18 const RUNNING: usize = 0b0001; 19 20 /// The task is complete. 21 /// 22 /// Once this bit is set, it is never unset 23 const COMPLETE: usize = 0b0010; 24 25 /// Extracts the task's lifecycle value from the state 26 const LIFECYCLE_MASK: usize = 0b11; 27 28 /// Flag tracking if the task has been pushed into a run queue. 29 const NOTIFIED: usize = 0b100; 30 31 /// The join handle is still around 32 #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 33 const JOIN_INTEREST: usize = 0b1_000; 34 35 /// A join handle waker has been set 36 #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 37 const JOIN_WAKER: usize = 0b10_000; 38 39 /// The task has been forcibly cancelled. 40 #[allow(clippy::unusual_byte_groupings)] // https://github.com/rust-lang/rust-clippy/issues/6556 41 const CANCELLED: usize = 0b100_000; 42 43 /// All bits 44 const STATE_MASK: usize = LIFECYCLE_MASK | NOTIFIED | JOIN_INTEREST | JOIN_WAKER | CANCELLED; 45 46 /// Bits used by the ref count portion of the state. 47 const REF_COUNT_MASK: usize = !STATE_MASK; 48 49 /// Number of positions to shift the ref count 50 const REF_COUNT_SHIFT: usize = REF_COUNT_MASK.count_zeros() as usize; 51 52 /// One ref count 53 const REF_ONE: usize = 1 << REF_COUNT_SHIFT; 54 55 /// State a task is initialized with 56 /// 57 /// A task is initialized with three references: 58 /// 59 /// * A reference that will be stored in an OwnedTasks or LocalOwnedTasks. 60 /// * A reference that will be sent to the scheduler as an ordinary notification. 61 /// * A reference for the JoinHandle. 62 /// 63 /// As the task starts with a `JoinHandle`, `JOIN_INTEREST` is set. 64 /// As the task starts with a `Notified`, `NOTIFIED` is set. 65 const INITIAL_STATE: usize = (REF_ONE * 3) | JOIN_INTEREST | NOTIFIED; 66 67 #[must_use] 68 pub(super) enum TransitionToRunning { 69 Success, 70 Cancelled, 71 Failed, 72 Dealloc, 73 } 74 75 #[must_use] 76 pub(super) enum TransitionToIdle { 77 Ok, 78 OkNotified, 79 OkDealloc, 80 Cancelled, 81 } 82 83 #[must_use] 84 pub(super) enum TransitionToNotifiedByVal { 85 DoNothing, 86 Submit, 87 Dealloc, 88 } 89 90 #[must_use] 91 pub(super) enum TransitionToNotifiedByRef { 92 DoNothing, 93 Submit, 94 } 95 96 /// All transitions are performed via RMW operations. This establishes an 97 /// unambiguous modification order. 98 impl State { 99 /// Return a task's initial state new() -> State100 pub(super) fn new() -> State { 101 // The raw task returned by this method has a ref-count of three. See 102 // the comment on INITIAL_STATE for more. 103 State { 104 val: AtomicUsize::new(INITIAL_STATE), 105 } 106 } 107 108 /// Loads the current state, establishes `Acquire` ordering. load(&self) -> Snapshot109 pub(super) fn load(&self) -> Snapshot { 110 Snapshot(self.val.load(Acquire)) 111 } 112 113 /// Attempt to transition the lifecycle to `Running`. This sets the 114 /// notified bit to false so notifications during the poll can be detected. transition_to_running(&self) -> TransitionToRunning115 pub(super) fn transition_to_running(&self) -> TransitionToRunning { 116 self.fetch_update_action(|mut next| { 117 let action; 118 assert!(next.is_notified()); 119 120 if !next.is_idle() { 121 // This happens if the task is either currently running or if it 122 // has already completed, e.g. if it was cancelled during 123 // shutdown. Consume the ref-count and return. 124 next.ref_dec(); 125 if next.ref_count() == 0 { 126 action = TransitionToRunning::Dealloc; 127 } else { 128 action = TransitionToRunning::Failed; 129 } 130 } else { 131 // We are able to lock the RUNNING bit. 132 next.set_running(); 133 next.unset_notified(); 134 135 if next.is_cancelled() { 136 action = TransitionToRunning::Cancelled; 137 } else { 138 action = TransitionToRunning::Success; 139 } 140 } 141 (action, Some(next)) 142 }) 143 } 144 145 /// Transitions the task from `Running` -> `Idle`. 146 /// 147 /// Returns `true` if the transition to `Idle` is successful, `false` otherwise. 148 /// The transition to `Idle` fails if the task has been flagged to be 149 /// cancelled. transition_to_idle(&self) -> TransitionToIdle150 pub(super) fn transition_to_idle(&self) -> TransitionToIdle { 151 self.fetch_update_action(|curr| { 152 assert!(curr.is_running()); 153 154 if curr.is_cancelled() { 155 return (TransitionToIdle::Cancelled, None); 156 } 157 158 let mut next = curr; 159 let action; 160 next.unset_running(); 161 162 if !next.is_notified() { 163 // Polling the future consumes the ref-count of the Notified. 164 next.ref_dec(); 165 if next.ref_count() == 0 { 166 action = TransitionToIdle::OkDealloc; 167 } else { 168 action = TransitionToIdle::Ok; 169 } 170 } else { 171 // The caller will schedule a new notification, so we create a 172 // new ref-count for the notification. Our own ref-count is kept 173 // for now, and the caller will drop it shortly. 174 next.ref_inc(); 175 action = TransitionToIdle::OkNotified; 176 } 177 178 (action, Some(next)) 179 }) 180 } 181 182 /// Transitions the task from `Running` -> `Complete`. transition_to_complete(&self) -> Snapshot183 pub(super) fn transition_to_complete(&self) -> Snapshot { 184 const DELTA: usize = RUNNING | COMPLETE; 185 186 let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); 187 assert!(prev.is_running()); 188 assert!(!prev.is_complete()); 189 190 Snapshot(prev.0 ^ DELTA) 191 } 192 193 /// Transition from `Complete` -> `Terminal`, decrementing the reference 194 /// count the specified number of times. 195 /// 196 /// Returns true if the task should be deallocated. transition_to_terminal(&self, count: usize) -> bool197 pub(super) fn transition_to_terminal(&self, count: usize) -> bool { 198 let prev = Snapshot(self.val.fetch_sub(count * REF_ONE, AcqRel)); 199 assert!( 200 prev.ref_count() >= count, 201 "current: {}, sub: {}", 202 prev.ref_count(), 203 count 204 ); 205 prev.ref_count() == count 206 } 207 208 /// Transitions the state to `NOTIFIED`. 209 /// 210 /// If no task needs to be submitted, a ref-count is consumed. 211 /// 212 /// If a task needs to be submitted, the ref-count is incremented for the 213 /// new Notified. transition_to_notified_by_val(&self) -> TransitionToNotifiedByVal214 pub(super) fn transition_to_notified_by_val(&self) -> TransitionToNotifiedByVal { 215 self.fetch_update_action(|mut snapshot| { 216 let action; 217 218 if snapshot.is_running() { 219 // If the task is running, we mark it as notified, but we should 220 // not submit anything as the thread currently running the 221 // future is responsible for that. 222 snapshot.set_notified(); 223 snapshot.ref_dec(); 224 225 // The thread that set the running bit also holds a ref-count. 226 assert!(snapshot.ref_count() > 0); 227 228 action = TransitionToNotifiedByVal::DoNothing; 229 } else if snapshot.is_complete() || snapshot.is_notified() { 230 // We do not need to submit any notifications, but we have to 231 // decrement the ref-count. 232 snapshot.ref_dec(); 233 234 if snapshot.ref_count() == 0 { 235 action = TransitionToNotifiedByVal::Dealloc; 236 } else { 237 action = TransitionToNotifiedByVal::DoNothing; 238 } 239 } else { 240 // We create a new notified that we can submit. The caller 241 // retains ownership of the ref-count they passed in. 242 snapshot.set_notified(); 243 snapshot.ref_inc(); 244 action = TransitionToNotifiedByVal::Submit; 245 } 246 247 (action, Some(snapshot)) 248 }) 249 } 250 251 /// Transitions the state to `NOTIFIED`. transition_to_notified_by_ref(&self) -> TransitionToNotifiedByRef252 pub(super) fn transition_to_notified_by_ref(&self) -> TransitionToNotifiedByRef { 253 self.fetch_update_action(|mut snapshot| { 254 if snapshot.is_complete() || snapshot.is_notified() { 255 // There is nothing to do in this case. 256 (TransitionToNotifiedByRef::DoNothing, None) 257 } else if snapshot.is_running() { 258 // If the task is running, we mark it as notified, but we should 259 // not submit as the thread currently running the future is 260 // responsible for that. 261 snapshot.set_notified(); 262 (TransitionToNotifiedByRef::DoNothing, Some(snapshot)) 263 } else { 264 // The task is idle and not notified. We should submit a 265 // notification. 266 snapshot.set_notified(); 267 snapshot.ref_inc(); 268 (TransitionToNotifiedByRef::Submit, Some(snapshot)) 269 } 270 }) 271 } 272 273 /// Set the cancelled bit and transition the state to `NOTIFIED` if idle. 274 /// 275 /// Returns `true` if the task needs to be submitted to the pool for 276 /// execution transition_to_notified_and_cancel(&self) -> bool277 pub(super) fn transition_to_notified_and_cancel(&self) -> bool { 278 self.fetch_update_action(|mut snapshot| { 279 if snapshot.is_cancelled() || snapshot.is_complete() { 280 // Aborts to completed or cancelled tasks are no-ops. 281 (false, None) 282 } else if snapshot.is_running() { 283 // If the task is running, we mark it as cancelled. The thread 284 // running the task will notice the cancelled bit when it 285 // stops polling and it will kill the task. 286 // 287 // The set_notified() call is not strictly necessary but it will 288 // in some cases let a wake_by_ref call return without having 289 // to perform a compare_exchange. 290 snapshot.set_notified(); 291 snapshot.set_cancelled(); 292 (false, Some(snapshot)) 293 } else { 294 // The task is idle. We set the cancelled and notified bits and 295 // submit a notification if the notified bit was not already 296 // set. 297 snapshot.set_cancelled(); 298 if !snapshot.is_notified() { 299 snapshot.set_notified(); 300 snapshot.ref_inc(); 301 (true, Some(snapshot)) 302 } else { 303 (false, Some(snapshot)) 304 } 305 } 306 }) 307 } 308 309 /// Set the `CANCELLED` bit and attempt to transition to `Running`. 310 /// 311 /// Returns `true` if the transition to `Running` succeeded. transition_to_shutdown(&self) -> bool312 pub(super) fn transition_to_shutdown(&self) -> bool { 313 let mut prev = Snapshot(0); 314 315 let _ = self.fetch_update(|mut snapshot| { 316 prev = snapshot; 317 318 if snapshot.is_idle() { 319 snapshot.set_running(); 320 } 321 322 // If the task was not idle, the thread currently running the task 323 // will notice the cancelled bit and cancel it once the poll 324 // completes. 325 snapshot.set_cancelled(); 326 Some(snapshot) 327 }); 328 329 prev.is_idle() 330 } 331 332 /// Optimistically tries to swap the state assuming the join handle is 333 /// __immediately__ dropped on spawn drop_join_handle_fast(&self) -> Result<(), ()>334 pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> { 335 use std::sync::atomic::Ordering::Relaxed; 336 337 // Relaxed is acceptable as if this function is called and succeeds, 338 // then nothing has been done w/ the join handle. 339 // 340 // The moment the join handle is used (polled), the `JOIN_WAKER` flag is 341 // set, at which point the CAS will fail. 342 // 343 // Given this, there is no risk if this operation is reordered. 344 self.val 345 .compare_exchange_weak( 346 INITIAL_STATE, 347 (INITIAL_STATE - REF_ONE) & !JOIN_INTEREST, 348 Release, 349 Relaxed, 350 ) 351 .map(|_| ()) 352 .map_err(|_| ()) 353 } 354 355 /// Try to unset the JOIN_INTEREST flag. 356 /// 357 /// Returns `Ok` if the operation happens before the task transitions to a 358 /// completed state, `Err` otherwise. unset_join_interested(&self) -> UpdateResult359 pub(super) fn unset_join_interested(&self) -> UpdateResult { 360 self.fetch_update(|curr| { 361 assert!(curr.is_join_interested()); 362 363 if curr.is_complete() { 364 return None; 365 } 366 367 let mut next = curr; 368 next.unset_join_interested(); 369 370 Some(next) 371 }) 372 } 373 374 /// Set the `JOIN_WAKER` bit. 375 /// 376 /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if 377 /// the task has completed. set_join_waker(&self) -> UpdateResult378 pub(super) fn set_join_waker(&self) -> UpdateResult { 379 self.fetch_update(|curr| { 380 assert!(curr.is_join_interested()); 381 assert!(!curr.has_join_waker()); 382 383 if curr.is_complete() { 384 return None; 385 } 386 387 let mut next = curr; 388 next.set_join_waker(); 389 390 Some(next) 391 }) 392 } 393 394 /// Unsets the `JOIN_WAKER` bit. 395 /// 396 /// Returns `Ok` has been unset, `Err` otherwise. This operation fails if 397 /// the task has completed. unset_waker(&self) -> UpdateResult398 pub(super) fn unset_waker(&self) -> UpdateResult { 399 self.fetch_update(|curr| { 400 assert!(curr.is_join_interested()); 401 assert!(curr.has_join_waker()); 402 403 if curr.is_complete() { 404 return None; 405 } 406 407 let mut next = curr; 408 next.unset_join_waker(); 409 410 Some(next) 411 }) 412 } 413 ref_inc(&self)414 pub(super) fn ref_inc(&self) { 415 use std::process; 416 use std::sync::atomic::Ordering::Relaxed; 417 418 // Using a relaxed ordering is alright here, as knowledge of the 419 // original reference prevents other threads from erroneously deleting 420 // the object. 421 // 422 // As explained in the [Boost documentation][1], Increasing the 423 // reference counter can always be done with memory_order_relaxed: New 424 // references to an object can only be formed from an existing 425 // reference, and passing an existing reference from one thread to 426 // another must already provide any required synchronization. 427 // 428 // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) 429 let prev = self.val.fetch_add(REF_ONE, Relaxed); 430 431 // If the reference count overflowed, abort. 432 if prev > isize::MAX as usize { 433 process::abort(); 434 } 435 } 436 437 /// Returns `true` if the task should be released. ref_dec(&self) -> bool438 pub(super) fn ref_dec(&self) -> bool { 439 let prev = Snapshot(self.val.fetch_sub(REF_ONE, AcqRel)); 440 assert!(prev.ref_count() >= 1); 441 prev.ref_count() == 1 442 } 443 444 /// Returns `true` if the task should be released. ref_dec_twice(&self) -> bool445 pub(super) fn ref_dec_twice(&self) -> bool { 446 let prev = Snapshot(self.val.fetch_sub(2 * REF_ONE, AcqRel)); 447 assert!(prev.ref_count() >= 2); 448 prev.ref_count() == 2 449 } 450 fetch_update_action<F, T>(&self, mut f: F) -> T where F: FnMut(Snapshot) -> (T, Option<Snapshot>),451 fn fetch_update_action<F, T>(&self, mut f: F) -> T 452 where 453 F: FnMut(Snapshot) -> (T, Option<Snapshot>), 454 { 455 let mut curr = self.load(); 456 457 loop { 458 let (output, next) = f(curr); 459 let next = match next { 460 Some(next) => next, 461 None => return output, 462 }; 463 464 let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); 465 466 match res { 467 Ok(_) => return output, 468 Err(actual) => curr = Snapshot(actual), 469 } 470 } 471 } 472 fetch_update<F>(&self, mut f: F) -> Result<Snapshot, Snapshot> where F: FnMut(Snapshot) -> Option<Snapshot>,473 fn fetch_update<F>(&self, mut f: F) -> Result<Snapshot, Snapshot> 474 where 475 F: FnMut(Snapshot) -> Option<Snapshot>, 476 { 477 let mut curr = self.load(); 478 479 loop { 480 let next = match f(curr) { 481 Some(next) => next, 482 None => return Err(curr), 483 }; 484 485 let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); 486 487 match res { 488 Ok(_) => return Ok(next), 489 Err(actual) => curr = Snapshot(actual), 490 } 491 } 492 } 493 } 494 495 // ===== impl Snapshot ===== 496 497 impl Snapshot { 498 /// Returns `true` if the task is in an idle state. is_idle(self) -> bool499 pub(super) fn is_idle(self) -> bool { 500 self.0 & (RUNNING | COMPLETE) == 0 501 } 502 503 /// Returns `true` if the task has been flagged as notified. is_notified(self) -> bool504 pub(super) fn is_notified(self) -> bool { 505 self.0 & NOTIFIED == NOTIFIED 506 } 507 unset_notified(&mut self)508 fn unset_notified(&mut self) { 509 self.0 &= !NOTIFIED 510 } 511 set_notified(&mut self)512 fn set_notified(&mut self) { 513 self.0 |= NOTIFIED 514 } 515 is_running(self) -> bool516 pub(super) fn is_running(self) -> bool { 517 self.0 & RUNNING == RUNNING 518 } 519 set_running(&mut self)520 fn set_running(&mut self) { 521 self.0 |= RUNNING; 522 } 523 unset_running(&mut self)524 fn unset_running(&mut self) { 525 self.0 &= !RUNNING; 526 } 527 is_cancelled(self) -> bool528 pub(super) fn is_cancelled(self) -> bool { 529 self.0 & CANCELLED == CANCELLED 530 } 531 set_cancelled(&mut self)532 fn set_cancelled(&mut self) { 533 self.0 |= CANCELLED; 534 } 535 536 /// Returns `true` if the task's future has completed execution. is_complete(self) -> bool537 pub(super) fn is_complete(self) -> bool { 538 self.0 & COMPLETE == COMPLETE 539 } 540 is_join_interested(self) -> bool541 pub(super) fn is_join_interested(self) -> bool { 542 self.0 & JOIN_INTEREST == JOIN_INTEREST 543 } 544 unset_join_interested(&mut self)545 fn unset_join_interested(&mut self) { 546 self.0 &= !JOIN_INTEREST 547 } 548 has_join_waker(self) -> bool549 pub(super) fn has_join_waker(self) -> bool { 550 self.0 & JOIN_WAKER == JOIN_WAKER 551 } 552 set_join_waker(&mut self)553 fn set_join_waker(&mut self) { 554 self.0 |= JOIN_WAKER; 555 } 556 unset_join_waker(&mut self)557 fn unset_join_waker(&mut self) { 558 self.0 &= !JOIN_WAKER 559 } 560 ref_count(self) -> usize561 pub(super) fn ref_count(self) -> usize { 562 (self.0 & REF_COUNT_MASK) >> REF_COUNT_SHIFT 563 } 564 ref_inc(&mut self)565 fn ref_inc(&mut self) { 566 assert!(self.0 <= isize::MAX as usize); 567 self.0 += REF_ONE; 568 } 569 ref_dec(&mut self)570 pub(super) fn ref_dec(&mut self) { 571 assert!(self.ref_count() > 0); 572 self.0 -= REF_ONE 573 } 574 } 575 576 impl fmt::Debug for State { fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result577 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { 578 let snapshot = self.load(); 579 snapshot.fmt(fmt) 580 } 581 } 582 583 impl fmt::Debug for Snapshot { fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result584 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { 585 fmt.debug_struct("Snapshot") 586 .field("is_running", &self.is_running()) 587 .field("is_complete", &self.is_complete()) 588 .field("is_notified", &self.is_notified()) 589 .field("is_cancelled", &self.is_cancelled()) 590 .field("is_join_interested", &self.is_join_interested()) 591 .field("has_join_waker", &self.has_join_waker()) 592 .field("ref_count", &self.ref_count()) 593 .finish() 594 } 595 } 596