Auto merge of #54951 - alexcrichton:more-wasm-threads, r=sfackler

std: Implement TLS for wasm32-unknown-unknown

This adds an implementation of thread local storage for the
`wasm32-unknown-unknown` target when the `atomics` feature is
implemented. This, however, comes with a notable caveat of that it
requires a new feature of the standard library, `wasm-bindgen-threads`,
to be enabled.

Thread local storage for wasm (when `atomics` are enabled and there's
actually more than one thread) is powered by the assumption that an
external entity can fill in some information for us. It's not currently
clear who will fill in this information nor whose responsibility it
should be long-term. In the meantime there's a strategy being gamed out
in the `wasm-bindgen` project specifically, and the hope is that we can
continue to test and iterate on the standard library without committing
to a particular strategy yet.

As to the details of `wasm-bindgen`'s strategy, LLVM doesn't currently
have the ability to emit custom `global` values (thread locals in a
`WebAssembly.Module`) so we leverage the `wasm-bindgen` CLI tool to do
it for us. To that end we have a few intrinsics, assuming two global values:

* `__wbindgen_current_id` - gets the current thread id as a 32-bit
  integer. It's `wasm-bindgen`'s responsibility to initialize this
  per-thread and then inform libstd of the id. Currently `wasm-bindgen`
  performs this initialization as part of the `start` function.
* `__wbindgen_tcb_{get,set}` - in addition to a thread id it's assumed
  that there's a global available for simply storing a pointer's worth
  of information (a thread control block, which currently only contains
  thread local storage). This would ideally be a native `global`
  injected by LLVM, but we don't have a great way to support that right
  now.

To reiterate, this is all intended to be unstable and purely intended
for testing out Rust on the web with threads. The story is very likely
to change in the future and we want to make sure that we're able to do
that!

Author: bors

Diff for: src/libstd/Cargo.toml

@@ -48,4 +48,13 @@ jemalloc = ["alloc_jemalloc"]
 force_alloc_system = []
 panic-unwind = ["panic_unwind"]
 profiler = ["profiler_builtins"]
+
+# An off-by-default feature which enables a linux-syscall-like ABI for libstd to
+# interoperate with the host environment. Currently not well documented and
+# requires rebuilding the standard library to use it.
 wasm_syscall = []
+
+# An off-by-default features to enable libstd to assume that wasm-bindgen is in
+# the environment for hooking up some thread-related information like the
+# current thread id and accessing/getting the current thread's TCB
+wasm-bindgen-threads = []

Diff for: src/libstd/sys/wasm/mutex_atomics.rs

@@ -11,7 +11,8 @@
 use arch::wasm32::atomic;
 use cell::UnsafeCell;
 use mem;
-use sync::atomic::{AtomicUsize, AtomicU64, Ordering::SeqCst};
+use sync::atomic::{AtomicUsize, AtomicU32, Ordering::SeqCst};
+use sys::thread;
 
 pub struct Mutex {
     locked: AtomicUsize,
@@ -70,7 +71,7 @@ impl Mutex {
 }
 
 pub struct ReentrantMutex {
-    owner: AtomicU64,
+    owner: AtomicU32,
     recursions: UnsafeCell<u32>,
 }
 
@@ -91,7 +92,7 @@ unsafe impl Sync for ReentrantMutex {}
 impl ReentrantMutex {
     pub unsafe fn uninitialized() -> ReentrantMutex {
         ReentrantMutex {
-            owner: AtomicU64::new(0),
+            owner: AtomicU32::new(0),
             recursions: UnsafeCell::new(0),
         }
     }
@@ -101,20 +102,20 @@ impl ReentrantMutex {
     }
 
     pub unsafe fn lock(&self) {
-        let me = thread_id();
+        let me = thread::my_id();
         while let Err(owner) = self._try_lock(me) {
-            let val = atomic::wait_i64(self.ptr(), owner as i64, -1);
+            let val = atomic::wait_i32(self.ptr(), owner as i32, -1);
             debug_assert!(val == 0 || val == 1);
         }
     }
 
     #[inline]
     pub unsafe fn try_lock(&self) -> bool {
-        self._try_lock(thread_id()).is_ok()
+        self._try_lock(thread::my_id()).is_ok()
     }
 
     #[inline]
-    unsafe fn _try_lock(&self, id: u64) -> Result<(), u64> {
+    unsafe fn _try_lock(&self, id: u32) -> Result<(), u32> {
         let id = id.checked_add(1).unwrap(); // make sure `id` isn't 0
         match self.owner.compare_exchange(0, id, SeqCst, SeqCst) {
             // we transitioned from unlocked to locked
@@ -153,11 +154,7 @@ impl ReentrantMutex {
     }
 
     #[inline]
-    fn ptr(&self) -> *mut i64 {
-        &self.owner as *const AtomicU64 as *mut i64
+    fn ptr(&self) -> *mut i32 {
+        &self.owner as *const AtomicU32 as *mut i32
     }
 }
-
-fn thread_id() -> u64 {
-    panic!("thread ids not implemented on wasm with atomics yet")
-}

Diff for: src/libstd/sys/wasm/thread.rs

@@ -69,3 +69,49 @@ pub mod guard {
     pub unsafe fn init() -> Option<Guard> { None }
     pub unsafe fn deinit() {}
 }
+
+cfg_if! {
+    if #[cfg(all(target_feature = "atomics", feature = "wasm-bindgen-threads"))] {
+        #[link(wasm_import_module = "__wbindgen_thread_xform__")]
+        extern {
+            fn __wbindgen_current_id() -> u32;
+            fn __wbindgen_tcb_get() -> u32;
+            fn __wbindgen_tcb_set(ptr: u32);
+        }
+        pub fn my_id() -> u32 {
+            unsafe { __wbindgen_current_id() }
+        }
+
+        // These are currently only ever used in `thread_local_atomics.rs`, if
+        // you'd like to use them be sure to update that and make sure everyone
+        // agrees what's what.
+        pub fn tcb_get() -> *mut u8 {
+            use mem;
+            assert_eq!(mem::size_of::<*mut u8>(), mem::size_of::<u32>());
+            unsafe { __wbindgen_tcb_get() as *mut u8 }
+        }
+
+        pub fn tcb_set(ptr: *mut u8) {
+            unsafe { __wbindgen_tcb_set(ptr as u32); }
+        }
+
+        // FIXME: still need something for hooking exiting a thread to free
+        // data...
+
+    } else if #[cfg(target_feature = "atomics")] {
+        pub fn my_id() -> u32 {
+            panic!("thread ids not implemented on wasm with atomics yet")
+        }
+
+        pub fn tcb_get() -> *mut u8 {
+            panic!("thread local data not implemented on wasm with atomics yet")
+        }
+
+        pub fn tcb_set(ptr: *mut u8) {
+            panic!("thread local data not implemented on wasm with atomics yet")
+        }
+    } else {
+        // stubbed out because no functions actually access these intrinsics
+        // unless atomics are enabled
+    }
+}

Diff for: src/libstd/sys/wasm/thread_local_atomics.rs

@@ -8,22 +8,61 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use sys::thread;
+use sync::atomic::{AtomicUsize, Ordering::SeqCst};
+
+const MAX_KEYS: usize = 128;
+static NEXT_KEY: AtomicUsize = AtomicUsize::new(0);
+
+struct ThreadControlBlock {
+    keys: [*mut u8; MAX_KEYS],
+}
+
+impl ThreadControlBlock {
+    fn new() -> ThreadControlBlock {
+        ThreadControlBlock {
+            keys: [0 as *mut u8; MAX_KEYS],
+        }
+    }
+
+    fn get() -> *mut ThreadControlBlock {
+        let ptr = thread::tcb_get();
+        if !ptr.is_null() {
+            return ptr as *mut ThreadControlBlock
+        }
+        let tcb = Box::into_raw(Box::new(ThreadControlBlock::new()));
+        thread::tcb_set(tcb as *mut u8);
+        tcb
+    }
+}
+
 pub type Key = usize;
 
-pub unsafe fn create(_dtor: Option<unsafe extern fn(*mut u8)>) -> Key {
-    panic!("TLS on wasm with atomics not implemented yet");
+pub unsafe fn create(dtor: Option<unsafe extern fn(*mut u8)>) -> Key {
+    drop(dtor); // FIXME: need to figure out how to hook thread exit to run this
+    let key = NEXT_KEY.fetch_add(1, SeqCst);
+    if key >= MAX_KEYS {
+        NEXT_KEY.store(MAX_KEYS, SeqCst);
+        panic!("cannot allocate space for more TLS keys");
+    }
+    // offset by 1 so we never hand out 0. This is currently required by
+    // `sys_common/thread_local.rs` where it can't cope with keys of value 0
+    // because it messes up the atomic management.
+    return key + 1
 }
 
-pub unsafe fn set(_key: Key, _value: *mut u8) {
-    panic!("TLS on wasm with atomics not implemented yet");
+pub unsafe fn set(key: Key, value: *mut u8) {
+    (*ThreadControlBlock::get()).keys[key - 1] = value;
 }
 
-pub unsafe fn get(_key: Key) -> *mut u8 {
-    panic!("TLS on wasm with atomics not implemented yet");
+pub unsafe fn get(key: Key) -> *mut u8 {
+    (*ThreadControlBlock::get()).keys[key - 1]
 }
 
 pub unsafe fn destroy(_key: Key) {
-    panic!("TLS on wasm with atomics not implemented yet");
+    // FIXME: should implement this somehow, this isn't typically called but it
+    // can be called if two threads race to initialize a TLS slot and one ends
+    // up not being needed.
 }
 
 #[inline]

Diff for: src/libstd/thread/local.rs

@@ -172,16 +172,22 @@ macro_rules! __thread_local_inner {
                 &'static $crate::cell::UnsafeCell<
                     $crate::option::Option<$t>>>
             {
-                #[cfg(target_arch = "wasm32")]
+                #[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))]
                 static __KEY: $crate::thread::__StaticLocalKeyInner<$t> =
                     $crate::thread::__StaticLocalKeyInner::new();
 
                 #[thread_local]
-                #[cfg(all(target_thread_local, not(target_arch = "wasm32")))]
+                #[cfg(all(
+                    target_thread_local,
+                    not(all(target_arch = "wasm32", not(target_feature = "atomics"))),
+                ))]
                 static __KEY: $crate::thread::__FastLocalKeyInner<$t> =
                     $crate::thread::__FastLocalKeyInner::new();
 
-                #[cfg(all(not(target_thread_local), not(target_arch = "wasm32")))]
+                #[cfg(all(
+                    not(target_thread_local),
+                    not(all(target_arch = "wasm32", not(target_feature = "atomics"))),
+                ))]
                 static __KEY: $crate::thread::__OsLocalKeyInner<$t> =
                     $crate::thread::__OsLocalKeyInner::new();
 
@@ -302,7 +308,7 @@ impl<T: 'static> LocalKey<T> {
 /// On some platforms like wasm32 there's no threads, so no need to generate
 /// thread locals and we can instead just use plain statics!
 #[doc(hidden)]
-#[cfg(target_arch = "wasm32")]
+#[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))]
 pub mod statik {
     use cell::UnsafeCell;
     use fmt;

Diff for: src/libstd/thread/mod.rs

@@ -203,7 +203,7 @@ pub use self::local::{LocalKey, AccessError};
 // where available, but both are needed.
 
 #[unstable(feature = "libstd_thread_internals", issue = "0")]
-#[cfg(target_arch = "wasm32")]
+#[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))]
 #[doc(hidden)] pub use self::local::statik::Key as __StaticLocalKeyInner;
 #[unstable(feature = "libstd_thread_internals", issue = "0")]
 #[cfg(target_thread_local)]