1use crate::cell::Cell;
2use crate::sync as public;
3use crate::sync::atomic::Ordering::{Acquire, Relaxed, Release};
4use crate::sync::poison::once::ExclusiveState;
5use crate::sys::futex::{Futex, Primitive, futex_wait, futex_wake_all};
67// On some platforms, the OS is very nice and handles the waiter queue for us.
8// This means we only need one atomic value with 4 states:
910/// No initialization has run yet, and no thread is currently using the Once.
11const INCOMPLETE: Primitive = 3;
12/// Some thread has previously attempted to initialize the Once, but it panicked,
13/// so the Once is now poisoned. There are no other threads currently accessing
14/// this Once.
15const POISONED: Primitive = 2;
16/// Some thread is currently attempting to run initialization. It may succeed,
17/// so all future threads need to wait for it to finish.
18const RUNNING: Primitive = 1;
19/// Initialization has completed and all future calls should finish immediately.
20/// By choosing this state as the all-zero state the `is_completed` check can be
21/// a bit faster on some platforms.
22const COMPLETE: Primitive = 0;
2324// An additional bit indicates whether there are waiting threads:
2526/// May only be set if the state is not COMPLETE.
27const QUEUED: Primitive = 4;
2829// Threads wait by setting the QUEUED bit and calling `futex_wait` on the state
30// variable. When the running thread finishes, it will wake all waiting threads using
31// `futex_wake_all`.
3233const STATE_MASK: Primitive = 0b11;
3435pub struct OnceState {
36 poisoned: bool,
37 set_state_to: Cell<Primitive>,
38}
3940impl OnceState {
41#[inline]
42pub fn is_poisoned(&self) -> bool {
43self.poisoned
44 }
4546#[inline]
47pub fn poison(&self) {
48self.set_state_to.set(POISONED);
49 }
50}
5152struct CompletionGuard<'a> {
53 state_and_queued: &'a Futex,
54 set_state_on_drop_to: Primitive,
55}
5657impl<'a> Drop for CompletionGuard<'a> {
58fn drop(&mut self) {
59// Use release ordering to propagate changes to all threads checking
60 // up on the Once. `futex_wake_all` does its own synchronization, hence
61 // we do not need `AcqRel`.
62if self.state_and_queued.swap(self.set_state_on_drop_to, Release) & QUEUED != 0 {
63 futex_wake_all(self.state_and_queued);
64 }
65 }
66}
6768pub struct Once {
69 state_and_queued: Futex,
70}
7172impl Once {
73#[inline]
74pub const fn new() -> Once {
75 Once { state_and_queued: Futex::new(INCOMPLETE) }
76 }
7778#[inline]
79pub fn is_completed(&self) -> bool {
80// Use acquire ordering to make all initialization changes visible to the
81 // current thread.
82self.state_and_queued.load(Acquire) == COMPLETE
83 }
8485#[inline]
86pub(crate) fn state(&mut self) -> ExclusiveState {
87match *self.state_and_queued.get_mut() {
88 INCOMPLETE => ExclusiveState::Incomplete,
89 POISONED => ExclusiveState::Poisoned,
90 COMPLETE => ExclusiveState::Complete,
91_ => unreachable!("invalid Once state"),
92 }
93 }
9495#[inline]
96pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
97*self.state_and_queued.get_mut() = match new_state {
98 ExclusiveState::Incomplete => INCOMPLETE,
99 ExclusiveState::Poisoned => POISONED,
100 ExclusiveState::Complete => COMPLETE,
101 };
102 }
103104#[cold]
105 #[track_caller]
106pub fn wait(&self, ignore_poisoning: bool) {
107let mut state_and_queued = self.state_and_queued.load(Acquire);
108loop {
109let state = state_and_queued & STATE_MASK;
110let queued = state_and_queued & QUEUED != 0;
111match state {
112 COMPLETE => return,
113 POISONED if !ignore_poisoning => {
114// Panic to propagate the poison.
115panic!("Once instance has previously been poisoned");
116 }
117_ => {
118// Set the QUEUED bit if it has not already been set.
119if !queued {
120 state_and_queued += QUEUED;
121if let Err(new) = self.state_and_queued.compare_exchange_weak(
122 state,
123 state_and_queued,
124 Relaxed,
125 Acquire,
126 ) {
127 state_and_queued = new;
128continue;
129 }
130 }
131132 futex_wait(&self.state_and_queued, state_and_queued, None);
133 state_and_queued = self.state_and_queued.load(Acquire);
134 }
135 }
136 }
137 }
138139#[cold]
140 #[track_caller]
141pub fn call(&self, ignore_poisoning: bool, f: &mut dyn FnMut(&public::OnceState)) {
142let mut state_and_queued = self.state_and_queued.load(Acquire);
143loop {
144let state = state_and_queued & STATE_MASK;
145let queued = state_and_queued & QUEUED != 0;
146match state {
147 COMPLETE => return,
148 POISONED if !ignore_poisoning => {
149// Panic to propagate the poison.
150panic!("Once instance has previously been poisoned");
151 }
152 INCOMPLETE | POISONED => {
153// Try to register the current thread as the one running.
154let next = RUNNING + if queued { QUEUED } else { 0 };
155if let Err(new) = self.state_and_queued.compare_exchange_weak(
156 state_and_queued,
157 next,
158 Acquire,
159 Acquire,
160 ) {
161 state_and_queued = new;
162continue;
163 }
164165// `waiter_queue` will manage other waiting threads, and
166 // wake them up on drop.
167let mut waiter_queue = CompletionGuard {
168 state_and_queued: &self.state_and_queued,
169 set_state_on_drop_to: POISONED,
170 };
171// Run the function, letting it know if we're poisoned or not.
172let f_state = public::OnceState {
173 inner: OnceState {
174 poisoned: state == POISONED,
175 set_state_to: Cell::new(COMPLETE),
176 },
177 };
178 f(&f_state);
179 waiter_queue.set_state_on_drop_to = f_state.inner.set_state_to.get();
180return;
181 }
182_ => {
183// All other values must be RUNNING.
184assert!(state == RUNNING);
185186// Set the QUEUED bit if it is not already set.
187if !queued {
188 state_and_queued += QUEUED;
189if let Err(new) = self.state_and_queued.compare_exchange_weak(
190 state,
191 state_and_queued,
192 Relaxed,
193 Acquire,
194 ) {
195 state_and_queued = new;
196continue;
197 }
198 }
199200 futex_wait(&self.state_and_queued, state_and_queued, None);
201 state_and_queued = self.state_and_queued.load(Acquire);
202 }
203 }
204 }
205 }
206}