Support thread-local-storage

Author: NathanRoyer

library/std/src/os/custom/mod.rs

@@ -26,6 +26,8 @@
 //! - Initially, locking primitives are functional and behave as spinlocks.
 //! - Before `os::set_impl` has been called, any call to `abort_internal`
 //!   will result in an infinite loop.
+//! - Before `thread::set_impl` has been called, `None` is used where
+//!   `current_thread_id` would be called.
 //!
 //! These should be set/changed as soon as possible, as they are used internally.
 //!
@@ -272,6 +274,8 @@ pub mod thread {
 
     static_rwlock_box_impl!(ThreadManager);
 
+    pub type ThreadId = usize;
+
     /// Platform-specific management of threads
     pub trait ThreadManager: Send + Sync {
         /// unsafe: see thread::Builder::spawn_unchecked for safety requirements
@@ -282,6 +286,9 @@ pub mod thread {
         fn join(&self, thread: &Thread);
         fn available_parallelism(&self) -> io::Result<NonZeroUsize>;
 
+        /// Must return None for the initial thread
+        fn current_thread_id(&self) -> Option<ThreadId>;
+
         // todo: thread parking
     }
 }

library/std/src/sync/mutex.rs

@@ -275,6 +275,21 @@ impl<T: ?Sized> Mutex<T> {
         }
     }
 
+    /// On the "custom" platform, allocator and thread local storage
+    /// implementations use Mutex & RwLock, sometimes while a panic is
+    /// being handled.
+    ///
+    /// The poison checks that these primitives do are deadlock sources
+    /// in these cases; to prevent these deadlocks, it is preferable to
+    /// bypass the poison checks.
+    #[cfg(target_os = "custom")]
+    #[stable(feature = "rust1", since = "1.0.0")]
+    pub(crate) fn lock_no_poison_check(&self) -> MutexGuard<'_, T> {
+        self.inner.lock();
+        let poison = poison::Guard::no_check();
+        MutexGuard { lock: self, poison }
+    }
+
     /// Attempts to acquire this lock.
     ///
     /// If the lock could not be acquired at this time, then [`Err`] is returned.

library/std/src/sync/poison.rs

@@ -59,6 +59,23 @@ pub struct Guard {
     panicking: bool,
 }
 
+#[cfg(target_os = "custom")]
+impl Guard {
+    /// On the "custom" platform, allocator and thread local storage
+    /// implementations use Mutex & RwLock, sometimes while a panic is
+    /// being handled.
+    ///
+    /// The poison checks that these primitives do are deadlock sources
+    /// in these cases; to prevent these deadlocks, it is preferable to
+    /// bypass the poison checks.
+    pub fn no_check() -> Self {
+        // by setting `panicking` to true, `poison::Flag::done` doesn't call `thread::panicking`
+        Self {
+            panicking: true,
+        }
+    }
+}
+
 /// A type of error which can be returned whenever a lock is acquired.
 ///
 /// Both [`Mutex`]es and [`RwLock`]s are poisoned whenever a thread fails while the lock

library/std/src/sync/rwlock.rs

@@ -212,6 +212,26 @@ impl<T: ?Sized> RwLock<T> {
         }
     }
 
+    /// On the "custom" platform, allocator and thread local storage
+    /// implementations use Mutex & RwLock, sometimes while a panic is
+    /// being handled.
+    ///
+    /// The poison checks that these primitives do are deadlock sources
+    /// in these cases; to prevent these deadlocks, it is preferable to
+    /// bypass the poison checks.
+    #[inline]
+    #[cfg(target_os = "custom")]
+    #[stable(feature = "rust1", since = "1.0.0")]
+    pub(crate) fn read_no_poison_check(&self) -> RwLockReadGuard<'_, T> {
+        unsafe {
+            self.inner.read();
+            RwLockReadGuard {
+                data: NonNull::new_unchecked(self.data.get()),
+                inner_lock: &self.inner,
+            }
+        }
+    }
+
     /// Attempts to acquire this `RwLock` with shared read access.
     ///
     /// If the access could not be granted at this time, then `Err` is returned.
@@ -300,6 +320,22 @@ impl<T: ?Sized> RwLock<T> {
         }
     }
 
+    /// On the "custom" platform, allocator and thread local storage
+    /// implementations use Mutex & RwLock, sometimes while a panic is
+    /// being handled.
+    ///
+    /// The poison checks that these primitives do are deadlock sources
+    /// in these cases; to prevent these deadlocks, it is preferable to
+    /// bypass the poison checks.
+    #[inline]
+    #[cfg(target_os = "custom")]
+    #[stable(feature = "rust1", since = "1.0.0")]
+    pub(crate) fn write_no_poison_check(&self) -> RwLockWriteGuard<'_, T> {
+        self.inner.write();
+        let poison = poison::Guard::no_check();
+        RwLockWriteGuard { lock: self, poison }
+    }
+
     /// Attempts to lock this `RwLock` with exclusive write access.
     ///
     /// If the lock could not be acquired at this time, then `Err` is returned.

library/std/src/sys/common/thread_local/os_local.rs

@@ -11,7 +11,7 @@ use crate::{fmt, marker, panic, ptr};
 pub macro thread_local_inner {
     // used to generate the `LocalKey` value for const-initialized thread locals
     (@key $t:ty, const $init:expr) => {{
-        #[inline]
+        #[cfg_attr(not(target_os = "custom"), inline)]
         #[deny(unsafe_op_in_unsafe_fn)]
         unsafe fn __getit(
             _init: $crate::option::Option<&mut $crate::option::Option<$t>>,
@@ -20,7 +20,7 @@ pub macro thread_local_inner {
 
             // On platforms without `#[thread_local]` we fall back to the
             // same implementation as below for os thread locals.
-            #[inline]
+            #[cfg_attr(not(target_os = "custom"), inline)]
             const fn __init() -> $t { INIT_EXPR }
             static __KEY: $crate::thread::local_impl::Key<$t> =
                 $crate::thread::local_impl::Key::new();
@@ -46,12 +46,12 @@ pub macro thread_local_inner {
     // used to generate the `LocalKey` value for `thread_local!`
     (@key $t:ty, $init:expr) => {
         {
-            #[inline]
+            #[cfg_attr(not(target_os = "custom"), inline)]
             fn __init() -> $t { $init }
 
             // `#[inline] does not work on windows-gnu due to linking errors around dllimports.
             // See https://github.com/rust-lang/rust/issues/109797.
-            #[cfg_attr(not(windows), inline)]
+            #[cfg_attr(all(not(windows), not(target_os = "custom")), inline)]
             unsafe fn __getit(
                 init: $crate::option::Option<&mut $crate::option::Option<$t>>,
             ) -> $crate::option::Option<&'static $t> {

library/std/src/sys/custom/alloc.rs

@@ -125,7 +125,7 @@ fn encode_slot(i: usize, len: usize, next: usize) {
 fn free(i: usize, len: usize) {
     assert!(len >= 2);
 
-    let mut first_slot = FIRST_SLOT.lock().unwrap();
+    let mut first_slot = FIRST_SLOT.lock_no_poison_check();
     encode_slot(i, len, *first_slot);
     *first_slot = i;
 }
@@ -164,7 +164,7 @@ unsafe impl GlobalAlloc for DefaultAlloc {
         let (req_align, req_size) = prepare_layout(layout);
 
         let (mut filler, leftover);
-        let mut first_slot = FIRST_SLOT.lock().unwrap();
+        let mut first_slot = FIRST_SLOT.lock_no_poison_check();
         let mut i = *first_slot;
         let mut prev = None;
 

library/std/src/sys/custom/mod.rs

@@ -24,10 +24,6 @@ pub mod args;
 #[path = "../unsupported/once.rs"]
 pub mod once;
 
-// "custom" doesn't support thread local storage
-#[path = "../unsupported/thread_local_key.rs"]
-pub mod thread_local_key;
-
 #[path = "../unsupported/common.rs"]
 #[deny(unsafe_op_in_unsafe_fn)]
 #[allow(unused)]
@@ -47,6 +43,9 @@ pub mod thread;
 pub mod thread_parking;
 pub mod time;
 
+// really bad implementation
+pub mod thread_local_key;
+
 pub fn decode_error_kind(errno: i32) -> std_io::ErrorKind {
     custom_os_impl!(os, decode_error_kind, errno)
 }

library/std/src/sys/custom/thread_local_key.rs

@@ -0,0 +1,117 @@
+use crate::sync::{RwLock, Mutex};
+use crate::os::custom::thread::{ThreadId, IMPL};
+
+pub type Key = usize;
+
+const NULL: *mut u8 = core::ptr::null_mut();
+
+type Destructor = Option<unsafe extern "C" fn(*mut u8)>;
+
+fn thread_id() -> Option<ThreadId> {
+    let reader = IMPL.read_no_poison_check();
+    reader.as_ref()?.current_thread_id()
+}
+
+// safety: each ThreadStorage is virtually owned by one thread
+unsafe impl Sync for ThreadStorage {}
+unsafe impl Send for ThreadStorage {}
+
+struct ThreadStorage(Mutex<Vec<*mut u8>>);
+
+impl ThreadStorage {
+    pub const fn new() -> Self {
+        Self(Mutex::new(Vec::new()))
+    }
+
+    pub fn get(&self, key: Key) -> *mut u8 {
+        let locked = self.0.lock_no_poison_check();
+        match locked.get(key - 1) {
+            Some(pointer_ref) => *pointer_ref,
+            None => NULL,
+        }
+    }
+
+    pub fn set(&self, key: Key, value: *mut u8) {
+        let mut locked = self.0.lock_no_poison_check();
+        locked.resize(key, NULL);
+        locked[key - 1] = value;
+    }
+}
+
+struct Slots {
+    destructors: RwLock<Vec<Destructor>>,
+    initial_thread: ThreadStorage,
+    other_threads: RwLock<Vec<(ThreadId, ThreadStorage)>>,
+}
+
+impl Slots {
+    pub const fn new() -> Self {
+        Self {
+            destructors: RwLock::new(Vec::new()),
+            initial_thread: ThreadStorage::new(),
+            other_threads: RwLock::new(Vec::new()),
+        }
+    }
+
+    pub fn push(&self, destructor: Destructor) -> Key {
+        let mut destructors = self.destructors.write_no_poison_check();
+        destructors.push(destructor);
+        destructors.len()
+    }
+
+    pub fn get_dtor(&self, key: Key) -> Destructor {
+        let destructors = self.destructors.read_no_poison_check();
+        destructors[key - 1]
+    }
+
+    pub fn with_thread_storage<T, F: FnOnce(&ThreadStorage) -> T>(&self, callback: F) -> T {
+        if let Some(id) = thread_id() {
+            loop {
+                let finder = |(tid, _): &(ThreadId, ThreadStorage)| tid.cmp(&id);
+
+                {
+                    let other_threads = self.other_threads.read_no_poison_check();
+                    let position = other_threads.binary_search_by(finder);
+
+                    if let Ok(i) = position {
+                        return callback(&other_threads[i].1);
+                    }
+                }
+
+                // cannot re-use `position` (race condition)
+                let mut other_threads = self.other_threads.write_no_poison_check();
+                let i = other_threads.binary_search_by(finder).unwrap_err();
+                other_threads.insert(i, (id, ThreadStorage::new()));
+            }
+        } else {
+            callback(&self.initial_thread)
+        }
+    }
+}
+
+static SLOTS: Slots = Slots::new();
+
+#[inline]
+pub unsafe fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {
+    SLOTS.push(dtor)
+}
+
+#[inline]
+pub unsafe fn set(key: Key, value: *mut u8) {
+    SLOTS.with_thread_storage(|storage| storage.set(key, value))
+}
+
+#[inline]
+pub unsafe fn get(key: Key) -> *mut u8 {
+    SLOTS.with_thread_storage(|storage| storage.get(key))
+}
+
+#[inline]
+pub unsafe fn destroy(key: Key) {
+    let value = get(key);
+    if !value.is_null() {
+        if let Some(destructor) = SLOTS.get_dtor(key) {
+            destructor(value)
+        }
+    }
+}