std: refactor Windows TLS destructor support

joboet · joboet · commit afc6eb2fcd82 · 2023-11-21T18:42:29.000+01:00
Now that the fallback code has been removed, there are no users of `StaticKey` left for targets with native TLS. Therefore, move the at-exit hack to the thread-local guard module, where it can be shared by both implementations, and cfg-out the key-based TLS when it's not needed.
diff --git a/library/std/src/sys/windows/c.rs b/library/std/src/sys/windows/c.rs
@@ -52,6 +52,7 @@ pub const EXIT_FAILURE: u32 = 1;
 
 pub const CONDITION_VARIABLE_INIT: CONDITION_VARIABLE = CONDITION_VARIABLE { Ptr: ptr::null_mut() };
 pub const SRWLOCK_INIT: SRWLOCK = SRWLOCK { Ptr: ptr::null_mut() };
+#[cfg(not(target_thread_local))] // Only used by key-based TLS.
 pub const INIT_ONCE_STATIC_INIT: INIT_ONCE = INIT_ONCE { Ptr: ptr::null_mut() };
 
 // Some windows_sys types have different signs than the types we use.
diff --git a/library/std/src/sys/windows/mod.rs b/library/std/src/sys/windows/mod.rs
@@ -31,7 +31,7 @@ pub mod process;
 pub mod rand;
 pub mod stdio;
 pub mod thread;
-pub mod thread_local_dtor;
+pub mod thread_local_guard;
 pub mod thread_local_key;
 pub mod thread_parking;
 pub mod time;
diff --git a/library/std/src/sys/windows/thread_local_dtor.rs b/library/std/src/sys/windows/thread_local_dtor.rs
diff --git a/library/std/src/sys/windows/thread_local_guard.rs b/library/std/src/sys/windows/thread_local_guard.rs
@@ -0,0 +1,120 @@
+//! A TLS destructor system.
+//!
+//! Turns out, like pretty much everything, Windows is pretty close the
+//! functionality that Unix provides, but slightly different! In the case of
+//! TLS, Windows does not provide an API to provide a destructor for a TLS
+//! variable. This ends up being pretty crucial to this implementation, so we
+//! need a way around this.
+//!
+//! The solution here ended up being a little obscure, but fear not, the
+//! internet has informed me [1][2] that this solution is not unique (no way
+//! I could have thought of it as well!). The key idea is to insert some hook
+//! somewhere to run arbitrary code on thread termination. With this in place
+//! we'll be able to run anything we like, including all TLS destructors!
+//!
+//! If you're looking at this code, and wondering "what is this doing?",
+//! you're not alone! I'll try to break this down step by step:
+//!
+//! # What's up with CRT$XLB?
+//!
+//! For anything about TLS destructors to work on Windows, we have to be able
+//! to run *something* when a thread exits. To do so, we place a very special
+//! static in a very special location. If this is encoded in just the right
+//! way, the kernel's loader is apparently nice enough to run some function
+//! of ours whenever a thread exits! How nice of the kernel!
+//!
+//! Lots of detailed information can be found in source [1] above, but the
+//! gist of it is that this is leveraging a feature of Microsoft's PE format
+//! (executable format) which is not actually used by any compilers today.
+//! This apparently translates to any callbacks in the ".CRT$XLB" section
+//! being run on certain events.
+//!
+//! So after all that, we use the compiler's #[link_section] feature to place
+//! a callback pointer into the magic section so it ends up being called.
+//!
+//! # What's up with this callback?
+//!
+//! The callback specified receives a number of parameters from... someone!
+//! (the kernel? the runtime? I'm not quite sure!) There are a few events that
+//! this gets invoked for, but we're currently only interested on when a
+//! thread or a process "detaches" (exits). The process part happens for the
+//! last thread and the thread part happens for any normal thread.
+//!
+//! # The article mentions weird stuff about "/INCLUDE"?
+//!
+//! It sure does! Specifically we're talking about this quote:
+//!
+//! > The Microsoft run-time library facilitates this process by defining a
+//! > memory image of the TLS Directory and giving it the special name
+//! > “__tls_used” (Intel x86 platforms) or “_tls_used” (other platforms). The
+//! > linker looks for this memory image and uses the data there to create the
+//! > TLS Directory. Other compilers that support TLS and work with the
+//! > Microsoft linker must use this same technique.
+//!
+//! Basically what this means is that if we want support for our TLS
+//! destructors/our hook being called then we need to make sure the linker does
+//! not omit this symbol. Otherwise it will omit it and our callback won't be
+//! wired up.
+//!
+//! We don't actually use the `/INCLUDE` linker flag here like the article
+//! mentions because the Rust compiler doesn't propagate linker flags, but
+//! instead we use a shim function which performs a volatile 1-byte load from
+//! the address of the symbol to ensure it sticks around.
+//!
+//! [1]: https://www.codeproject.com/Articles/8113/Thread-Local-Storage-The-C-Way
+//! [2]: https://github.com/ChromiumWebApps/chromium/blob/master/base/threading/thread_local_storage_win.cc#L42
+
+#![unstable(feature = "thread_local_internals", issue = "none")]
+
+use crate::ptr;
+use crate::sync::atomic::{
+    AtomicBool,
+    Ordering::{Acquire, Relaxed},
+};
+use crate::sys::c;
+
+// If the target uses native TLS, run its destructors.
+#[cfg(target_thread_local)]
+use crate::sys::common::thread_local::run_dtors;
+// Otherwise, run the destructors for the key-based variant.
+#[cfg(not(target_thread_local))]
+use super::thread_local_key::run_dtors;
+
+/// An optimization hint. The compiler is often smart enough to know if an atomic
+/// is never set and can remove dead code based on that fact.
+static HAS_DTORS: AtomicBool = AtomicBool::new(false);
+
+/// Ensure that thread-locals are destroyed when the thread exits.
+pub fn activate() {
+    HAS_DTORS.store(true, Relaxed);
+}
+
+#[link_section = ".CRT$XLB"]
+#[allow(dead_code, unused_variables)]
+#[used] // we don't want LLVM eliminating this symbol for any reason, and
+// when the symbol makes it to the linker the linker will take over
+pub static p_thread_callback: unsafe extern "system" fn(c::LPVOID, c::DWORD, c::LPVOID) =
+    on_tls_callback;
+
+#[allow(dead_code, unused_variables)]
+unsafe extern "system" fn on_tls_callback(h: c::LPVOID, dwReason: c::DWORD, pv: c::LPVOID) {
+    if !HAS_DTORS.load(Acquire) {
+        return;
+    }
+    if dwReason == c::DLL_THREAD_DETACH || dwReason == c::DLL_PROCESS_DETACH {
+        run_dtors(ptr::null_mut());
+    }
+
+    // See comments above for what this is doing. Note that we don't need this
+    // trickery on GNU windows, just on MSVC.
+    reference_tls_used();
+    #[cfg(target_env = "msvc")]
+    unsafe fn reference_tls_used() {
+        extern "C" {
+            static _tls_used: u8;
+        }
+        crate::intrinsics::volatile_load(&_tls_used);
+    }
+    #[cfg(not(target_env = "msvc"))]
+    unsafe fn reference_tls_used() {}
+}
diff --git a/library/std/src/sys/windows/thread_local_key.rs b/library/std/src/sys/windows/thread_local_key.rs
@@ -1,89 +1,19 @@
+#![cfg(not(target_thread_local))]
+
 use crate::cell::UnsafeCell;
 use crate::ptr;
 use crate::sync::atomic::{
-    AtomicBool, AtomicPtr, AtomicU32,
+    AtomicPtr, AtomicU32,
     Ordering::{AcqRel, Acquire, Relaxed, Release},
 };
 use crate::sys::c;
 
 #[cfg(test)]
 mod tests;
 
-/// An optimization hint. The compiler is often smart enough to know if an atomic
-/// is never set and can remove dead code based on that fact.
-static HAS_DTORS: AtomicBool = AtomicBool::new(false);
-
-// Using a per-thread list avoids the problems in synchronizing global state.
-#[thread_local]
-#[cfg(target_thread_local)]
-static DESTRUCTORS: crate::cell::RefCell<Vec<(*mut u8, unsafe extern "C" fn(*mut u8))>> =
-    crate::cell::RefCell::new(Vec::new());
-
-// Ensure this can never be inlined because otherwise this may break in dylibs.
-// See #44391.
-#[inline(never)]
-#[cfg(target_thread_local)]
-pub unsafe fn register_keyless_dtor(t: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
-    match DESTRUCTORS.try_borrow_mut() {
-        Ok(mut dtors) => dtors.push((t, dtor)),
-        Err(_) => rtabort!("global allocator may not use TLS"),
-    }
-
-    HAS_DTORS.store(true, Relaxed);
-}
-
-#[inline(never)] // See comment above
-#[cfg(target_thread_local)]
-/// Runs destructors. This should not be called until thread exit.
-unsafe fn run_keyless_dtors() {
-    // Drop all the destructors.
-    //
-    // Note: While this is potentially an infinite loop, it *should* be
-    // the case that this loop always terminates because we provide the
-    // guarantee that a TLS key cannot be set after it is flagged for
-    // destruction.
-    loop {
-        // Use a let-else binding to ensure the `RefCell` guard is dropped
-        // immediately. Otherwise, a panic would occur if a TLS destructor
-        // tries to access the list.
-        let Some((ptr, dtor)) = DESTRUCTORS.borrow_mut().pop() else {
-            break;
-        };
-        (dtor)(ptr);
-    }
-    // We're done so free the memory.
-    DESTRUCTORS.replace(Vec::new());
-}
-
 type Key = c::DWORD;
 type Dtor = unsafe extern "C" fn(*mut u8);
 
-// Turns out, like pretty much everything, Windows is pretty close the
-// functionality that Unix provides, but slightly different! In the case of
-// TLS, Windows does not provide an API to provide a destructor for a TLS
-// variable. This ends up being pretty crucial to this implementation, so we
-// need a way around this.
-//
-// The solution here ended up being a little obscure, but fear not, the
-// internet has informed me [1][2] that this solution is not unique (no way
-// I could have thought of it as well!). The key idea is to insert some hook
-// somewhere to run arbitrary code on thread termination. With this in place
-// we'll be able to run anything we like, including all TLS destructors!
-//
-// To accomplish this feat, we perform a number of threads, all contained
-// within this module:
-//
-// * All TLS destructors are tracked by *us*, not the Windows runtime. This
-//   means that we have a global list of destructors for each TLS key that
-//   we know about.
-// * When a thread exits, we run over the entire list and run dtors for all
-//   non-null keys. This attempts to match Unix semantics in this regard.
-//
-// For more details and nitty-gritty, see the code sections below!
-//
-// [1]: https://www.codeproject.com/Articles/8113/Thread-Local-Storage-The-C-Way
-// [2]: https://github.com/ChromiumWebApps/chromium/blob/master/base/threading/thread_local_storage_win.cc#L42
-
 pub struct StaticKey {
     /// The key value shifted up by one. Since TLS_OUT_OF_INDEXES == DWORD::MAX
     /// is not a valid key value, this allows us to use zero as sentinel value
@@ -215,41 +145,10 @@ unsafe fn register_dtor(key: &'static StaticKey) {
             Err(new) => head = new,
         }
     }
-    HAS_DTORS.store(true, Release);
+    super::thread_local_guard::activate();
 }
 
-// -------------------------------------------------------------------------
-// Where the Magic (TM) Happens
-//
-// If you're looking at this code, and wondering "what is this doing?",
-// you're not alone! I'll try to break this down step by step:
-//
-// # What's up with CRT$XLB?
-//
-// For anything about TLS destructors to work on Windows, we have to be able
-// to run *something* when a thread exits. To do so, we place a very special
-// static in a very special location. If this is encoded in just the right
-// way, the kernel's loader is apparently nice enough to run some function
-// of ours whenever a thread exits! How nice of the kernel!
-//
-// Lots of detailed information can be found in source [1] above, but the
-// gist of it is that this is leveraging a feature of Microsoft's PE format
-// (executable format) which is not actually used by any compilers today.
-// This apparently translates to any callbacks in the ".CRT$XLB" section
-// being run on certain events.
-//
-// So after all that, we use the compiler's #[link_section] feature to place
-// a callback pointer into the magic section so it ends up being called.
-//
-// # What's up with this callback?
-//
-// The callback specified receives a number of parameters from... someone!
-// (the kernel? the runtime? I'm not quite sure!) There are a few events that
-// this gets invoked for, but we're currently only interested on when a
-// thread or a process "detaches" (exits). The process part happens for the
-// last thread and the thread part happens for any normal thread.
-//
-// # Ok, what's up with running all these destructors?
+// What's up with running all these destructors?
 //
 // This will likely need to be improved over time, but this function
 // attempts a "poor man's" destructor callback system. Once we've got a list
@@ -258,63 +157,7 @@ unsafe fn register_dtor(key: &'static StaticKey) {
 // beforehand). We do this a few times in a loop to basically match Unix
 // semantics. If we don't reach a fixed point after a short while then we just
 // inevitably leak something most likely.
-//
-// # The article mentions weird stuff about "/INCLUDE"?
-//
-// It sure does! Specifically we're talking about this quote:
-//
-//      The Microsoft run-time library facilitates this process by defining a
-//      memory image of the TLS Directory and giving it the special name
-//      “__tls_used” (Intel x86 platforms) or “_tls_used” (other platforms). The
-//      linker looks for this memory image and uses the data there to create the
-//      TLS Directory. Other compilers that support TLS and work with the
-//      Microsoft linker must use this same technique.
-//
-// Basically what this means is that if we want support for our TLS
-// destructors/our hook being called then we need to make sure the linker does
-// not omit this symbol. Otherwise it will omit it and our callback won't be
-// wired up.
-//
-// We don't actually use the `/INCLUDE` linker flag here like the article
-// mentions because the Rust compiler doesn't propagate linker flags, but
-// instead we use a shim function which performs a volatile 1-byte load from
-// the address of the symbol to ensure it sticks around.
-
-#[link_section = ".CRT$XLB"]
-#[allow(dead_code, unused_variables)]
-#[used] // we don't want LLVM eliminating this symbol for any reason, and
-// when the symbol makes it to the linker the linker will take over
-pub static p_thread_callback: unsafe extern "system" fn(c::LPVOID, c::DWORD, c::LPVOID) =
-    on_tls_callback;
-
-#[allow(dead_code, unused_variables)]
-unsafe extern "system" fn on_tls_callback(h: c::LPVOID, dwReason: c::DWORD, pv: c::LPVOID) {
-    if !HAS_DTORS.load(Acquire) {
-        return;
-    }
-    if dwReason == c::DLL_THREAD_DETACH || dwReason == c::DLL_PROCESS_DETACH {
-        #[cfg(not(target_thread_local))]
-        run_dtors();
-        #[cfg(target_thread_local)]
-        run_keyless_dtors();
-    }
-
-    // See comments above for what this is doing. Note that we don't need this
-    // trickery on GNU windows, just on MSVC.
-    reference_tls_used();
-    #[cfg(target_env = "msvc")]
-    unsafe fn reference_tls_used() {
-        extern "C" {
-            static _tls_used: u8;
-        }
-        crate::intrinsics::volatile_load(&_tls_used);
-    }
-    #[cfg(not(target_env = "msvc"))]
-    unsafe fn reference_tls_used() {}
-}
-
-#[allow(dead_code)] // actually called below
-unsafe fn run_dtors() {
+pub(super) unsafe fn run_dtors(_ptr: *mut u8) {
     for _ in 0..5 {
         let mut any_run = false;
 
diff --git a/library/std/src/sys_common/mod.rs b/library/std/src/sys_common/mod.rs
@@ -35,6 +35,7 @@ pub mod wtf8;
 
 cfg_if::cfg_if! {
     if #[cfg(target_os = "windows")] {
+        #[cfg(not(target_thread_local))]
         pub use crate::sys::thread_local_key;
     } else {
         pub mod thread_local_key;
