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 = 0;
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 = 1;
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 = 2;
19/// Initialization has completed and all future calls should finish immediately.
20const COMPLETE: Primitive = 3;
2122// An additional bit indicates whether there are waiting threads:
2324/// May only be set if the state is not COMPLETE.
25const QUEUED: Primitive = 4;
2627// Threads wait by setting the QUEUED bit and calling `futex_wait` on the state
28// variable. When the running thread finishes, it will wake all waiting threads using
29// `futex_wake_all`.
3031const STATE_MASK: Primitive = 0b11;
3233pub struct OnceState {
34 poisoned: bool,
35 set_state_to: Cell<Primitive>,
36}
3738impl OnceState {
39#[inline]
40pub fn is_poisoned(&self) -> bool {
41self.poisoned
42 }
4344#[inline]
45pub fn poison(&self) {
46self.set_state_to.set(POISONED);
47 }
48}
4950struct CompletionGuard<'a> {
51 state_and_queued: &'a Futex,
52 set_state_on_drop_to: Primitive,
53}
5455impl<'a> Drop for CompletionGuard<'a> {
56fn drop(&mut self) {
57// Use release ordering to propagate changes to all threads checking
58 // up on the Once. `futex_wake_all` does its own synchronization, hence
59 // we do not need `AcqRel`.
60if self.state_and_queued.swap(self.set_state_on_drop_to, Release) & QUEUED != 0 {
61 futex_wake_all(self.state_and_queued);
62 }
63 }
64}
6566pub struct Once {
67 state_and_queued: Futex,
68}
6970impl Once {
71#[inline]
72pub const fn new() -> Once {
73 Once { state_and_queued: Futex::new(INCOMPLETE) }
74 }
7576#[inline]
77pub fn is_completed(&self) -> bool {
78// Use acquire ordering to make all initialization changes visible to the
79 // current thread.
80self.state_and_queued.load(Acquire) == COMPLETE
81 }
8283#[inline]
84pub(crate) fn state(&mut self) -> ExclusiveState {
85match *self.state_and_queued.get_mut() {
86 INCOMPLETE => ExclusiveState::Incomplete,
87 POISONED => ExclusiveState::Poisoned,
88 COMPLETE => ExclusiveState::Complete,
89_ => unreachable!("invalid Once state"),
90 }
91 }
9293#[inline]
94pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
95*self.state_and_queued.get_mut() = match new_state {
96 ExclusiveState::Incomplete => INCOMPLETE,
97 ExclusiveState::Poisoned => POISONED,
98 ExclusiveState::Complete => COMPLETE,
99 };
100 }
101102#[cold]
103 #[track_caller]
104pub fn wait(&self, ignore_poisoning: bool) {
105let mut state_and_queued = self.state_and_queued.load(Acquire);
106loop {
107let state = state_and_queued & STATE_MASK;
108let queued = state_and_queued & QUEUED != 0;
109match state {
110 COMPLETE => return,
111 POISONED if !ignore_poisoning => {
112// Panic to propagate the poison.
113panic!("Once instance has previously been poisoned");
114 }
115_ => {
116// Set the QUEUED bit if it has not already been set.
117if !queued {
118 state_and_queued += QUEUED;
119if let Err(new) = self.state_and_queued.compare_exchange_weak(
120 state,
121 state_and_queued,
122 Relaxed,
123 Acquire,
124 ) {
125 state_and_queued = new;
126continue;
127 }
128 }
129130 futex_wait(&self.state_and_queued, state_and_queued, None);
131 state_and_queued = self.state_and_queued.load(Acquire);
132 }
133 }
134 }
135 }
136137#[cold]
138 #[track_caller]
139pub fn call(&self, ignore_poisoning: bool, f: &mut dyn FnMut(&public::OnceState)) {
140let mut state_and_queued = self.state_and_queued.load(Acquire);
141loop {
142let state = state_and_queued & STATE_MASK;
143let queued = state_and_queued & QUEUED != 0;
144match state {
145 COMPLETE => return,
146 POISONED if !ignore_poisoning => {
147// Panic to propagate the poison.
148panic!("Once instance has previously been poisoned");
149 }
150 INCOMPLETE | POISONED => {
151// Try to register the current thread as the one running.
152let next = RUNNING + if queued { QUEUED } else { 0 };
153if let Err(new) = self.state_and_queued.compare_exchange_weak(
154 state_and_queued,
155 next,
156 Acquire,
157 Acquire,
158 ) {
159 state_and_queued = new;
160continue;
161 }
162163// `waiter_queue` will manage other waiting threads, and
164 // wake them up on drop.
165let mut waiter_queue = CompletionGuard {
166 state_and_queued: &self.state_and_queued,
167 set_state_on_drop_to: POISONED,
168 };
169// Run the function, letting it know if we're poisoned or not.
170let f_state = public::OnceState {
171 inner: OnceState {
172 poisoned: state == POISONED,
173 set_state_to: Cell::new(COMPLETE),
174 },
175 };
176 f(&f_state);
177 waiter_queue.set_state_on_drop_to = f_state.inner.set_state_to.get();
178return;
179 }
180_ => {
181// All other values must be RUNNING.
182assert!(state == RUNNING);
183184// Set the QUEUED bit if it is not already set.
185if !queued {
186 state_and_queued += QUEUED;
187if let Err(new) = self.state_and_queued.compare_exchange_weak(
188 state,
189 state_and_queued,
190 Relaxed,
191 Acquire,
192 ) {
193 state_and_queued = new;
194continue;
195 }
196 }
197198 futex_wait(&self.state_and_queued, state_and_queued, None);
199 state_and_queued = self.state_and_queued.load(Acquire);
200 }
201 }
202 }
203 }
204}