std/sys/sync/once/queue.rs
1// Each `Once` has one word of atomic state, and this state is CAS'd on to
2// determine what to do. There are four possible state of a `Once`:
3//
4// * Incomplete - no initialization has run yet, and no thread is currently
5// using the Once.
6// * Poisoned - some thread has previously attempted to initialize the Once, but
7// it panicked, so the Once is now poisoned. There are no other
8// threads currently accessing this Once.
9// * Running - some thread is currently attempting to run initialization. It may
10// succeed, so all future threads need to wait for it to finish.
11// Note that this state is accompanied with a payload, described
12// below.
13// * Complete - initialization has completed and all future calls should finish
14// immediately.
15//
16// With 4 states we need 2 bits to encode this, and we use the remaining bits
17// in the word we have allocated as a queue of threads waiting for the thread
18// responsible for entering the RUNNING state. This queue is just a linked list
19// of Waiter nodes which is monotonically increasing in size. Each node is
20// allocated on the stack, and whenever the running closure finishes it will
21// consume the entire queue and notify all waiters they should try again.
22//
23// You'll find a few more details in the implementation, but that's the gist of
24// it!
25//
26// Futex orderings:
27// When running `Once` we deal with multiple atomics:
28// `Once.state_and_queue` and an unknown number of `Waiter.signaled`.
29// * `state_and_queue` is used (1) as a state flag, (2) for synchronizing the
30// result of the `Once`, and (3) for synchronizing `Waiter` nodes.
31// - At the end of the `call` function we have to make sure the result
32// of the `Once` is acquired. So every load which can be the only one to
33// load COMPLETED must have at least acquire ordering, which means all
34// three of them.
35// - `WaiterQueue::drop` is the only place that may store COMPLETED, and
36// must do so with release ordering to make the result available.
37// - `wait` inserts `Waiter` nodes as a pointer in `state_and_queue`, and
38// needs to make the nodes available with release ordering. The load in
39// its `compare_exchange` can be relaxed because it only has to compare
40// the atomic, not to read other data.
41// - `WaiterQueue::drop` must see the `Waiter` nodes, so it must load
42// `state_and_queue` with acquire ordering.
43// - There is just one store where `state_and_queue` is used only as a
44// state flag, without having to synchronize data: switching the state
45// from INCOMPLETE to RUNNING in `call`. This store can be Relaxed,
46// but the read has to be Acquire because of the requirements mentioned
47// above.
48// * `Waiter.signaled` is both used as a flag, and to protect a field with
49// interior mutability in `Waiter`. `Waiter.thread` is changed in
50// `WaiterQueue::drop` which then sets `signaled` with release ordering.
51// After `wait` loads `signaled` with acquire ordering and sees it is true,
52// it needs to see the changes to drop the `Waiter` struct correctly.
53// * There is one place where the two atomics `Once.state_and_queue` and
54// `Waiter.signaled` come together, and might be reordered by the compiler or
55// processor. Because both use acquire ordering such a reordering is not
56// allowed, so no need for `SeqCst`.
57
58use crate::cell::Cell;
59use crate::sync::atomic::Ordering::{AcqRel, Acquire, Release};
60use crate::sync::atomic::{AtomicBool, AtomicPtr};
61use crate::sync::poison::once::ExclusiveState;
62use crate::thread::{self, Thread};
63use crate::{fmt, ptr, sync as public};
64
65type StateAndQueue = *mut ();
66
67pub struct Once {
68 state_and_queue: AtomicPtr<()>,
69}
70
71pub struct OnceState {
72 poisoned: bool,
73 set_state_on_drop_to: Cell<StateAndQueue>,
74}
75
76// Four states that a Once can be in, encoded into the lower bits of
77// `state_and_queue` in the Once structure.
78const INCOMPLETE: usize = 0x0;
79const POISONED: usize = 0x1;
80const RUNNING: usize = 0x2;
81const COMPLETE: usize = 0x3;
82
83// Mask to learn about the state. All other bits are the queue of waiters if
84// this is in the RUNNING state.
85const STATE_MASK: usize = 0b11;
86const QUEUE_MASK: usize = !STATE_MASK;
87
88// Representation of a node in the linked list of waiters, used while in the
89// RUNNING state.
90// Note: `Waiter` can't hold a mutable pointer to the next thread, because then
91// `wait` would both hand out a mutable reference to its `Waiter` node, and keep
92// a shared reference to check `signaled`. Instead we hold shared references and
93// use interior mutability.
94#[repr(align(4))] // Ensure the two lower bits are free to use as state bits.
95struct Waiter {
96 thread: Thread,
97 signaled: AtomicBool,
98 next: Cell<*const Waiter>,
99}
100
101// Head of a linked list of waiters.
102// Every node is a struct on the stack of a waiting thread.
103// Will wake up the waiters when it gets dropped, i.e. also on panic.
104struct WaiterQueue<'a> {
105 state_and_queue: &'a AtomicPtr<()>,
106 set_state_on_drop_to: StateAndQueue,
107}
108
109fn to_queue(current: StateAndQueue) -> *const Waiter {
110 current.mask(QUEUE_MASK).cast()
111}
112
113fn to_state(current: StateAndQueue) -> usize {
114 current.addr() & STATE_MASK
115}
116
117impl Once {
118 #[inline]
119 pub const fn new() -> Once {
120 Once { state_and_queue: AtomicPtr::new(ptr::without_provenance_mut(INCOMPLETE)) }
121 }
122
123 #[inline]
124 pub fn is_completed(&self) -> bool {
125 // An `Acquire` load is enough because that makes all the initialization
126 // operations visible to us, and, this being a fast path, weaker
127 // ordering helps with performance. This `Acquire` synchronizes with
128 // `Release` operations on the slow path.
129 self.state_and_queue.load(Acquire).addr() == COMPLETE
130 }
131
132 #[inline]
133 pub(crate) fn state(&mut self) -> ExclusiveState {
134 match self.state_and_queue.get_mut().addr() {
135 INCOMPLETE => ExclusiveState::Incomplete,
136 POISONED => ExclusiveState::Poisoned,
137 COMPLETE => ExclusiveState::Complete,
138 _ => unreachable!("invalid Once state"),
139 }
140 }
141
142 #[inline]
143 pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
144 *self.state_and_queue.get_mut() = match new_state {
145 ExclusiveState::Incomplete => ptr::without_provenance_mut(INCOMPLETE),
146 ExclusiveState::Poisoned => ptr::without_provenance_mut(POISONED),
147 ExclusiveState::Complete => ptr::without_provenance_mut(COMPLETE),
148 };
149 }
150
151 #[cold]
152 #[track_caller]
153 pub fn wait(&self, ignore_poisoning: bool) {
154 let mut current = self.state_and_queue.load(Acquire);
155 loop {
156 let state = to_state(current);
157 match state {
158 COMPLETE => return,
159 POISONED if !ignore_poisoning => {
160 // Panic to propagate the poison.
161 panic!("Once instance has previously been poisoned");
162 }
163 _ => {
164 current = wait(&self.state_and_queue, current, !ignore_poisoning);
165 }
166 }
167 }
168 }
169
170 // This is a non-generic function to reduce the monomorphization cost of
171 // using `call_once` (this isn't exactly a trivial or small implementation).
172 //
173 // Additionally, this is tagged with `#[cold]` as it should indeed be cold
174 // and it helps let LLVM know that calls to this function should be off the
175 // fast path. Essentially, this should help generate more straight line code
176 // in LLVM.
177 //
178 // Finally, this takes an `FnMut` instead of a `FnOnce` because there's
179 // currently no way to take an `FnOnce` and call it via virtual dispatch
180 // without some allocation overhead.
181 #[cold]
182 #[track_caller]
183 pub fn call(&self, ignore_poisoning: bool, init: &mut dyn FnMut(&public::OnceState)) {
184 let mut current = self.state_and_queue.load(Acquire);
185 loop {
186 let state = to_state(current);
187 match state {
188 COMPLETE => break,
189 POISONED if !ignore_poisoning => {
190 // Panic to propagate the poison.
191 panic!("Once instance has previously been poisoned");
192 }
193 POISONED | INCOMPLETE => {
194 // Try to register this thread as the one RUNNING.
195 if let Err(new) = self.state_and_queue.compare_exchange_weak(
196 current,
197 current.mask(QUEUE_MASK).wrapping_byte_add(RUNNING),
198 Acquire,
199 Acquire,
200 ) {
201 current = new;
202 continue;
203 }
204
205 // `waiter_queue` will manage other waiting threads, and
206 // wake them up on drop.
207 let mut waiter_queue = WaiterQueue {
208 state_and_queue: &self.state_and_queue,
209 set_state_on_drop_to: ptr::without_provenance_mut(POISONED),
210 };
211 // Run the initialization function, letting it know if we're
212 // poisoned or not.
213 let init_state = public::OnceState {
214 inner: OnceState {
215 poisoned: state == POISONED,
216 set_state_on_drop_to: Cell::new(ptr::without_provenance_mut(COMPLETE)),
217 },
218 };
219 init(&init_state);
220 waiter_queue.set_state_on_drop_to = init_state.inner.set_state_on_drop_to.get();
221 return;
222 }
223 _ => {
224 // All other values must be RUNNING with possibly a
225 // pointer to the waiter queue in the more significant bits.
226 assert!(state == RUNNING);
227 current = wait(&self.state_and_queue, current, true);
228 }
229 }
230 }
231 }
232}
233
234fn wait(
235 state_and_queue: &AtomicPtr<()>,
236 mut current: StateAndQueue,
237 return_on_poisoned: bool,
238) -> StateAndQueue {
239 let node = &Waiter {
240 thread: thread::current_or_unnamed(),
241 signaled: AtomicBool::new(false),
242 next: Cell::new(ptr::null()),
243 };
244
245 loop {
246 let state = to_state(current);
247 let queue = to_queue(current);
248
249 // If initialization has finished, return.
250 if state == COMPLETE || (return_on_poisoned && state == POISONED) {
251 return current;
252 }
253
254 // Update the node for our current thread.
255 node.next.set(queue);
256
257 // Try to slide in the node at the head of the linked list, making sure
258 // that another thread didn't just replace the head of the linked list.
259 if let Err(new) = state_and_queue.compare_exchange_weak(
260 current,
261 ptr::from_ref(node).wrapping_byte_add(state) as StateAndQueue,
262 Release,
263 Acquire,
264 ) {
265 current = new;
266 continue;
267 }
268
269 // We have enqueued ourselves, now lets wait.
270 // It is important not to return before being signaled, otherwise we
271 // would drop our `Waiter` node and leave a hole in the linked list
272 // (and a dangling reference). Guard against spurious wakeups by
273 // reparking ourselves until we are signaled.
274 while !node.signaled.load(Acquire) {
275 // If the managing thread happens to signal and unpark us before we
276 // can park ourselves, the result could be this thread never gets
277 // unparked. Luckily `park` comes with the guarantee that if it got
278 // an `unpark` just before on an unparked thread it does not park.
279 // SAFETY: we retrieved this handle on the current thread above.
280 unsafe { node.thread.park() }
281 }
282
283 return state_and_queue.load(Acquire);
284 }
285}
286
287#[stable(feature = "std_debug", since = "1.16.0")]
288impl fmt::Debug for Once {
289 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
290 f.debug_struct("Once").finish_non_exhaustive()
291 }
292}
293
294impl Drop for WaiterQueue<'_> {
295 fn drop(&mut self) {
296 // Swap out our state with however we finished.
297 let current = self.state_and_queue.swap(self.set_state_on_drop_to, AcqRel);
298
299 // We should only ever see an old state which was RUNNING.
300 assert_eq!(current.addr() & STATE_MASK, RUNNING);
301
302 // Walk the entire linked list of waiters and wake them up (in lifo
303 // order, last to register is first to wake up).
304 unsafe {
305 // Right after setting `node.signaled = true` the other thread may
306 // free `node` if there happens to be has a spurious wakeup.
307 // So we have to take out the `thread` field and copy the pointer to
308 // `next` first.
309 let mut queue = to_queue(current);
310 while !queue.is_null() {
311 let next = (*queue).next.get();
312 let thread = (*queue).thread.clone();
313 (*queue).signaled.store(true, Release);
314 thread.unpark();
315 queue = next;
316 }
317 }
318 }
319}
320
321impl OnceState {
322 #[inline]
323 pub fn is_poisoned(&self) -> bool {
324 self.poisoned
325 }
326
327 #[inline]
328 pub fn poison(&self) {
329 self.set_state_on_drop_to.set(ptr::without_provenance_mut(POISONED));
330 }
331}