Auto merge of #144465 - orlp:system-alloc-tls, r=joboet

Allow the global allocator to use thread-local storage and std::thread::current()

Fixes #115209.

Currently the thread-local storage implementation uses the `Global` allocator if it needs to allocate memory in some places. This effectively means the global allocator can not use thread-local variables. This is a shame as an allocator is precisely one of the locations where you'd *really* want to use thread-locals. We also see that this lead to hacks such as #116402, where we detect re-entrance and abort.

So I've made the places where I could find allocation happening in the TLS implementation use the `System` allocator instead. I also applied this change to the storage allocated for a `Thread` handle so that it may be used care-free in the global allocator as well, for e.g. registering it to a central place or parking primitives.

r? `@joboet`

Author: bors

library/core/src/alloc/global.rs

@@ -115,6 +115,31 @@ use crate::{cmp, ptr};
 ///   Whether allocations happen or not is not part of the program behavior, even if it
 ///   could be detected via an allocator that tracks allocations by printing or otherwise
 ///   having side effects.
+///
+/// # Re-entrance
+///
+/// When implementing a global allocator one has to be careful not to create an infinitely recursive
+/// implementation by accident, as many constructs in the Rust standard library may allocate in
+/// their implementation. For example, on some platforms [`std::sync::Mutex`] may allocate, so using
+/// it is highly problematic in a global allocator.
+///
+/// Generally speaking for this reason one should stick to library features available through
+/// [`core`], and avoid using [`std`] in a global allocator. A few features from [`std`] are
+/// guaranteed to not use `#[global_allocator]` to allocate:
+///
+///  - [`std::thread_local`],
+///  - [`std::thread::current`],
+///  - [`std::thread::park`] and [`std::thread::Thread`]'s [`unpark`] method and
+/// [`Clone`] implementation.
+///
+/// [`std`]: ../../std/index.html
+/// [`std::sync::Mutex`]: ../../std/sync/struct.Mutex.html
+/// [`std::thread_local`]: ../../std/macro.thread_local.html
+/// [`std::thread::current`]: ../../std/thread/fn.current.html
+/// [`std::thread::park`]: ../../std/thread/fn.park.html
+/// [`std::thread::Thread`]: ../../std/thread/struct.Thread.html
+/// [`unpark`]: ../../std/thread/struct.Thread.html#method.unpark
+
 #[stable(feature = "global_alloc", since = "1.28.0")]
 pub unsafe trait GlobalAlloc {
     /// Allocates memory as described by the given `layout`.

library/std/src/alloc.rs

@@ -11,7 +11,7 @@
 //!
 //! This attribute allows configuring the choice of global allocator.
 //! You can use this to implement a completely custom global allocator
-//! to route all default allocation requests to a custom object.
+//! to route all[^system-alloc] default allocation requests to a custom object.
 //!
 //! ```rust
 //! use std::alloc::{GlobalAlloc, System, Layout};
@@ -52,6 +52,13 @@
 //!
 //! The `#[global_allocator]` can only be used once in a crate
 //! or its recursive dependencies.
+//!
+//! [^system-alloc]: Note that the Rust standard library internals may still
+//! directly call [`System`] when necessary (for example for the runtime
+//! support typically required to implement a global allocator, see [re-entrance] on [`GlobalAlloc`]
+//! for more details).
+//!
+//! [re-entrance]: trait.GlobalAlloc.html#re-entrance
 
 #![deny(unsafe_op_in_unsafe_fn)]
 #![stable(feature = "alloc_module", since = "1.28.0")]

library/std/src/sys/pal/sgx/abi/mod.rs

@@ -3,6 +3,7 @@
 use core::arch::global_asm;
 use core::sync::atomic::{Atomic, AtomicUsize, Ordering};
 
+use crate::alloc::System;
 use crate::io::Write;
 
 // runtime features
@@ -63,7 +64,9 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 #[unsafe(no_mangle)]
 extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64) -> EntryReturn {
     // FIXME: how to support TLS in library mode?
-    let tls = Box::new(tls::Tls::new());
+    // We use the System allocator here such that the global allocator may use
+    // thread-locals.
+    let tls = Box::new_in(tls::Tls::new(), System);
     let tls_guard = unsafe { tls.activate() };
 
     if secondary {

library/std/src/sys/pal/sgx/abi/tls/mod.rs

@@ -89,13 +89,6 @@ impl Tls {
         ActiveTls { tls: self }
     }
 
-    #[allow(unused)]
-    pub unsafe fn activate_persistent(self: Box<Self>) {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        let ptr = Box::into_raw(self).cast_const().cast::<u8>();
-        unsafe { set_tls_ptr(ptr) };
-    }
-
     unsafe fn current<'a>() -> &'a Tls {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
         unsafe { &*(get_tls_ptr() as *const Tls) }

library/std/src/sys/thread/hermit.rs

@@ -1,4 +1,5 @@
 use crate::num::NonZero;
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 use crate::{io, ptr};
 
@@ -16,50 +17,49 @@ pub const DEFAULT_MIN_STACK_SIZE: usize = 1 << 20;
 impl Thread {
     pub unsafe fn new_with_coreid(
         stack: usize,
-        p: Box<dyn FnOnce()>,
+        init: Box<ThreadInit>,
         core_id: isize,
     ) -> io::Result<Thread> {
-        let p = Box::into_raw(Box::new(p));
+        let data = Box::into_raw(init);
         let tid = unsafe {
             hermit_abi::spawn2(
                 thread_start,
-                p.expose_provenance(),
+                data.expose_provenance(),
                 hermit_abi::Priority::into(hermit_abi::NORMAL_PRIO),
                 stack,
                 core_id,
             )
         };
 
         return if tid == 0 {
-            // The thread failed to start and as a result p was not consumed. Therefore, it is
+            // The thread failed to start and as a result data was not consumed. Therefore, it is
             // safe to reconstruct the box so that it gets deallocated.
             unsafe {
-                drop(Box::from_raw(p));
+                drop(Box::from_raw(data));
             }
             Err(io::const_error!(io::ErrorKind::Uncategorized, "unable to create thread!"))
         } else {
             Ok(Thread { tid })
         };
 
-        extern "C" fn thread_start(main: usize) {
-            unsafe {
-                // Finally, let's run some code.
-                Box::from_raw(ptr::with_exposed_provenance::<Box<dyn FnOnce()>>(main).cast_mut())();
+        extern "C" fn thread_start(data: usize) {
+            // SAFETY: we are simply recreating the box that was leaked earlier.
+            let init =
+                unsafe { Box::from_raw(ptr::with_exposed_provenance_mut::<ThreadInit>(data)) };
+            let rust_start = init.init();
+            rust_start();
 
-                // run all destructors
+            // Run all destructors.
+            unsafe {
                 crate::sys::thread_local::destructors::run();
-                crate::rt::thread_cleanup();
             }
+            crate::rt::thread_cleanup();
         }
     }
 
-    pub unsafe fn new(
-        stack: usize,
-        _name: Option<&str>,
-        p: Box<dyn FnOnce()>,
-    ) -> io::Result<Thread> {
+    pub unsafe fn new(stack: usize, init: Box<ThreadInit>) -> io::Result<Thread> {
         unsafe {
-            Thread::new_with_coreid(stack, p, -1 /* = no specific core */)
+            Thread::new_with_coreid(stack, init, -1 /* = no specific core */)
         }
     }
 

library/std/src/sys/thread/sgx.rs

@@ -2,6 +2,7 @@
 
 use crate::io;
 use crate::sys::pal::abi::{thread, usercalls};
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 
 pub struct Thread(task_queue::JoinHandle);
@@ -13,6 +14,7 @@ pub use self::task_queue::JoinNotifier;
 mod task_queue {
     use super::wait_notify;
     use crate::sync::{Mutex, MutexGuard};
+    use crate::thread::ThreadInit;
 
     pub type JoinHandle = wait_notify::Waiter;
 
@@ -25,19 +27,20 @@ mod task_queue {
     }
 
     pub(super) struct Task {
-        p: Box<dyn FnOnce() + Send>,
+        init: Box<ThreadInit>,
         done: JoinNotifier,
     }
 
     impl Task {
-        pub(super) fn new(p: Box<dyn FnOnce() + Send>) -> (Task, JoinHandle) {
+        pub(super) fn new(init: Box<ThreadInit>) -> (Task, JoinHandle) {
             let (done, recv) = wait_notify::new();
             let done = JoinNotifier(Some(done));
-            (Task { p, done }, recv)
+            (Task { init, done }, recv)
         }
 
         pub(super) fn run(self) -> JoinNotifier {
-            (self.p)();
+            let rust_start = self.init.init();
+            rust_start();
             self.done
         }
     }
@@ -93,14 +96,10 @@ pub mod wait_notify {
 
 impl Thread {
     // unsafe: see thread::Builder::spawn_unchecked for safety requirements
-    pub unsafe fn new(
-        _stack: usize,
-        _name: Option<&str>,
-        p: Box<dyn FnOnce() + Send>,
-    ) -> io::Result<Thread> {
+    pub unsafe fn new(_stack: usize, init: Box<ThreadInit>) -> io::Result<Thread> {
         let mut queue_lock = task_queue::lock();
         unsafe { usercalls::launch_thread()? };
-        let (task, handle) = task_queue::Task::new(p);
+        let (task, handle) = task_queue::Task::new(init);
         queue_lock.push(task);
         Ok(Thread(handle))
     }

library/std/src/sys/thread/solid.rs

@@ -8,6 +8,7 @@ use crate::sync::atomic::{Atomic, AtomicUsize, Ordering};
 use crate::sys::pal::itron::error::{ItronError, expect_success, expect_success_aborting};
 use crate::sys::pal::itron::time::dur2reltims;
 use crate::sys::pal::itron::{abi, task};
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 use crate::{hint, io};
 
@@ -27,9 +28,9 @@ unsafe impl Sync for Thread {}
 /// State data shared between a parent thread and child thread. It's dropped on
 /// a transition to one of the final states.
 struct ThreadInner {
-    /// This field is used on thread creation to pass a closure from
+    /// This field is used on thread creation to pass initialization data from
     /// `Thread::new` to the created task.
-    start: UnsafeCell<ManuallyDrop<Box<dyn FnOnce()>>>,
+    init: UnsafeCell<ManuallyDrop<Box<ThreadInit>>>,
 
     /// A state machine. Each transition is annotated with `[...]` in the
     /// source code.
@@ -65,7 +66,7 @@ struct ThreadInner {
     lifecycle: Atomic<usize>,
 }
 
-// Safety: The only `!Sync` field, `ThreadInner::start`, is only touched by
+// Safety: The only `!Sync` field, `ThreadInner::init`, is only touched by
 //         the task represented by `ThreadInner`.
 unsafe impl Sync for ThreadInner {}
 
@@ -84,13 +85,9 @@ impl Thread {
     /// # Safety
     ///
     /// See `thread::Builder::spawn_unchecked` for safety requirements.
-    pub unsafe fn new(
-        stack: usize,
-        _name: Option<&str>,
-        p: Box<dyn FnOnce()>,
-    ) -> io::Result<Thread> {
+    pub unsafe fn new(stack: usize, init: Box<ThreadInit>) -> io::Result<Thread> {
         let inner = Box::new(ThreadInner {
-            start: UnsafeCell::new(ManuallyDrop::new(p)),
+            init: UnsafeCell::new(ManuallyDrop::new(init)),
             lifecycle: AtomicUsize::new(LIFECYCLE_INIT),
         });
 
@@ -100,10 +97,11 @@ impl Thread {
             let inner = unsafe { &*p_inner };
 
             // Safety: Since `trampoline` is called only once for each
-            //         `ThreadInner` and only `trampoline` touches `start`,
-            //         `start` contains contents and is safe to mutably borrow.
-            let p = unsafe { ManuallyDrop::take(&mut *inner.start.get()) };
-            p();
+            //         `ThreadInner` and only `trampoline` touches `init`,
+            //         `init` contains contents and is safe to mutably borrow.
+            let init = unsafe { ManuallyDrop::take(&mut *inner.init.get()) };
+            let rust_start = init.init();
+            rust_start();
 
             // Fix the current thread's state just in case, so that the
             // destructors won't abort

library/std/src/sys/thread/teeos.rs

@@ -1,5 +1,6 @@
 use crate::mem::{self, ManuallyDrop};
 use crate::sys::os;
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 use crate::{cmp, io, ptr};
 
@@ -24,12 +25,8 @@ unsafe impl Sync for Thread {}
 
 impl Thread {
     // unsafe: see thread::Builder::spawn_unchecked for safety requirements
-    pub unsafe fn new(
-        stack: usize,
-        _name: Option<&str>,
-        p: Box<dyn FnOnce()>,
-    ) -> io::Result<Thread> {
-        let p = Box::into_raw(Box::new(p));
+    pub unsafe fn new(stack: usize, init: Box<ThreadInit>) -> io::Result<Thread> {
+        let data = Box::into_raw(init);
         let mut native: libc::pthread_t = unsafe { mem::zeroed() };
         let mut attr: libc::pthread_attr_t = unsafe { mem::zeroed() };
         assert_eq!(unsafe { libc::pthread_attr_init(&mut attr) }, 0);
@@ -62,16 +59,16 @@ impl Thread {
             }
         };
 
-        let ret = unsafe { libc::pthread_create(&mut native, &attr, thread_start, p as *mut _) };
-        // Note: if the thread creation fails and this assert fails, then p will
+        let ret = unsafe { libc::pthread_create(&mut native, &attr, thread_start, data as *mut _) };
+        // Note: if the thread creation fails and this assert fails, then data will
         // be leaked. However, an alternative design could cause double-free
         // which is clearly worse.
         assert_eq!(unsafe { libc::pthread_attr_destroy(&mut attr) }, 0);
 
         return if ret != 0 {
-            // The thread failed to start and as a result p was not consumed. Therefore, it is
+            // The thread failed to start and as a result data was not consumed. Therefore, it is
             // safe to reconstruct the box so that it gets deallocated.
-            drop(unsafe { Box::from_raw(p) });
+            drop(unsafe { Box::from_raw(data) });
             Err(io::Error::from_raw_os_error(ret))
         } else {
             // The new thread will start running earliest after the next yield.
@@ -80,15 +77,11 @@ impl Thread {
             Ok(Thread { id: native })
         };
 
-        extern "C" fn thread_start(main: *mut libc::c_void) -> *mut libc::c_void {
-            unsafe {
-                // Next, set up our stack overflow handler which may get triggered if we run
-                // out of stack.
-                // this is not necessary in TEE.
-                //let _handler = stack_overflow::Handler::new();
-                // Finally, let's run some code.
-                Box::from_raw(main as *mut Box<dyn FnOnce()>)();
-            }
+        extern "C" fn thread_start(data: *mut libc::c_void) -> *mut libc::c_void {
+            // SAFETY: we are simply recreating the box that was leaked earlier.
+            let init = unsafe { Box::from_raw(data as *mut ThreadInit) };
+            let rust_start = init.init();
+            rust_start();
             ptr::null_mut()
         }
     }