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,5 +1,6 @@
 use super::key::{Key, LazyKey, get, set};
 use super::{abort_on_dtor_unwind, guard};
+use crate::alloc::System;
 use crate::cell::Cell;
 use crate::marker::PhantomData;
 use crate::ptr;
@@ -95,8 +96,11 @@ impl<T: 'static> Storage<T> {
             return ptr::null();
         }
 
-        let value = Box::new(Value { value: i.and_then(Option::take).unwrap_or_else(f), key });
-        let ptr = Box::into_raw(value);
+        // We use the System allocator here to avoid interfering with a potential
+        // Global allocator using thread-local storage.
+        let value =
+            Box::new_in(Value { value: i.and_then(Option::take).unwrap_or_else(f), key }, System);
+        let ptr = Box::into_raw_with_allocator(value).0;
 
         // SAFETY:
         // * key came from a `LazyKey` and is thus correct.
@@ -114,7 +118,7 @@ impl<T: 'static> Storage<T> {
             // 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 { Box::from_raw(old) });
+            drop(unsafe { Box::from_raw_in(old, System) });
         }
 
         // SAFETY: We just created this value above.
@@ -133,7 +137,7 @@ unsafe extern "C" fn destroy_value<T: 'static>(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 { Box::from_raw(ptr as *mut Value<T>) };
+        let ptr = unsafe { Box::from_raw_in(ptr as *mut Value<T>, System) };
         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;
@@ -1432,7 +1433,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 {
@@ -1445,7 +1449,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);
@@ -1602,7 +1606,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.
@@ -1624,7 +1628,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);
+}