Allow the global alloc one TLS slot with a destructor

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

@@ -1,20 +1,26 @@
-use crate::cell::RefCell;
+use crate::cell::{Cell, RefCell};
 use crate::sys::thread_local::guard;
 
+#[thread_local]
+static REENTRANT_DTOR: Cell<Option<(*mut u8, unsafe extern "C" fn(*mut u8))>> = Cell::new(None);
+
 #[thread_local]
 static DTORS: RefCell<Vec<(*mut u8, unsafe extern "C" fn(*mut u8))>> = RefCell::new(Vec::new());
 
 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");
-    };
-
-    guard::enable();
-
-    dtors.push((t, dtor));
+    // Borrowing DTORS can only fail if the global allocator calls this
+    // function again.
+    if let Ok(mut dtors) = DTORS.try_borrow_mut() {
+        guard::enable();
+        dtors.push((t, dtor));
+    } else if REENTRANT_DTOR.get().is_none() {
+        guard::enable();
+        REENTRANT_DTOR.set(Some((t, dtor)));
+    } else {
+        rtabort!(
+            "the global allocator may only create one thread-local variable with a destructor"
+        );
+    }
 }
 
 /// The [`guard`] module contains platform-specific functions which will run this
@@ -27,18 +33,27 @@ pub unsafe fn register(t: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
 pub unsafe fn run() {
     loop {
         let mut dtors = DTORS.borrow_mut();
-        match dtors.pop() {
-            Some((t, dtor)) => {
-                drop(dtors);
-                unsafe {
-                    dtor(t);
-                }
-            }
-            None => {
-                // Free the list memory.
-                *dtors = Vec::new();
-                break;
-            }
+        let Some((t, dtor)) = dtors.pop() else { break };
+
+        // If the global allocator has allocated a thread-local variable
+        // it will always be the first in the dtors list (if an alloc came
+        // before the first regular register()), or be the REENTRANT_DTOR
+        // (if a register() came before the first alloc). In the former
+        // case we have to make sure not to touch the global (de)allocator again
+        // on this thread, so free dtors now, before calling the last dtor.
+        if dtors.is_empty() {
+            *dtors = Vec::new();
+        }
+
+        drop(dtors);
+        unsafe {
+            dtor(t);
+        }
+    }
+
+    if let Some((t, dtor)) = REENTRANT_DTOR.replace(None) {
+        unsafe {
+            dtor(t);
         }
     }
 }

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

@@ -0,0 +1,48 @@
+//@ run-pass
+//@ needs-threads
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+static GLOBAL: AtomicUsize = AtomicUsize::new(0);
+
+struct Local;
+
+thread_local! {
+    static LOCAL: Local = {
+        GLOBAL.fetch_or(1, Ordering::Relaxed);
+        Local
+    };
+}
+
+impl Drop for Local {
+    fn drop(&mut self) {
+        GLOBAL.fetch_or(2, Ordering::Relaxed);
+    }
+}
+
+#[global_allocator]
+static ALLOC: Alloc = Alloc;
+struct Alloc;
+
+unsafe impl GlobalAlloc for Alloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        LOCAL.with(|_local| {});
+        unsafe { System.alloc(layout) }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        LOCAL.with(|_local| {});
+        unsafe { System.dealloc(ptr, layout) }
+    }
+}
+
+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);
+}