Handle wasip3 task-specific information differently

ci/rebuild-libwit-bindgen-cabi.sh

@@ -113,11 +113,22 @@ cat >./crates/guest-rust/src/rt/wit_bindgen_cabi_wasip3_tls.c <<-EOF
 
 static _Thread_local void *task_context = NULL;
 
+__attribute__((weak))
+extern void *__wasilibc_task_context_get(void);
+__attribute__((weak))
+extern void __wasilibc_task_context_set(void*);
+
 void *$wasip3_get(void) {
+  if (__wasilibc_task_context_get)
+    return __wasilibc_task_context_get();
   return task_context;
 }
 
 void $wasip3_set(void *ptr) {
+  if (__wasilibc_task_context_set) {
+    __wasilibc_task_context_set(ptr);
+    return;
+  }
   task_context = ptr;
 }
 EOF

crates/guest-rust/build.rs

@@ -18,10 +18,7 @@ fn main() {
         link_lib("cabi_realloc");
     }
     link_lib("cabi_wasip3");
-
-    // The no-std async runtime's task context uses WASIp3 TLS. Keep this
-    // separate from the MVP task helper needed by every wasm32 target.
-    if target_env == "p3" && !cfg!(feature = "std") {
+    if target_env == "p3" {
         link_lib("cabi_wasip3_tls");
     }
 }

crates/guest-rust/src/rt/async_support.rs

@@ -579,22 +579,14 @@ mod task_state {
     }
 }
 
-#[cfg(all(target_env = "p3", feature = "std"))]
-mod task_state {
-    std::thread_local!(static TASK_STATE: std::cell::Cell<*mut u8> = std::cell::Cell::new(std::ptr::null_mut()));
-
-    pub fn get() -> *mut u8 {
-        TASK_STATE.with(|c| c.get())
-    }
-
-    pub unsafe fn set(value: *mut u8) {
-        TASK_STATE.with(|c| c.set(value))
-    }
-}
-
-#[cfg(all(target_env = "p3", not(feature = "std")))]
+// TODO: in the future once wasi-sdk-35 is released with
+// WebAssembly/wasi-libc#884 this `task_state` module can unconditionally work
+// with `__wasilibc_task_context_{get,set}`. Until that time though detecting
+// wasi-sdk-35 is done with with a weak symbol import which is not possible to
+// write on stable Rust right now.
+#[cfg(target_env = "p3")]
 mod wasip3_context;
-#[cfg(all(target_env = "p3", not(feature = "std")))]
+#[cfg(target_env = "p3")]
 mod task_state {
 
     pub fn get() -> *mut u8 {

crates/guest-rust/src/rt/wit_bindgen_cabi_wasip3_tls.c

@@ -4,10 +4,21 @@
 
 static _Thread_local void *task_context = NULL;
 
+__attribute__((weak))
+extern void *__wasilibc_task_context_get(void);
+__attribute__((weak))
+extern void __wasilibc_task_context_set(void*);
+
 void *wit_bindgen_0_63_0_get_task_context(void) {
+  if (__wasilibc_task_context_get)
+    return __wasilibc_task_context_get();
   return task_context;
 }
 
 void wit_bindgen_0_63_0_set_task_context(void *ptr) {
+  if (__wasilibc_task_context_set) {
+    __wasilibc_task_context_set(ptr);
+    return;
+  }
   task_context = ptr;
 }