Added a thread_local guard to trigger dtor hook before fibers are used (#12426)

* Added a thread_local guard to trigger dtor hook before fibers are used

* Add a test for async function that use a tls variable with a destructor

Author: ohadravid

crates/fiber/src/windows.rs

@@ -3,6 +3,7 @@ use alloc::boxed::Box;
 use std::cell::Cell;
 use std::ffi::c_void;
 use std::io;
+use std::mem::needs_drop;
 use std::ops::Range;
 use std::ptr;
 use windows_sys::Win32::Foundation::*;
@@ -128,6 +129,21 @@ impl Fiber {
             let parent_fiber = if is_fiber {
                 wasmtime_fiber_get_current()
             } else {
+                // Newer Rust versions use fiber local storage to register an internal hook that
+                // calls thread locals' destructors on thread exit.
+                // This has a limitation: the hook only runs in a regular thread (not in a fiber).
+                // We convert back into a thread once execution returns to this function,
+                // but we must also ensure that the hook is registered before converting into a fiber.
+                // Otherwise, a different fiber could be the first to register the hook,
+                // causing the hook to be called (and skipped) prematurely when that fiber is deleted.
+                struct Guard;
+
+                impl Drop for Guard {
+                    fn drop(&mut self) {}
+                }
+                assert!(needs_drop::<Guard>());
+                thread_local!(static GUARD: Guard = Guard);
+                GUARD.with(|_g| {});
                 ConvertThreadToFiber(ptr::null_mut())
             };
             assert!(

tests/all/async_functions.rs

@@ -1,5 +1,7 @@
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll, Waker};
 use wasmtime::*;
@@ -615,6 +617,49 @@ async fn resume_separate_thread3() {
     assert!(f.call(&mut store, &[], &mut []).is_err());
 }
 
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn resume_separate_thread_tls() {
+    static COUNTER: AtomicUsize = AtomicUsize::new(0);
+
+    struct IncOnDrop;
+    impl Drop for IncOnDrop {
+        fn drop(&mut self) {
+            COUNTER.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    thread_local!(static FOO: IncOnDrop = IncOnDrop);
+
+    // This test will poll the following future on two threads.
+    // We verify that TLS destructors are run correctly.
+    execute_across_threads(async move {
+        let mut store = async_store();
+        let module = Module::new(
+            store.engine(),
+            "
+            (module
+                (import \"\" \"\" (func))
+                (start 0)
+            )
+            ",
+        )
+        .unwrap();
+        let func = Func::wrap_async(&mut store, |_, _: ()| {
+            Box::new(async {
+                tokio::task::yield_now().await;
+                FOO.with(|_f| {});
+                Err::<(), _>(format_err!("test"))
+            })
+        });
+        let result = Instance::new_async(&mut store, &module, &[func.into()]).await;
+        assert!(result.is_err());
+    })
+    .await;
+
+    assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
+}
+
 #[tokio::test]
 #[cfg_attr(miri, ignore)]
 async fn recursive_async() -> Result<()> {