Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 1 | use core::alloc::Layout; |
| 2 | use core::cell::UnsafeCell; |
| 3 | use core::fmt; |
| 4 | use core::sync::atomic::{AtomicUsize, Ordering}; |
| 5 | use core::task::Waker; |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 6 | |
| 7 | use crate::raw::TaskVTable; |
| 8 | use crate::state::*; |
| 9 | use crate::utils::{abort_on_panic, extend}; |
| 10 | |
| 11 | /// The header of a task. |
| 12 | /// |
| 13 | /// This header is stored right at the beginning of every heap-allocated task. |
| 14 | pub(crate) struct Header { |
| 15 | /// Current state of the task. |
| 16 | /// |
| 17 | /// Contains flags representing the current state and the reference count. |
| 18 | pub(crate) state: AtomicUsize, |
| 19 | |
| 20 | /// The task that is blocked on the `JoinHandle`. |
| 21 | /// |
Stjepan Glavina | 5c398cf | 2019-08-20 15:29:43 +0200 | [diff] [blame] | 22 | /// This waker needs to be woken up once the task completes or is closed. |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 23 | pub(crate) awaiter: UnsafeCell<Option<Waker>>, |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 24 | |
| 25 | /// The virtual table. |
| 26 | /// |
| 27 | /// In addition to the actual waker virtual table, it also contains pointers to several other |
| 28 | /// methods necessary for bookkeeping the heap-allocated task. |
| 29 | pub(crate) vtable: &'static TaskVTable, |
| 30 | } |
| 31 | |
| 32 | impl Header { |
| 33 | /// Cancels the task. |
| 34 | /// |
Stjepan Glavina | 5c398cf | 2019-08-20 15:29:43 +0200 | [diff] [blame] | 35 | /// This method will mark the task as closed and notify the awaiter, but it won't reschedule |
| 36 | /// the task if it's not completed. |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 37 | pub(crate) fn cancel(&self) { |
| 38 | let mut state = self.state.load(Ordering::Acquire); |
| 39 | |
| 40 | loop { |
| 41 | // If the task has been completed or closed, it can't be cancelled. |
| 42 | if state & (COMPLETED | CLOSED) != 0 { |
| 43 | break; |
| 44 | } |
| 45 | |
| 46 | // Mark the task as closed. |
| 47 | match self.state.compare_exchange_weak( |
| 48 | state, |
| 49 | state | CLOSED, |
| 50 | Ordering::AcqRel, |
| 51 | Ordering::Acquire, |
| 52 | ) { |
| 53 | Ok(_) => { |
| 54 | // Notify the awaiter that the task has been closed. |
| 55 | if state & AWAITER != 0 { |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 56 | self.notify(None); |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 57 | } |
| 58 | |
| 59 | break; |
| 60 | } |
| 61 | Err(s) => state = s, |
| 62 | } |
| 63 | } |
| 64 | } |
| 65 | |
Stjepan Glavina | 5c398cf | 2019-08-20 15:29:43 +0200 | [diff] [blame] | 66 | /// Notifies the awaiter blocked on this task. |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 67 | /// |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 68 | /// If the awaiter is the same as the current waker, it will not be notified. |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 69 | #[inline] |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 70 | pub(crate) fn notify(&self, current: Option<&Waker>) { |
| 71 | // Mark the awaiter as being notified. |
| 72 | let state = self.state.fetch_or(NOTIFYING, Ordering::AcqRel); |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 73 | |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 74 | // If the awaiter was not being notified nor registered... |
| 75 | if state & (NOTIFYING | REGISTERING) == 0 { |
| 76 | // Take the waker out. |
| 77 | let waker = unsafe { (*self.awaiter.get()).take() }; |
| 78 | |
| 79 | // Mark the state as not being notified anymore nor containing an awaiter. |
| 80 | self.state |
| 81 | .fetch_and(!NOTIFYING & !AWAITER, Ordering::Release); |
| 82 | |
| 83 | if let Some(w) = waker { |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 84 | // We need a safeguard against panics because waking can panic. |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 85 | abort_on_panic(|| match current { |
| 86 | None => w.wake(), |
| 87 | Some(c) if !w.will_wake(c) => w.wake(), |
| 88 | Some(_) => {} |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 89 | }); |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 94 | /// Registers a new awaiter blocked on this task. |
| 95 | /// |
| 96 | /// This method is called when `JoinHandle` is polled and the task has not completed. |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 97 | #[inline] |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 98 | pub(crate) fn register(&self, waker: &Waker) { |
| 99 | // Load the state and synchronize with it. |
| 100 | let mut state = self.state.fetch_or(0, Ordering::Acquire); |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 101 | |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 102 | loop { |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 103 | // There can't be two concurrent registrations because `JoinHandle` can only be polled |
| 104 | // by a unique pinned reference. |
| 105 | debug_assert!(state & REGISTERING == 0); |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 106 | |
Stjepan Glavina | f4ae6f5 | 2020-01-06 21:49:10 +0100 | [diff] [blame] | 107 | // If we're in the notifying state at this moment, just wake and return without |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 108 | // registering. |
| 109 | if state & NOTIFYING != 0 { |
| 110 | waker.wake_by_ref(); |
| 111 | return; |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 112 | } |
| 113 | |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 114 | // Mark the state to let other threads know we're registering a new awaiter. |
| 115 | match self.state.compare_exchange_weak( |
| 116 | state, |
| 117 | state | REGISTERING, |
| 118 | Ordering::AcqRel, |
| 119 | Ordering::Acquire, |
| 120 | ) { |
| 121 | Ok(_) => { |
| 122 | state |= REGISTERING; |
| 123 | break; |
| 124 | } |
| 125 | Err(s) => state = s, |
| 126 | } |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 127 | } |
| 128 | |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 129 | // Put the waker into the awaiter field. |
| 130 | unsafe { |
| 131 | abort_on_panic(|| (*self.awaiter.get()) = Some(waker.clone())); |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 132 | } |
| 133 | |
Stjepan Glavina | 921e8a0 | 2020-01-06 14:31:28 -0600 | [diff] [blame] | 134 | // This variable will contain the newly registered waker if a notification comes in before |
| 135 | // we complete registration. |
| 136 | let mut waker = None; |
| 137 | |
| 138 | loop { |
| 139 | // If there was a notification, take the waker out of the awaiter field. |
| 140 | if state & NOTIFYING != 0 { |
| 141 | if let Some(w) = unsafe { (*self.awaiter.get()).take() } { |
| 142 | waker = Some(w); |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | // The new state is not being notified nor registered, but there might or might not be |
| 147 | // an awaiter depending on whether there was a concurrent notification. |
| 148 | let new = if waker.is_none() { |
| 149 | (state & !NOTIFYING & !REGISTERING) | AWAITER |
| 150 | } else { |
| 151 | state & !NOTIFYING & !REGISTERING & !AWAITER |
| 152 | }; |
| 153 | |
| 154 | match self |
| 155 | .state |
| 156 | .compare_exchange_weak(state, new, Ordering::AcqRel, Ordering::Acquire) |
| 157 | { |
| 158 | Ok(_) => break, |
| 159 | Err(s) => state = s, |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // If there was a notification during registration, wake the awaiter now. |
| 164 | if let Some(w) = waker { |
| 165 | abort_on_panic(|| w.wake()); |
| 166 | } |
Stjepan Glavina | 1479e86 | 2019-08-12 20:18:51 +0200 | [diff] [blame] | 167 | } |
| 168 | |
| 169 | /// Returns the offset at which the tag of type `T` is stored. |
| 170 | #[inline] |
| 171 | pub(crate) fn offset_tag<T>() -> usize { |
| 172 | let layout_header = Layout::new::<Header>(); |
| 173 | let layout_t = Layout::new::<T>(); |
| 174 | let (_, offset_t) = extend(layout_header, layout_t); |
| 175 | offset_t |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | impl fmt::Debug for Header { |
| 180 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 181 | let state = self.state.load(Ordering::SeqCst); |
| 182 | |
| 183 | f.debug_struct("Header") |
| 184 | .field("scheduled", &(state & SCHEDULED != 0)) |
| 185 | .field("running", &(state & RUNNING != 0)) |
| 186 | .field("completed", &(state & COMPLETED != 0)) |
| 187 | .field("closed", &(state & CLOSED != 0)) |
| 188 | .field("awaiter", &(state & AWAITER != 0)) |
| 189 | .field("handle", &(state & HANDLE != 0)) |
| 190 | .field("ref_count", &(state / REFERENCE)) |
| 191 | .finish() |
| 192 | } |
| 193 | } |