From 6631b73b8c242c4d4d9f936cd4d6d7efec95e26b Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Mon, 12 Dec 2022 23:04:05 +0000 Subject: [PATCH 1/2] Split the thread-id TLS into 2 variables This allows the fast path to avoid a branch which checks if the TLS destructor for the thread ID has been registered. --- src/lib.rs | 16 +++++++++------- src/thread_id.rs | 44 ++++++++++++++++++++++++++++++++------------ 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 3e30d66..7ba7102 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -189,8 +189,7 @@ impl ThreadLocal { /// Returns the element for the current thread, if it exists. pub fn get(&self) -> Option<&T> { - let thread = thread_id::get(); - self.get_inner(thread) + thread_id::try_get().and_then(|thread| self.get_inner(thread)) } /// Returns the element for the current thread, or creates it if it doesn't @@ -212,11 +211,13 @@ impl ThreadLocal { where F: FnOnce() -> Result, { - let thread = thread_id::get(); - match self.get_inner(thread) { - Some(x) => Ok(x), - None => Ok(self.insert(thread, create()?)), + let thread = thread_id::try_get(); + if let Some(thread) = thread { + if let Some(val) = self.get_inner(thread) { + return Ok(val); + } } + Ok(self.insert(create()?)) } fn get_inner(&self, thread: Thread) -> Option<&T> { @@ -237,7 +238,8 @@ impl ThreadLocal { } #[cold] - fn insert(&self, thread: Thread, data: T) -> &T { + fn insert(&self, data: T) -> &T { + let thread = thread_id::get(); let bucket_atomic_ptr = unsafe { self.buckets.get_unchecked(thread.bucket) }; let bucket_ptr: *const _ = bucket_atomic_ptr.load(Ordering::Acquire); diff --git a/src/thread_id.rs b/src/thread_id.rs index 205fada..6e7b848 100644 --- a/src/thread_id.rs +++ b/src/thread_id.rs @@ -7,6 +7,7 @@ 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; @@ -73,25 +74,44 @@ impl Thread { } } -/// Wrapper around `Thread` that allocates and deallocates the ID. -struct ThreadHolder(Thread); -impl ThreadHolder { - fn new() -> ThreadHolder { - ThreadHolder(Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc())) - } -} -impl Drop for ThreadHolder { +// 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) { - THREAD_ID_MANAGER.lock().unwrap().free(self.0.id); + let thread = THREAD.with(|thread| thread.get()).unwrap(); + THREAD_ID_MANAGER.lock().unwrap().free(thread.id); } } -thread_local!(static THREAD_HOLDER: ThreadHolder = ThreadHolder::new()); +/// Attempts to get the current thread if `get` has previously been +/// called. +#[inline] +pub(crate) fn try_get() -> Option { + THREAD.with(|thread| thread.get()) +} -/// Get the current thread. +/// Returns a thread ID for the current thread, allocating one if needed. #[inline] pub(crate) fn get() -> Thread { - THREAD_HOLDER.with(|holder| holder.0) + 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 + } + }) } #[test] From cedabb14305b05bd167b2aee3cc33b0c5a5c8f9c Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Fri, 3 Feb 2023 16:50:07 +0100 Subject: [PATCH 2/2] 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]