extension_host: Run extensions on the tokio threadpool (#40936)

Lukas Wirth created 2 months ago

Fixes ZED-12D

`wasmtime_wasi` might call into tokio futures (to sleep for example)
which requires access to the tokio runtime. So we are required to run
these extensions in the tokio thread pool

Release Notes:

- Fixed extensions causing zed to occasionally panic

Change summary

Cargo.lock                                                       |  1 
crates/extension_host/Cargo.toml                                 |  1 
crates/extension_host/benches/extension_compilation_benchmark.rs |  3 
crates/extension_host/src/extension_store_test.rs                |  1 
crates/extension_host/src/wasm_host.rs                           | 42 +
5 files changed, 31 insertions(+), 17 deletions(-)

Detailed changes

Cargo.lock

@@ -5883,6 +5883,7 @@ dependencies = [
  "fs",
  "futures 0.3.31",
  "gpui",
+ "gpui_tokio",
  "http_client",
  "language",
  "language_extension",

crates/extension_host/Cargo.toml

@@ -27,6 +27,7 @@ extension.workspace = true
 fs.workspace = true
 futures.workspace = true
 gpui.workspace = true
+gpui_tokio.workspace = true
 http_client.workspace = true
 language.workspace = true
 log.workspace = true

crates/extension_host/benches/extension_compilation_benchmark.rs

@@ -19,6 +19,7 @@ use util::test::TempTree;
 
 fn extension_benchmarks(c: &mut Criterion) {
     let cx = init();
+    cx.update(gpui_tokio::init);
 
     let mut group = c.benchmark_group("load");
 
@@ -37,7 +38,7 @@ fn extension_benchmarks(c: &mut Criterion) {
             |wasm_bytes| {
                 let _extension = cx
                     .executor()
-                    .block(wasm_host.load_extension(wasm_bytes, &manifest, cx.executor()))
+                    .block(wasm_host.load_extension(wasm_bytes, &manifest, &cx.to_async()))
                     .unwrap();
             },
             BatchSize::SmallInput,

crates/extension_host/src/extension_store_test.rs

@@ -868,5 +868,6 @@ fn init_test(cx: &mut TestAppContext) {
         Project::init_settings(cx);
         ExtensionSettings::register(cx);
         language::init(cx);
+        gpui_tokio::init(cx);
     });
 }

crates/extension_host/src/wasm_host.rs

@@ -591,11 +591,12 @@ impl WasmHost {
         self: &Arc<Self>,
         wasm_bytes: Vec<u8>,
         manifest: &Arc<ExtensionManifest>,
-        executor: BackgroundExecutor,
+        cx: &AsyncApp,
     ) -> Task<Result<WasmExtension>> {
         let this = self.clone();
         let manifest = manifest.clone();
-        executor.clone().spawn(async move {
+        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)
@@ -632,20 +633,29 @@ impl WasmHost {
                 .context("failed to initialize wasm extension")?;
 
             let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
-            executor
-                .spawn(async move {
-                    while let Some(call) = rx.next().await {
-                        (call)(&mut extension, &mut store).await;
-                    }
-                })
-                .detach();
+            let extension_task = async move {
+                while let Some(call) = rx.next().await {
+                    (call)(&mut extension, &mut store).await;
+                }
+            };
 
-            Ok(WasmExtension {
-                manifest: manifest.clone(),
-                work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
-                tx,
-                zed_api_version,
-            })
+            anyhow::Ok((
+                extension_task,
+                WasmExtension {
+                    manifest: manifest.clone(),
+                    work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
+                    tx,
+                    zed_api_version,
+                },
+            ))
+        };
+        cx.spawn(async move |cx| {
+            let (extension_task, extension) = 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
+            gpui_tokio::Tokio::spawn(cx, extension_task)?.detach();
+
+            Ok(extension)
         })
     }
 
@@ -747,7 +757,7 @@ impl WasmExtension {
             .context("failed to read wasm")?;
 
         wasm_host
-            .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
+            .load_extension(wasm_bytes, manifest, cx)
             .await
             .with_context(|| format!("failed to load wasm extension {}", manifest.id))
     }