extension_host: Avoid creating a Tokio runtime (#45990)

Marco Mihai Condrache created 3 months ago

While profiling Zed, I noticed around 12 `tokio-runtime-worker` threads.
Since `gpui_tokio` runs a Tokio runtime with only two worker threads,
this meant something else was spinning up a default Tokio runtime (My
machine has 10 cores).

After digging into it a bit, it looks like `wasmtime-wasi` needs a Tokio
runtime for some of its calls. That’s also why we run extension tasks on
`gpui-tokio` in the first place.

Reproduction case:

1. Run a clean Zed build.
2. Sample the process → only 2–3 Tokio threads.
3. Install a WASM extension (for example
https://github.com/zed-extensions/nix), which shouldn’t pull in anything
that creates its own runtime.
4. Sample the process again → 10+ Tokio worker threads show up.

```
marcocondrache@xawed ~/P/zed (main)> sample 34260 1 | grep -E "Thread_"
Sampling process 34260 for 1 second with 1 millisecond of run time between samples
Sampling completed, processing symbols...
Sample analysis of process 34260 written to file /tmp/zed_2026-01-03_112527_5pq7.sample.txt
    744 Thread_3086626   DispatchQueue_1: com.apple.main-thread  (serial)
    744 Thread_3086629: RayonWorker0
    ....
    744 Thread_3086651: MacWatcher
    744 Thread_3086653: MacWatcher
    744 Thread_3086661: com.apple.NSEventThread
    744 Thread_3086680: tokio-runtime-worker
    744 Thread_3086681: tokio-runtime-worker
    ...
    744 Thread_3087087: tokio-runtime-worker
    744 Thread_3087089: tokio-runtime-worker
    744 Thread_3087090: tokio-runtime-worker
    744 Thread_3087091: tokio-runtime-worker
    744 Thread_3087092: tokio-runtime-worker
    744 Thread_3087093: tokio-runtime-worker
    744 Thread_3087094: tokio-runtime-worker
    744 Thread_3087095: tokio-runtime-worker
    744 Thread_3087096: tokio-runtime-worker
    744 Thread_3087097: tokio-runtime-worker
    ...
```

Release Notes:

- Reduced number of threads when using extensions

---------

Signed-off-by: Marco Mihai Condrache <[e-mail]>

Change summary

crates/extension_host/src/wasm_host.rs | 35 +++++++++++++++++++++------
1 file changed, 27 insertions(+), 8 deletions(-)

Detailed changes

crates/extension_host/src/wasm_host.rs

@@ -604,15 +604,28 @@ impl WasmHost {
         let this = self.clone();
         let manifest = manifest.clone();
         let executor = cx.background_executor().clone();
-        let load_extension_task = async move {
-            let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
 
-            let component = Component::from_binary(&this.engine, &wasm_bytes)
-                .context("failed to compile wasm component")?;
+        // Parse version and compile component on gpui's background executor.
+        // These are cpu-bound operations that don't require a tokio runtime.
+        let compile_task = {
+            let manifest_id = manifest.id.clone();
+            let engine = this.engine.clone();
+
+            executor.spawn(async move {
+                let zed_api_version = parse_wasm_extension_version(&manifest_id, &wasm_bytes)?;
+                let component = Component::from_binary(&engine, &wasm_bytes)
+                    .context("failed to compile wasm component")?;
+
+                anyhow::Ok((zed_api_version, component))
+            })
+        };
+
+        let load_extension = |zed_api_version: Version, component| async move {
+            let wasi_ctx = this.build_wasi_ctx(&manifest).await?;
             let mut store = wasmtime::Store::new(
                 &this.engine,
                 WasmState {
-                    ctx: this.build_wasi_ctx(&manifest).await?,
+                    ctx: wasi_ctx,
                     manifest: manifest.clone(),
                     table: ResourceTable::new(),
                     host: this.clone(),
@@ -661,11 +674,17 @@ impl WasmHost {
                 zed_api_version,
             ))
         };
+
         cx.spawn(async move |cx| {
+            let (zed_api_version, component) = compile_task.await?;
+
+            // Run wasi-dependent operations on tokio.
+            // wasmtime_wasi internally uses tokio for I/O operations.
             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 a tokio context as wasmtime_wasi may
-            // call into tokio, accessing its runtime handle when we trigger the `engine.increment_epoch()` above.
+                gpui_tokio::Tokio::spawn(cx, load_extension(zed_api_version, component))?.await??;
+
+            // Run the extension message loop on tokio since extension
+            // calls may invoke wasi functions that require a tokio runtime.
             let task = Arc::new(gpui_tokio::Tokio::spawn(cx, extension_task)?);
 
             Ok(WasmExtension {