#![warn(missing_docs)]
#![warn(clippy::undocumented_unsafe_blocks)]
#![warn(unsafe_op_in_unsafe_fn)]
#![cfg_attr(feature = "nightly", feature(thread_local))]
mod cached;
mod thread_id;
#[allow(deprecated)]
pub use cached::{CachedIntoIter, CachedIterMut, CachedThreadLocal};
use std::cell::UnsafeCell;
use std::fmt;
use std::iter::FusedIterator;
use std::mem::MaybeUninit;
use std::panic::UnwindSafe;
use std::ptr;
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};
use thread_id::Thread;
const BUCKETS: usize = (usize::BITS - 1) as usize;
pub struct ThreadLocal<T: Send> {
buckets: [AtomicPtr<Entry<T>>; BUCKETS],
values: AtomicUsize,
}
struct Entry<T> {
present: AtomicBool,
value: UnsafeCell<MaybeUninit<T>>,
}
impl<T> Entry<T> {
fn get_value_cell(&self) -> Option<&UnsafeCell<MaybeUninit<T>>> {
self.present.load(Ordering::Acquire).then_some(&self.value)
}
unsafe fn as_ref<'a>(&self) -> Option<&'a T> {
self.get_value_cell()
.map(|cell| unsafe { (&*cell.get()).assume_init_ref() })
}
}
impl<T> Drop for Entry<T> {
fn drop(&mut self) {
if *self.present.get_mut() {
unsafe {
MaybeUninit::assume_init_drop(&mut *self.value.get());
}
}
}
}
unsafe impl<T: Send> Sync for ThreadLocal<T> {}
impl<T: Send> Default for ThreadLocal<T> {
fn default() -> ThreadLocal<T> {
ThreadLocal::new()
}
}
impl<T: Send> Drop for ThreadLocal<T> {
fn drop(&mut self) {
for (i, bucket) in self.buckets.iter_mut().enumerate() {
let bucket_ptr = *bucket.get_mut();
let this_bucket_size = 1 << i;
if bucket_ptr.is_null() {
continue;
}
unsafe { deallocate_bucket(bucket_ptr, this_bucket_size) };
}
}
}
impl<T: Send> ThreadLocal<T> {
#[allow(clippy::declare_interior_mutable_const)]
const NULL_BUCKET: AtomicPtr<Entry<T>> = AtomicPtr::new(ptr::null_mut());
pub const fn new() -> ThreadLocal<T> {
Self {
buckets: [Self::NULL_BUCKET; BUCKETS],
values: AtomicUsize::new(0),
}
}
pub fn with_capacity(capacity: usize) -> ThreadLocal<T> {
let allocated_buckets = (usize::BITS - capacity.leading_zeros()) as usize;
let mut buckets = [Self::NULL_BUCKET; BUCKETS];
for (i, bucket) in buckets[..allocated_buckets].iter_mut().enumerate() {
*bucket.get_mut() = allocate_bucket::<T>(1 << i);
}
Self {
buckets,
values: AtomicUsize::new(0),
}
}
pub fn get(&self) -> Option<&T> {
thread_id::try_get().and_then(|id| self.get_inner(id))
}
pub fn get_or<F>(&self, create: F) -> &T
where
F: FnOnce() -> T,
{
let result = self.get_or_try(|| Ok::<T, ()>(create()));
unsafe { result.unwrap_unchecked() }
}
pub fn get_or_try<F, E>(&self, create: F) -> Result<&T, E>
where
F: FnOnce() -> Result<T, E>,
{
let thread = thread_id::get();
if let Some(val) = self.get_inner(thread) {
return Ok(val);
}
Ok(self.insert(thread, create()?))
}
fn get_inner(&self, thread: Thread) -> Option<&T> {
let bucket_ptr = self.get_bucket(thread).load(Ordering::Acquire);
if bucket_ptr.is_null() {
return None;
}
let entry = unsafe { &*bucket_ptr.add(thread.index) };
unsafe { entry.as_ref() }
}
#[cold]
fn insert(&self, thread: Thread, data: T) -> &T {
let bucket_atomic_ptr = self.get_bucket(thread);
let bucket_ptr: *const _ = bucket_atomic_ptr.load(Ordering::Acquire);
let bucket_ptr = if bucket_ptr.is_null() {
let new_bucket = allocate_bucket(thread.bucket_size());
match bucket_atomic_ptr.compare_exchange(
ptr::null_mut(),
new_bucket,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => new_bucket,
Err(bucket_ptr) => {
unsafe { deallocate_bucket(new_bucket, thread.bucket_size()) }
bucket_ptr
}
}
} else {
bucket_ptr
};
let entry = unsafe { &*bucket_ptr.add(thread.index) };
if let Some(existing) = unsafe { entry.as_ref() } {
return existing;
}
let value_ptr = entry.value.get();
let data_ref = &*unsafe { &mut *value_ptr }.write(data);
entry.present.store(true, Ordering::Release);
self.values.fetch_add(1, Ordering::Release);
data_ref
}
#[inline]
fn get_bucket(&self, thread: Thread) -> &AtomicPtr<Entry<T>> {
unsafe { self.buckets.get_unchecked(thread.bucket) }
}
pub fn iter(&self) -> Iter<'_, T>
where
T: Sync,
{
Iter {
thread_local: self,
raw: RawIter::new(),
}
}
pub fn iter_mut(&mut self) -> IterMut<'_, T> {
IterMut {
thread_local: self,
raw: RawIter::new(),
}
}
pub fn clear(&mut self) {
*self = ThreadLocal::new();
}
}
impl<T: Send> IntoIterator for ThreadLocal<T> {
type Item = T;
type IntoIter = IntoIter<T>;
fn into_iter(self) -> IntoIter<T> {
IntoIter {
thread_local: self,
raw: RawIter::new(),
}
}
}
impl<'a, T: Send + Sync> IntoIterator for &'a ThreadLocal<T> {
type Item = &'a T;
type IntoIter = Iter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T: Send> IntoIterator for &'a mut ThreadLocal<T> {
type Item = &'a mut T;
type IntoIter = IterMut<'a, T>;
fn into_iter(self) -> IterMut<'a, T> {
self.iter_mut()
}
}
impl<T: Send + Default> ThreadLocal<T> {
pub fn get_or_default(&self) -> &T {
self.get_or(Default::default)
}
}
impl<T: Send + fmt::Debug> fmt::Debug for ThreadLocal<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "ThreadLocal {{ local_data: {:?} }}", self.get())
}
}
impl<T: Send + UnwindSafe> UnwindSafe for ThreadLocal<T> {}
#[derive(Debug)]
struct RawIter {
yielded: usize,
bucket: usize,
bucket_size: usize,
index: usize,
}
impl RawIter {
#[inline]
fn new() -> Self {
Self {
yielded: 0,
bucket: 0,
bucket_size: 1,
index: 0,
}
}
fn next<'a, T: Send + Sync>(&mut self, thread_local: &'a ThreadLocal<T>) -> Option<&'a T> {
while let Some(bucket) = thread_local.buckets.get(self.bucket) {
let bucket = bucket.load(Ordering::Acquire);
if !bucket.is_null() {
while self.index < self.bucket_size {
let entry = unsafe { &*bucket.add(self.index) };
self.index += 1;
if let Some(value) = unsafe { entry.as_ref() } {
self.yielded += 1;
return Some(value);
}
}
}
self.next_bucket();
}
None
}
fn next_mut<'a, T: Send>(
&mut self,
thread_local: &'a mut ThreadLocal<T>,
) -> Option<&'a mut Entry<T>> {
if *thread_local.values.get_mut() == self.yielded {
return None;
}
loop {
let bucket = unsafe { thread_local.buckets.get_unchecked_mut(self.bucket) };
let bucket = *bucket.get_mut();
if !bucket.is_null() {
while self.index < self.bucket_size {
let entry = unsafe { &mut *bucket.add(self.index) };
self.index += 1;
if *entry.present.get_mut() {
self.yielded += 1;
return Some(entry);
}
}
}
self.next_bucket();
}
}
#[inline]
fn next_bucket(&mut self) {
self.bucket_size <<= 1;
self.bucket += 1;
self.index = 0;
}
fn size_hint<T: Send>(&self, thread_local: &ThreadLocal<T>) -> (usize, Option<usize>) {
let total = thread_local.values.load(Ordering::Relaxed);
(total.saturating_sub(self.yielded), None)
}
fn size_hint_frozen<T: Send>(&self, thread_local: &ThreadLocal<T>) -> (usize, Option<usize>) {
let total = thread_local.values.load(Ordering::Relaxed);
let remaining = total - self.yielded;
(remaining, Some(remaining))
}
}
#[derive(Debug)]
pub struct Iter<'a, T: Send + Sync> {
thread_local: &'a ThreadLocal<T>,
raw: RawIter,
}
impl<'a, T: Send + Sync> Iterator for Iter<'a, T> {
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
self.raw.next(self.thread_local)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.raw.size_hint(self.thread_local)
}
}
impl<T: Send + Sync> FusedIterator for Iter<'_, T> {}
pub struct IterMut<'a, T: Send> {
thread_local: &'a mut ThreadLocal<T>,
raw: RawIter,
}
impl<'a, T: Send> Iterator for IterMut<'a, T> {
type Item = &'a mut T;
fn next(&mut self) -> Option<&'a mut T> {
self.raw
.next_mut(self.thread_local)
.map(|entry| unsafe { (&mut *entry.value.get()).assume_init_mut() })
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.raw.size_hint_frozen(self.thread_local)
}
}
impl<T: Send> ExactSizeIterator for IterMut<'_, T> {}
impl<T: Send> FusedIterator for IterMut<'_, T> {}
impl<'a, T: Send + fmt::Debug> fmt::Debug for IterMut<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IterMut").field("raw", &self.raw).finish()
}
}
#[derive(Debug)]
pub struct IntoIter<T: Send> {
thread_local: ThreadLocal<T>,
raw: RawIter,
}
impl<T: Send> Iterator for IntoIter<T> {
type Item = T;
fn next(&mut self) -> Option<T> {
self.raw.next_mut(&mut self.thread_local).map(|entry| {
*entry.present.get_mut() = false;
let cell = unsafe { &mut *entry.value.get() };
let old_value = std::mem::replace(cell, MaybeUninit::uninit());
unsafe { old_value.assume_init() }
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.raw.size_hint_frozen(&self.thread_local)
}
}
impl<T: Send> ExactSizeIterator for IntoIter<T> {}
impl<T: Send> FusedIterator for IntoIter<T> {}
fn allocate_bucket<T>(size: usize) -> *mut Entry<T> {
Box::into_raw(
(0..size)
.map(|_| Entry::<T> {
present: AtomicBool::new(false),
value: UnsafeCell::new(MaybeUninit::uninit()),
})
.collect(),
) as *mut _
}
unsafe fn deallocate_bucket<T>(bucket: *mut Entry<T>, size: usize) {
let slice = unsafe { std::slice::from_raw_parts_mut(bucket, size) };
drop(unsafe { Box::from_raw(slice) });
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::RefCell;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;
use std::thread;
fn make_create() -> Arc<dyn Fn() -> usize + Send + Sync> {
let count = AtomicUsize::new(0);
Arc::new(move || count.fetch_add(1, Relaxed))
}
#[test]
fn same_thread() {
let create = make_create();
let mut tls = ThreadLocal::new();
assert_eq!(None, tls.get());
assert_eq!("ThreadLocal { local_data: None }", format!("{:?}", &tls));
assert_eq!(0, *tls.get_or(|| create()));
assert_eq!(Some(&0), tls.get());
assert_eq!(0, *tls.get_or(|| create()));
assert_eq!(Some(&0), tls.get());
assert_eq!(0, *tls.get_or(|| create()));
assert_eq!(Some(&0), tls.get());
assert_eq!("ThreadLocal { local_data: Some(0) }", format!("{:?}", &tls));
tls.clear();
assert_eq!(None, tls.get());
}
#[test]
fn reentrant_get_or() {
let tls: ThreadLocal<Box<u32>> = ThreadLocal::new();
let outer = tls.get_or(|| {
let inner = tls.get_or(|| Box::new(1));
assert_eq!(**inner, 1);
Box::new(2)
});
assert_eq!(**outer, 1);
let collected: Vec<u32> = tls.into_iter().map(|b| *b).collect();
assert_eq!(collected, vec![1]);
}
#[test]
fn different_thread() {
let create = make_create();
let tls = Arc::new(ThreadLocal::new());
assert_eq!(None, tls.get());
assert_eq!(0, *tls.get_or(|| create()));
assert_eq!(Some(&0), tls.get());
let tls2 = tls.clone();
let create2 = create.clone();
thread::spawn(move || {
assert_eq!(None, tls2.get());
assert_eq!(1, *tls2.get_or(|| create2()));
assert_eq!(Some(&1), tls2.get());
})
.join()
.unwrap();
assert_eq!(Some(&0), tls.get());
assert_eq!(0, *tls.get_or(|| create()));
}
#[test]
fn iter() {
let tls = Arc::new(ThreadLocal::new());
tls.get_or(|| Box::new(1));
let tls2 = tls.clone();
thread::spawn(move || {
tls2.get_or(|| Box::new(2));
let tls3 = tls2.clone();
thread::spawn(move || {
tls3.get_or(|| Box::new(3));
})
.join()
.unwrap();
drop(tls2);
})
.join()
.unwrap();
let mut tls = Arc::try_unwrap(tls).unwrap();
let mut v = tls.iter().map(|x| **x).collect::<Vec<i32>>();
v.sort_unstable();
assert_eq!(vec![1, 2, 3], v);
let mut v = tls.iter_mut().map(|x| **x).collect::<Vec<i32>>();
v.sort_unstable();
assert_eq!(vec![1, 2, 3], v);
let mut v = tls.into_iter().map(|x| *x).collect::<Vec<i32>>();
v.sort_unstable();
assert_eq!(vec![1, 2, 3], v);
}
#[test]
fn miri_iter_soundness_check() {
let tls = Arc::new(ThreadLocal::new());
let _local = tls.get_or(|| Box::new(1));
let tls2 = tls.clone();
let join_1 = thread::spawn(move || {
let _tls = tls2.get_or(|| Box::new(2));
let iter = tls2.iter();
for item in iter {
println!("{:?}", item);
}
});
let iter = tls.iter();
for item in iter {
println!("{:?}", item);
}
join_1.join().ok();
}
#[test]
fn test_drop() {
let local = ThreadLocal::new();
struct Dropped(Arc<AtomicUsize>);
impl Drop for Dropped {
fn drop(&mut self) {
self.0.fetch_add(1, Relaxed);
}
}
let dropped = Arc::new(AtomicUsize::new(0));
local.get_or(|| Dropped(dropped.clone()));
assert_eq!(dropped.load(Relaxed), 0);
drop(local);
assert_eq!(dropped.load(Relaxed), 1);
}
#[test]
fn test_earlyreturn_buckets() {
struct Dropped(Arc<AtomicUsize>);
impl Drop for Dropped {
fn drop(&mut self) {
self.0.fetch_add(1, Relaxed);
}
}
let dropped = Arc::new(AtomicUsize::new(0));
let thread = Thread::new(1234);
assert!(thread.bucket > 1);
let mut local = ThreadLocal::new();
local.insert(thread, Dropped(dropped.clone()));
let item = local.iter().next().unwrap();
assert_eq!(item.0.load(Relaxed), 0);
let item = local.iter_mut().next().unwrap();
assert_eq!(item.0.load(Relaxed), 0);
drop(local);
assert_eq!(dropped.load(Relaxed), 1);
}
#[test]
fn is_sync() {
fn foo<T: Sync>() {}
foo::<ThreadLocal<String>>();
foo::<ThreadLocal<RefCell<String>>>();
}
}