Ensure set_current is called early during thread init

Author: orlp

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/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()
         }
     }

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

@@ -14,6 +14,7 @@ use crate::sys::weak::dlsym;
 #[cfg(any(target_os = "solaris", target_os = "illumos", target_os = "nto",))]
 use crate::sys::weak::weak;
 use crate::sys::{os, stack_overflow};
+use crate::thread::{ThreadInit, current};
 use crate::time::Duration;
 use crate::{cmp, io, ptr};
 #[cfg(not(any(
@@ -30,11 +31,6 @@ pub const DEFAULT_MIN_STACK_SIZE: usize = 256 * 1024;
 #[cfg(any(target_os = "espidf", target_os = "nuttx"))]
 pub const DEFAULT_MIN_STACK_SIZE: usize = 0; // 0 indicates that the stack size configured in the ESP-IDF/NuttX menuconfig system should be used
 
-struct ThreadData {
-    name: Option<Box<str>>,
-    f: Box<dyn FnOnce()>,
-}
-
 pub struct Thread {
     id: libc::pthread_t,
 }
@@ -47,12 +43,8 @@ unsafe impl Sync for Thread {}
 impl Thread {
     // unsafe: see thread::Builder::spawn_unchecked for safety requirements
     #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
-    pub unsafe fn new(
-        stack: usize,
-        name: Option<&str>,
-        f: Box<dyn FnOnce()>,
-    ) -> io::Result<Thread> {
-        let data = Box::into_raw(Box::new(ThreadData { name: name.map(Box::from), f }));
+    pub unsafe fn new(stack: usize, init: Box<ThreadInit>) -> io::Result<Thread> {
+        let data = Box::into_raw(init);
         let mut native: libc::pthread_t = mem::zeroed();
         let mut attr: mem::MaybeUninit<libc::pthread_attr_t> = mem::MaybeUninit::uninit();
         assert_eq!(libc::pthread_attr_init(attr.as_mut_ptr()), 0);
@@ -118,12 +110,16 @@ impl Thread {
 
         extern "C" fn thread_start(data: *mut libc::c_void) -> *mut libc::c_void {
             unsafe {
-                let data = Box::from_raw(data as *mut ThreadData);
-                // Next, set up our stack overflow handler which may get triggered if we run
-                // out of stack.
-                let _handler = stack_overflow::Handler::new(data.name);
-                // Finally, let's run some code.
-                (data.f)();
+                // SAFETY: we are simply recreating the box that was leaked earlier.
+                let init = Box::from_raw(data as *mut ThreadInit);
+                let rust_start = init.init();
+
+                // Set up our thread name and stack overflow handler which may get triggered
+                // if we run out of stack.
+                let thread = current();
+                let _handler = stack_overflow::Handler::new(thread.name().map(Box::from));
+
+                rust_start();
             }
             ptr::null_mut()
         }

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

@@ -1,6 +1,7 @@
 use crate::ffi::CStr;
 use crate::io;
 use crate::num::NonZero;
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 
 pub struct Thread(!);
@@ -9,11 +10,7 @@ pub const DEFAULT_MIN_STACK_SIZE: usize = 64 * 1024;
 
 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> {
+    pub unsafe fn new(_stack: usize, _init: Box<ThreadInit>) -> io::Result<Thread> {
         Err(io::Error::UNSUPPORTED_PLATFORM)
     }
 

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

@@ -7,6 +7,7 @@ use crate::mem;
 use crate::num::NonZero;
 #[cfg(target_feature = "atomics")]
 use crate::sys::os;
+use crate::thread::ThreadInit;
 use crate::time::Duration;
 #[cfg(target_feature = "atomics")]
 use crate::{cmp, ptr};
@@ -73,12 +74,8 @@ pub const DEFAULT_MIN_STACK_SIZE: usize = 1024 * 1024;
 #[cfg(target_feature = "atomics")]
 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);
@@ -100,28 +97,28 @@ 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.
             unsafe {
-                drop(Box::from_raw(p));
+                drop(Box::from_raw(data));
             }
             Err(io::Error::from_raw_os_error(ret))
         } else {
             Ok(Thread { id: native })
         };
 
-        extern "C" fn thread_start(main: *mut libc::c_void) -> *mut libc::c_void {
-            unsafe {
-                // 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()
         }
     }