Use System allocator for thread-local storage

Author: orlp

library/std/src/sys/thread_local/destructors/list.rs

@@ -1,19 +1,14 @@
+use crate::alloc::System;
 use crate::cell::RefCell;
 use crate::sys::thread_local::guard;
 
 #[thread_local]
-static DTORS: RefCell<Vec<(*mut u8, unsafe extern "C" fn(*mut u8))>> = RefCell::new(Vec::new());
+static DTORS: RefCell<Vec<(*mut u8, unsafe extern "C" fn(*mut u8)), System>> =
+    RefCell::new(Vec::new_in(System));
 
 pub unsafe fn register(t: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
-    let Ok(mut dtors) = DTORS.try_borrow_mut() else {
-        // This point can only be reached if the global allocator calls this
-        // function again.
-        // FIXME: maybe use the system allocator instead?
-        rtabort!("the global allocator may not use TLS with destructors");
-    };
-
+    let Ok(mut dtors) = DTORS.try_borrow_mut() else { rtabort!("unreachable") };
     guard::enable();
-
     dtors.push((t, dtor));
 }
 
@@ -36,7 +31,7 @@ pub unsafe fn run() {
             }
             None => {
                 // Free the list memory.
-                *dtors = Vec::new();
+                *dtors = Vec::new_in(System);
                 break;
             }
         }

library/std/src/sys/thread_local/key/xous.rs

@@ -38,6 +38,7 @@
 
 use core::arch::asm;
 
+use crate::alloc::System;
 use crate::mem::ManuallyDrop;
 use crate::os::xous::ffi::{MemoryFlags, map_memory, unmap_memory};
 use crate::ptr;
@@ -151,7 +152,10 @@ struct Node {
 }
 
 unsafe fn register_dtor(key: Key, dtor: Dtor) {
-    let mut node = ManuallyDrop::new(Box::new(Node { key, dtor, next: ptr::null_mut() }));
+    // We use the System allocator here to avoid interfering with a potential
+    // Global allocator using thread-local storage.
+    let mut node =
+        ManuallyDrop::new(Box::new_in(Node { key, dtor, next: ptr::null_mut() }, System));
 
     #[allow(unused_unsafe)]
     let mut head = unsafe { DTORS.load(Acquire) };

library/std/src/sys/thread_local/os.rs

@@ -1,6 +1,6 @@
 use super::key::{Key, LazyKey, get, set};
 use super::{abort_on_dtor_unwind, guard};
-use crate::alloc::{self, Layout};
+use crate::alloc::{self, Layout, System};
 use crate::cell::Cell;
 use crate::marker::PhantomData;
 use crate::mem::ManuallyDrop;
@@ -113,17 +113,19 @@ pub const fn value_align<T: 'static>() -> usize {
     crate::mem::align_of::<Value<T>>()
 }
 
-/// Equivalent to `Box<Value<T>>`, but potentially over-aligned.
-struct AlignedBox<T: 'static, const ALIGN: usize> {
+/// Equivalent to `Box<Value<T>, System>`, but potentially over-aligned.
+struct AlignedSystemBox<T: 'static, const ALIGN: usize> {
     ptr: NonNull<Value<T>>,
 }
 
-impl<T: 'static, const ALIGN: usize> AlignedBox<T, ALIGN> {
+impl<T: 'static, const ALIGN: usize> AlignedSystemBox<T, ALIGN> {
     #[inline]
     fn new(v: Value<T>) -> Self {
         let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
 
-        let ptr: *mut Value<T> = (unsafe { alloc::alloc(layout) }).cast();
+        // We use the System allocator here to avoid interfering with a potential
+        // Global allocator using thread-local storage.
+        let ptr: *mut Value<T> = (unsafe { System.alloc(layout) }).cast();
         let Some(ptr) = NonNull::new(ptr) else {
             alloc::handle_alloc_error(layout);
         };
@@ -143,7 +145,7 @@ impl<T: 'static, const ALIGN: usize> AlignedBox<T, ALIGN> {
     }
 }
 
-impl<T: 'static, const ALIGN: usize> Deref for AlignedBox<T, ALIGN> {
+impl<T: 'static, const ALIGN: usize> Deref for AlignedSystemBox<T, ALIGN> {
     type Target = Value<T>;
 
     #[inline]
@@ -152,14 +154,14 @@ impl<T: 'static, const ALIGN: usize> Deref for AlignedBox<T, ALIGN> {
     }
 }
 
-impl<T: 'static, const ALIGN: usize> Drop for AlignedBox<T, ALIGN> {
+impl<T: 'static, const ALIGN: usize> Drop for AlignedSystemBox<T, ALIGN> {
     #[inline]
     fn drop(&mut self) {
         let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
 
         unsafe {
             let unwind_result = catch_unwind(AssertUnwindSafe(|| self.ptr.drop_in_place()));
-            alloc::dealloc(self.ptr.as_ptr().cast(), layout);
+            System.dealloc(self.ptr.as_ptr().cast(), layout);
             if let Err(payload) = unwind_result {
                 resume_unwind(payload);
             }
@@ -205,11 +207,11 @@ impl<T: 'static, const ALIGN: usize> Storage<T, ALIGN> {
             return ptr::null();
         }
 
-        let value = AlignedBox::<T, ALIGN>::new(Value {
+        let value = AlignedSystemBox::<T, ALIGN>::new(Value {
             value: i.and_then(Option::take).unwrap_or_else(f),
             key,
         });
-        let ptr = AlignedBox::into_raw(value);
+        let ptr = AlignedSystemBox::into_raw(value);
 
         // SAFETY:
         // * key came from a `LazyKey` and is thus correct.
@@ -227,7 +229,7 @@ impl<T: 'static, const ALIGN: usize> Storage<T, ALIGN> {
             // initializer has already returned and the next scope only starts
             // after we return the pointer. Therefore, there can be no references
             // to the old value.
-            drop(unsafe { AlignedBox::<T, ALIGN>::from_raw(old) });
+            drop(unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(old) });
         }
 
         // SAFETY: We just created this value above.
@@ -246,7 +248,7 @@ unsafe extern "C" fn destroy_value<T: 'static, const ALIGN: usize>(ptr: *mut u8)
     // Note that to prevent an infinite loop we reset it back to null right
     // before we return from the destructor ourselves.
     abort_on_dtor_unwind(|| {
-        let ptr = unsafe { AlignedBox::<T, ALIGN>::from_raw(ptr as *mut Value<T>) };
+        let ptr = unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(ptr as *mut Value<T>) };
         let key = ptr.key;
         // SAFETY: `key` is the TLS key `ptr` was stored under.
         unsafe { set(key, ptr::without_provenance_mut(1)) };

library/std/src/thread/mod.rs

@@ -158,6 +158,7 @@
 #[cfg(all(test, not(any(target_os = "emscripten", target_os = "wasi"))))]
 mod tests;
 
+use crate::alloc::System;
 use crate::any::Any;
 use crate::cell::UnsafeCell;
 use crate::ffi::CStr;
@@ -1466,7 +1467,10 @@ impl Inner {
 ///
 /// [`thread::current`]: current::current
 pub struct Thread {
-    inner: Pin<Arc<Inner>>,
+    // We use the System allocator such that creating or dropping this handle
+    // does not interfere with a potential Global allocator using thread-local
+    // storage.
+    inner: Pin<Arc<Inner, System>>,
 }
 
 impl Thread {
@@ -1479,7 +1483,7 @@ impl Thread {
         // SAFETY: We pin the Arc immediately after creation, so its address never
         // changes.
         let inner = unsafe {
-            let mut arc = Arc::<Inner>::new_uninit();
+            let mut arc = Arc::<Inner, _>::new_uninit_in(System);
             let ptr = Arc::get_mut_unchecked(&mut arc).as_mut_ptr();
             (&raw mut (*ptr).name).write(name);
             (&raw mut (*ptr).id).write(id);
@@ -1650,7 +1654,7 @@ impl Thread {
     pub fn into_raw(self) -> *const () {
         // Safety: We only expose an opaque pointer, which maintains the `Pin` invariant.
         let inner = unsafe { Pin::into_inner_unchecked(self.inner) };
-        Arc::into_raw(inner) as *const ()
+        Arc::into_raw_with_allocator(inner).0 as *const ()
     }
 
     /// Constructs a `Thread` from a raw pointer.
@@ -1672,7 +1676,9 @@ impl Thread {
     #[unstable(feature = "thread_raw", issue = "97523")]
     pub unsafe fn from_raw(ptr: *const ()) -> Thread {
         // Safety: Upheld by caller.
-        unsafe { Thread { inner: Pin::new_unchecked(Arc::from_raw(ptr as *const Inner)) } }
+        unsafe {
+            Thread { inner: Pin::new_unchecked(Arc::from_raw_in(ptr as *const Inner, System)) }
+        }
     }
 
     fn cname(&self) -> Option<&CStr> {

tests/ui/threads-sendsync/tls-in-global-alloc.rs

@@ -0,0 +1,64 @@
+//@ run-pass
+//@ needs-threads
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::thread::Thread;
+use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};
+
+static GLOBAL: AtomicUsize = AtomicUsize::new(0);
+
+struct Local(Thread);
+
+thread_local! {
+    static LOCAL: Local = {
+        GLOBAL.fetch_or(1, Ordering::Relaxed);
+        Local(std::thread::current())
+    };
+}
+
+impl Drop for Local {
+    fn drop(&mut self) {
+        GLOBAL.fetch_or(2, Ordering::Relaxed);
+    }
+}
+
+static SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS: AtomicBool = AtomicBool::new(false);
+
+#[global_allocator]
+static ALLOC: Alloc = Alloc;
+struct Alloc;
+
+unsafe impl GlobalAlloc for Alloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // Make sure we aren't re-entrant.
+        assert!(!SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.load(Ordering::Relaxed));
+        SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(true, Ordering::Relaxed);
+        LOCAL.with(|local| {
+            assert!(local.0.id() == std::thread::current().id());
+        });
+        let ret = unsafe { System.alloc(layout) };
+        SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(false, Ordering::Relaxed);
+        ret
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Make sure we aren't re-entrant.
+        assert!(!SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.load(Ordering::Relaxed));
+        SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(true, Ordering::Relaxed);
+        LOCAL.with(|local| {
+            assert!(local.0.id() == std::thread::current().id());
+        });
+        unsafe { System.dealloc(ptr, layout) }
+        SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(false, Ordering::Relaxed);
+    }
+}
+
+fn main() {
+    std::thread::spawn(|| {
+        std::hint::black_box(vec![1, 2]);
+        assert!(GLOBAL.load(Ordering::Relaxed) == 1);
+    })
+    .join()
+    .unwrap();
+    assert!(GLOBAL.load(Ordering::Relaxed) == 3);
+}