From e0058c71c3661a17cb7dbfa7826fe4dc3fbd6066 Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Wed, 14 Dec 2022 15:07:11 +0000 Subject: [PATCH] Add a "nightly" feature to use `#[thread_local]` --- Cargo.toml | 8 +++- src/lib.rs | 1 + src/thread_id.rs | 118 +++++++++++++++++++++++++++++++++-------------- 3 files changed, 92 insertions(+), 35 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bcfa98d..00cb91e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,13 +8,19 @@ license = "MIT OR Apache-2.0" repository = "https://github.com/Amanieu/thread_local-rs" readme = "README.md" keywords = ["thread_local", "concurrent", "thread"] -edition = "2018" +edition = "2021" + +[features] +# this feature provides performance improvements using nightly features +nightly = [] [badges] travis-ci = { repository = "Amanieu/thread_local-rs" } [dependencies] once_cell = "1.5.2" +# this is required to gate `nightly` related code paths +cfg-if = "1.0.0" # This is actually a dev-dependency, see https://github.com/rust-lang/cargo/issues/1596 criterion = { version = "0.4.0", optional = true } diff --git a/src/lib.rs b/src/lib.rs index 7ba7102..2283261 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -65,6 +65,7 @@ #![warn(missing_docs)] #![allow(clippy::mutex_atomic)] +#![cfg_attr(feature = "nightly", feature(thread_local))] mod cached; mod thread_id; diff --git a/src/thread_id.rs b/src/thread_id.rs index 6e7b848..d04a45b 100644 --- a/src/thread_id.rs +++ b/src/thread_id.rs @@ -7,7 +7,6 @@ use crate::POINTER_WIDTH; use once_cell::sync::Lazy; -use std::cell::Cell; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::sync::Mutex; @@ -74,44 +73,95 @@ impl Thread { } } -// This is split into 2 thread-local variables so that we can check whether the -// thread is initialized without having to register a thread-local destructor. -// -// This makes the fast path smaller. -thread_local! { static THREAD: Cell> = const { Cell::new(None) }; } -thread_local! { static THREAD_GUARD: ThreadGuard = const { ThreadGuard }; } +cfg_if::cfg_if! { + if #[cfg(feature = "nightly")] { + // This is split into 2 thread-local variables so that we can check whether the + // thread is initialized without having to register a thread-local destructor. + // + // This makes the fast path smaller. + #[thread_local] + static mut THREAD: Option = None; + thread_local! { static THREAD_GUARD: ThreadGuard = const { ThreadGuard }; } -// Guard to ensure the thread ID is released on thread exit. -struct ThreadGuard; + // Guard to ensure the thread ID is released on thread exit. + struct ThreadGuard; -impl Drop for ThreadGuard { - fn drop(&mut self) { - let thread = THREAD.with(|thread| thread.get()).unwrap(); - THREAD_ID_MANAGER.lock().unwrap().free(thread.id); - } -} + impl Drop for ThreadGuard { + fn drop(&mut self) { + // SAFETY: this is safe because we know that we (the current thread) + // are the only one who can be accessing our `THREAD` and thus + // it's safe for us to access and drop it. + if let Some(thread) = unsafe { THREAD.take() } { + THREAD_ID_MANAGER.lock().unwrap().free(thread.id); + } + } + } -/// Attempts to get the current thread if `get` has previously been -/// called. -#[inline] -pub(crate) fn try_get() -> Option { - THREAD.with(|thread| thread.get()) -} + /// Attempts to get the current thread if `get` has previously been + /// called. + #[inline] + pub(crate) fn try_get() -> Option { + unsafe { + THREAD + } + } -/// Returns a thread ID for the current thread, allocating one if needed. -#[inline] -pub(crate) fn get() -> Thread { - THREAD.with(|thread| { - if let Some(thread) = thread.get() { - thread - } else { - debug_assert!(thread.get().is_none()); - let new = Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc()); - thread.set(Some(new)); - THREAD_GUARD.with(|_| {}); - new + /// Returns a thread ID for the current thread, allocating one if needed. + #[inline] + pub(crate) fn get() -> Thread { + if let Some(thread) = unsafe { THREAD } { + thread + } else { + let new = Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc()); + unsafe { + THREAD = Some(new); + } + THREAD_GUARD.with(|_| {}); + new + } + } + } else { + use std::cell::Cell; + + // This is split into 2 thread-local variables so that we can check whether the + // thread is initialized without having to register a thread-local destructor. + // + // This makes the fast path smaller. + thread_local! { static THREAD: Cell> = const { Cell::new(None) }; } + thread_local! { static THREAD_GUARD: ThreadGuard = const { ThreadGuard }; } + + // Guard to ensure the thread ID is released on thread exit. + struct ThreadGuard; + + impl Drop for ThreadGuard { + fn drop(&mut self) { + let thread = THREAD.with(|thread| thread.get()).unwrap(); + THREAD_ID_MANAGER.lock().unwrap().free(thread.id); + } + } + + /// Attempts to get the current thread if `get` has previously been + /// called. + #[inline] + pub(crate) fn try_get() -> Option { + THREAD.with(|thread| thread.get()) + } + + /// Returns a thread ID for the current thread, allocating one if needed. + #[inline] + pub(crate) fn get() -> Thread { + THREAD.with(|thread| { + if let Some(thread) = thread.get() { + thread + } else { + let new = Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc()); + thread.set(Some(new)); + THREAD_GUARD.with(|_| {}); + new + } + }) } - }) + } } #[test]