extension_host: Fix `IS_WASM_THREAD` being set for wrong threads (#43005)

Lukas Wirth created 1 month ago

https://github.com/zed-industries/zed/pull/40883 implemented this
incorrectly. It was marking a random background thread as a wasm thread
(whatever thread picked up the wasm epoch timer background task),
instead of marking the threads that actually run the wasm extension.

This has two implications:
1. it didn't prevent extension panics from tearing down as planned
2. Worse, it actually made us hide legit panics in sentry for one of our
background workers.

Now 2 still technically applies for all tokio threads after this, but we
basically only use these for wasm extensions in the main zed binary.

Release Notes:

- Fixed extension panics crashing Zed on Linux

Change summary

crates/crashes/src/crashes.rs          |  8 ++------
crates/extension_host/src/wasm_host.rs | 11 ++++++++---
2 files changed, 10 insertions(+), 9 deletions(-)

Detailed changes

crates/crashes/src/crashes.rs

@@ -289,15 +289,11 @@ impl minidumper::ServerHandler for CrashServer {
 pub fn panic_hook(info: &PanicHookInfo) {
     // Don't handle a panic on threads that are not relevant to the main execution.
     if extension_host::wasm_host::IS_WASM_THREAD.with(|v| v.load(Ordering::Acquire)) {
+        log::error!("wasm thread panicked!");
         return;
     }
 
-    let message = info
-        .payload()
-        .downcast_ref::<&str>()
-        .map(|s| s.to_string())
-        .or_else(|| info.payload().downcast_ref::<String>().cloned())
-        .unwrap_or_else(|| "Box<Any>".to_string());
+    let message = info.payload_as_str().unwrap_or("Box<Any>").to_owned();
 
     let span = info
         .location()

crates/extension_host/src/wasm_host.rs

@@ -537,7 +537,6 @@ fn wasm_engine(executor: &BackgroundExecutor) -> wasmtime::Engine {
             let engine_ref = engine.weak();
             executor
                 .spawn(async move {
-                    IS_WASM_THREAD.with(|v| v.store(true, Ordering::Release));
                     // Somewhat arbitrary interval, as it isn't a guaranteed interval.
                     // But this is a rough upper bound for how long the extension execution can block on
                     // `Future::poll`.
@@ -643,6 +642,12 @@ impl WasmHost {
 
             let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
             let extension_task = async move {
+                // note: Setting the thread local here will slowly "poison" all tokio threads
+                // causing us to not record their panics any longer.
+                //
+                // This is fine though, the main zed binary only uses tokio for livekit and wasm extensions.
+                // Livekit seldom (if ever) panics 🤞 so the likelihood of us missing a panic in sentry is very low.
+                IS_WASM_THREAD.with(|v| v.store(true, Ordering::Release));
                 while let Some(call) = rx.next().await {
                     (call)(&mut extension, &mut store).await;
                 }
@@ -659,8 +664,8 @@ impl WasmHost {
         cx.spawn(async move |cx| {
             let (extension_task, manifest, work_dir, tx, zed_api_version) =
                 cx.background_executor().spawn(load_extension_task).await?;
-            // we need to run run the task in an extension context as wasmtime_wasi may
-            // call into tokio, accessing its runtime handle
+            // we need to run run the task in a tokio context as wasmtime_wasi may
+            // call into tokio, accessing its runtime handle when we trigger the `engine.increment_epoch()` above.
             let task = Arc::new(gpui_tokio::Tokio::spawn(cx, extension_task)?);
 
             Ok(WasmExtension {