Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use a more efficient Once on platforms without threads #105698

Merged
merged 1 commit into from
Dec 19, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions library/std/src/sys/unsupported/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ pub mod fs;
pub mod io;
pub mod locks;
pub mod net;
pub mod once;
pub mod os;
#[path = "../unix/os_str.rs"]
pub mod os_str;
Expand Down
89 changes: 89 additions & 0 deletions library/std/src/sys/unsupported/once.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
use crate::cell::Cell;
use crate::sync as public;

pub struct Once {
state: Cell<State>,
}

pub struct OnceState {
poisoned: bool,
set_state_to: Cell<State>,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum State {
Incomplete,
Poisoned,
Running,
Complete,
}

struct CompletionGuard<'a> {
state: &'a Cell<State>,
set_state_on_drop_to: State,
}

impl<'a> Drop for CompletionGuard<'a> {
fn drop(&mut self) {
self.state.set(self.set_state_on_drop_to);
}
}

// Safety: threads are not supported on this platform.
unsafe impl Sync for Once {}

impl Once {
#[inline]
#[rustc_const_stable(feature = "const_once_new", since = "1.32.0")]
pub const fn new() -> Once {
Once { state: Cell::new(State::Incomplete) }
}

#[inline]
pub fn is_completed(&self) -> bool {
self.state.get() == State::Complete
}

#[cold]
#[track_caller]
pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) {
let state = self.state.get();
match state {
State::Poisoned if !ignore_poisoning => {
// Panic to propagate the poison.
panic!("Once instance has previously been poisoned");
}
State::Incomplete | State::Poisoned => {
self.state.set(State::Running);
// `guard` will set the new state on drop.
let mut guard =
CompletionGuard { state: &self.state, set_state_on_drop_to: State::Poisoned };
// Run the function, letting it know if we're poisoned or not.
let f_state = public::OnceState {
inner: OnceState {
poisoned: state == State::Poisoned,
set_state_to: Cell::new(State::Complete),
},
};
f(&f_state);
guard.set_state_on_drop_to = f_state.inner.set_state_to.get();
}
State::Running => {
panic!("one-time initialization may not be performed recursively");
}
State::Complete => {}
}
}
}

impl OnceState {
#[inline]
pub fn is_poisoned(&self) -> bool {
self.poisoned
}

#[inline]
pub fn poison(&self) {
self.set_state_to.set(State::Poisoned)
}
}
2 changes: 2 additions & 0 deletions library/std/src/sys/wasi/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ pub mod io;
#[path = "../unsupported/locks/mod.rs"]
pub mod locks;
pub mod net;
#[path = "../unsupported/once.rs"]
pub mod once;
pub mod os;
#[path = "../unix/os_str.rs"]
pub mod os_str;
Expand Down
2 changes: 2 additions & 0 deletions library/std/src/sys/wasm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ cfg_if::cfg_if! {
} else {
#[path = "../unsupported/locks/mod.rs"]
pub mod locks;
#[path = "../unsupported/once.rs"]
pub mod once;
#[path = "../unsupported/thread.rs"]
pub mod thread;
}
Expand Down
27 changes: 9 additions & 18 deletions library/std/src/sys_common/once/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,6 @@
// As a result, we end up implementing it ourselves in the standard library.
// This also gives us the opportunity to optimize the implementation a bit which
// should help the fast path on call sites.
//
// So to recap, the guarantees of a Once are that it will call the
// initialization closure at most once, and it will never return until the one
// that's running has finished running. This means that we need some form of
// blocking here while the custom callback is running at the very least.
// Additionally, we add on the restriction of **poisoning**. Whenever an
// initialization closure panics, the Once enters a "poisoned" state which means
// that all future calls will immediately panic as well.
//
// So to implement this, one might first reach for a `Mutex`, but those cannot
// be put into a `static`. It also gets a lot harder with poisoning to figure
// out when the mutex needs to be deallocated because it's not after the closure
// finishes, but after the first successful closure finishes.
//
// All in all, this is instead implemented with atomics and lock-free
// operations! Whee!

cfg_if::cfg_if! {
if #[cfg(any(
Expand All @@ -36,8 +20,15 @@ cfg_if::cfg_if! {
))] {
mod futex;
pub use futex::{Once, OnceState};
} else if #[cfg(any(
windows,
target_family = "unix",
all(target_vendor = "fortanix", target_env = "sgx"),
target_os = "solid_asp3",
))] {
mod queue;
pub use queue::{Once, OnceState};
} else {
mod generic;
pub use generic::{Once, OnceState};
pub use crate::sys::once::{Once, OnceState};
}
}