std/sys/sync/once/
futex.rs1use 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};
6
7const INCOMPLETE: Primitive = 0;
12const POISONED: Primitive = 1;
16const RUNNING: Primitive = 2;
19const COMPLETE: Primitive = 3;
21
22const QUEUED: Primitive = 4;
26
27const STATE_MASK: Primitive = 0b11;
32
33pub struct OnceState {
34    poisoned: bool,
35    set_state_to: Cell<Primitive>,
36}
37
38impl OnceState {
39    #[inline]
40    pub fn is_poisoned(&self) -> bool {
41        self.poisoned
42    }
43
44    #[inline]
45    pub fn poison(&self) {
46        self.set_state_to.set(POISONED);
47    }
48}
49
50struct CompletionGuard<'a> {
51    state_and_queued: &'a Futex,
52    set_state_on_drop_to: Primitive,
53}
54
55impl<'a> Drop for CompletionGuard<'a> {
56    fn drop(&mut self) {
57        if 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}
65
66pub struct Once {
67    state_and_queued: Futex,
68}
69
70impl Once {
71    #[inline]
72    pub const fn new() -> Once {
73        Once { state_and_queued: Futex::new(INCOMPLETE) }
74    }
75
76    #[inline]
77    pub fn is_completed(&self) -> bool {
78        self.state_and_queued.load(Acquire) == COMPLETE
81    }
82
83    #[inline]
84    pub(crate) fn state(&mut self) -> ExclusiveState {
85        match *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    }
92
93    #[inline]
94    pub(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    }
101
102    #[cold]
103    #[track_caller]
104    pub fn wait(&self, ignore_poisoning: bool) {
105        let mut state_and_queued = self.state_and_queued.load(Acquire);
106        loop {
107            let state = state_and_queued & STATE_MASK;
108            let queued = state_and_queued & QUEUED != 0;
109            match state {
110                COMPLETE => return,
111                POISONED if !ignore_poisoning => {
112                    panic!("Once instance has previously been poisoned");
114                }
115                _ => {
116                    if !queued {
118                        state_and_queued += QUEUED;
119                        if 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;
126                            continue;
127                        }
128                    }
129
130                    futex_wait(&self.state_and_queued, state_and_queued, None);
131                    state_and_queued = self.state_and_queued.load(Acquire);
132                }
133            }
134        }
135    }
136
137    #[cold]
138    #[track_caller]
139    pub fn call(&self, ignore_poisoning: bool, f: &mut dyn FnMut(&public::OnceState)) {
140        let mut state_and_queued = self.state_and_queued.load(Acquire);
141        loop {
142            let state = state_and_queued & STATE_MASK;
143            let queued = state_and_queued & QUEUED != 0;
144            match state {
145                COMPLETE => return,
146                POISONED if !ignore_poisoning => {
147                    panic!("Once instance has previously been poisoned");
149                }
150                INCOMPLETE | POISONED => {
151                    let next = RUNNING + if queued { QUEUED } else { 0 };
153                    if 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;
160                        continue;
161                    }
162
163                    let mut waiter_queue = CompletionGuard {
166                        state_and_queued: &self.state_and_queued,
167                        set_state_on_drop_to: POISONED,
168                    };
169                    let 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();
178                    return;
179                }
180                _ => {
181                    assert!(state == RUNNING);
183
184                    if !queued {
186                        state_and_queued += QUEUED;
187                        if 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;
194                            continue;
195                        }
196                    }
197
198                    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}