Skip to content

Commit

Permalink
Auto merge of #105698 - joboet:unsupported_threads_once, r=thomcc
Browse files Browse the repository at this point in the history
Use a more efficient `Once` on platforms without threads

The current implementation uses an atomic queue and spins rather than panicking when calling `call_once` recursively. Since concurrency is not supported on platforms like WASM, `Once` can be implemented much more efficiently using just a single non-atomic state variable.
  • Loading branch information
bors committed Dec 19, 2022
2 parents 4653c93 + f9b5684 commit 7ab8038
Show file tree
Hide file tree
Showing 6 changed files with 103 additions and 18 deletions.
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};
}
}
File renamed without changes.

0 comments on commit 7ab8038

Please sign in to comment.