extension_host: Turn on parallel compilation (#30942)

Precursor to other optimizations, but this already gets us a big
improvement.

Wasm compilation can easily be parallelized, and with all of the cores
on my M4 Max this already gets us an 86% improvement, bringing loading
an extension down to <9ms.

Not all setups will see this much improvement, but it will use the cores
available (it just uses rayon under the hood like we do elsewhere).
Since we load extensions in sequence, this should have a nice impact for
users with a lot of extensions.

#### Before

```
Benchmarking load: Warming up for 3.0000 s
Warning: Unable to complete 100 samples in 5.0s. You may wish to increase target time to 6.5s, or reduce sample count to 70.
load                    time:   [64.859 ms 64.935 ms 65.027 ms]
Found 8 outliers among 100 measurements (8.00%)
  2 (2.00%) low mild
  3 (3.00%) high mild
  3 (3.00%) high severe
```

#### After

```
load                    time:   [8.8685 ms 8.9012 ms 8.9344 ms]
                        change: [-86.347% -86.292% -86.237%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 2 outliers among 100 measurements (2.00%)
  2 (2.00%) high mild
```

Release Notes:

- N/A

Cargo.lock

@@ -5137,6 +5137,7 @@ dependencies = [
  "async-trait",
  "client",
  "collections",
+ "criterion",
  "ctor",
  "dap",
  "env_logger 0.11.8",
@@ -5153,6 +5154,7 @@ dependencies = [
  "parking_lot",
  "paths",
  "project",
+ "rand 0.8.5",
  "release_channel",
  "remote",
  "reqwest_client",
@@ -17533,6 +17535,7 @@ dependencies = [
  "postcard",
  "psm",
  "pulley-interpreter",
+ "rayon",
  "rustix 0.38.44",
  "semver",
  "serde",

Cargo.toml

@@ -430,6 +430,7 @@ convert_case = "0.8.0"
 core-foundation = "0.10.0"
 core-foundation-sys = "0.8.6"
 core-video = { version = "0.4.3", features = ["metal"] }
+criterion = { version = "0.5", features = ["html_reports"] }
 ctor = "0.4.0"
 dap-types = { git = "https://github.com/zed-industries/dap-types", rev = "be69a016ba710191b9fdded28c8b042af4b617f7" }
 dashmap = "6.0"
@@ -608,6 +609,7 @@ wasmtime = { version = "29", default-features = false, features = [
     "runtime",
     "cranelift",
     "component-model",
+    "parallel-compilation",
 ] }
 wasmtime-wasi = "29"
 which = "6.0.0"

crates/extension/src/extension_builder.rs

@@ -398,7 +398,7 @@ impl ExtensionBuilder {
 
     async fn install_wasi_sdk_if_needed(&self) -> Result<PathBuf> {
         let url = if let Some(asset_name) = WASI_SDK_ASSET_NAME {
-            format!("{WASI_SDK_URL}/{asset_name}")
+            format!("{WASI_SDK_URL}{asset_name}")
         } else {
             bail!("wasi-sdk is not available for platform {}", env::consts::OS);
         };

crates/extension/src/extension_manifest.rs

@@ -162,7 +162,7 @@ pub struct GrammarManifestEntry {
     pub path: Option<String>,
 }
 
-#[derive(Clone, PartialEq, Eq, Debug, Deserialize, Serialize)]
+#[derive(Clone, Default, PartialEq, Eq, Debug, Deserialize, Serialize)]
 pub struct LanguageServerManifestEntry {
     /// Deprecated in favor of `languages`.
     #[serde(default)]

crates/extension_host/Cargo.toml

@@ -54,6 +54,7 @@ wasmtime.workspace = true
 workspace-hack.workspace = true
 
 [dev-dependencies]
+criterion.workspace = true
 ctor.workspace = true
 env_logger.workspace = true
 fs = { workspace = true, features = ["test-support"] }
@@ -62,6 +63,11 @@ language = { workspace = true, features = ["test-support"] }
 language_extension.workspace = true
 parking_lot.workspace = true
 project = { workspace = true, features = ["test-support"] }
+rand.workspace = true
 reqwest_client.workspace = true
 theme = { workspace = true, features = ["test-support"] }
 theme_extension.workspace = true
+
+[[bench]]
+name = "extension_compilation_benchmark"
+harness = false

crates/extension_host/benches/extension_compilation_benchmark.rs

@@ -0,0 +1,145 @@
+use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
+
+use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
+use extension::{
+    ExtensionCapability, ExtensionHostProxy, ExtensionLibraryKind, ExtensionManifest,
+    LanguageServerManifestEntry, LibManifestEntry, SchemaVersion,
+    extension_builder::{CompileExtensionOptions, ExtensionBuilder},
+};
+use extension_host::wasm_host::WasmHost;
+use fs::RealFs;
+use gpui::{SemanticVersion, TestAppContext, TestDispatcher};
+use http_client::{FakeHttpClient, Response};
+use node_runtime::NodeRuntime;
+use rand::{SeedableRng, rngs::StdRng};
+use reqwest_client::ReqwestClient;
+use serde_json::json;
+use settings::SettingsStore;
+use util::test::TempTree;
+
+fn extension_benchmarks(c: &mut Criterion) {
+    let cx = init();
+
+    let mut group = c.benchmark_group("load");
+
+    let mut manifest = manifest();
+    let wasm_bytes = wasm_bytes(&cx, &mut manifest);
+    let manifest = Arc::new(manifest);
+    let extensions_dir = TempTree::new(json!({
+        "installed": {},
+        "work": {}
+    }));
+    let wasm_host = wasm_host(&cx, &extensions_dir);
+
+    group.bench_function(BenchmarkId::from_parameter(1), |b| {
+        b.iter_batched(
+            || wasm_bytes.clone(),
+            |wasm_bytes| {
+                let _extension = cx
+                    .executor()
+                    .block(wasm_host.load_extension(wasm_bytes, &manifest, cx.executor()))
+                    .unwrap();
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+fn init() -> TestAppContext {
+    const SEED: u64 = 9999;
+    let dispatcher = TestDispatcher::new(StdRng::seed_from_u64(SEED));
+    let cx = TestAppContext::build(dispatcher, None);
+    cx.executor().allow_parking();
+    cx.update(|cx| {
+        let store = SettingsStore::test(cx);
+        cx.set_global(store);
+        release_channel::init(SemanticVersion::default(), cx);
+    });
+
+    cx
+}
+
+fn wasm_bytes(cx: &TestAppContext, manifest: &mut ExtensionManifest) -> Vec<u8> {
+    let extension_builder = extension_builder();
+    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .parent()
+        .unwrap()
+        .join("extensions/test-extension");
+    cx.executor()
+        .block(extension_builder.compile_extension(
+            &path,
+            manifest,
+            CompileExtensionOptions { release: true },
+        ))
+        .unwrap();
+    std::fs::read(path.join("extension.wasm")).unwrap()
+}
+
+fn extension_builder() -> ExtensionBuilder {
+    let user_agent = format!(
+        "Zed Extension CLI/{} ({}; {})",
+        env!("CARGO_PKG_VERSION"),
+        std::env::consts::OS,
+        std::env::consts::ARCH
+    );
+    let http_client = Arc::new(ReqwestClient::user_agent(&user_agent).unwrap());
+    // Local dir so that we don't have to download it on every run
+    let build_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("benches/.build");
+    ExtensionBuilder::new(http_client, build_dir)
+}
+
+fn wasm_host(cx: &TestAppContext, extensions_dir: &TempTree) -> Arc<WasmHost> {
+    let http_client = FakeHttpClient::create(async |_| {
+        Ok(Response::builder().status(404).body("not found".into())?)
+    });
+    let extensions_dir = extensions_dir.path().canonicalize().unwrap();
+    let work_dir = extensions_dir.join("work");
+    let fs = Arc::new(RealFs::new(None, cx.executor()));
+
+    cx.update(|cx| {
+        WasmHost::new(
+            fs,
+            http_client,
+            NodeRuntime::unavailable(),
+            Arc::new(ExtensionHostProxy::new()),
+            work_dir,
+            cx,
+        )
+    })
+}
+
+fn manifest() -> ExtensionManifest {
+    ExtensionManifest {
+        id: "test-extension".into(),
+        name: "Test Extension".into(),
+        version: "0.1.0".into(),
+        schema_version: SchemaVersion(1),
+        description: Some("An extension for use in tests.".into()),
+        authors: Vec::new(),
+        repository: None,
+        themes: Default::default(),
+        icon_themes: Vec::new(),
+        lib: LibManifestEntry {
+            kind: Some(ExtensionLibraryKind::Rust),
+            version: Some(SemanticVersion::new(0, 1, 0)),
+        },
+        languages: Vec::new(),
+        grammars: BTreeMap::default(),
+        language_servers: [("gleam".into(), LanguageServerManifestEntry::default())]
+            .into_iter()
+            .collect(),
+        context_servers: BTreeMap::default(),
+        slash_commands: BTreeMap::default(),
+        indexed_docs_providers: BTreeMap::default(),
+        snippets: None,
+        capabilities: vec![ExtensionCapability::ProcessExec {
+            command: "echo".into(),
+            args: vec!["hello!".into()],
+        }],
+    }
+}
+
+criterion_group!(benches, extension_benchmarks);
+criterion_main!(benches);

crates/rope/Cargo.toml

@@ -27,7 +27,7 @@ env_logger.workspace = true
 gpui = { workspace = true, features = ["test-support"] }
 rand.workspace = true
 util = { workspace = true, features = ["test-support"] }
-criterion = { version = "0.5", features = ["html_reports"] }
+criterion.workspace = true
 
 [[bench]]
 name = "rope_benchmark"

tooling/workspace-hack/Cargo.toml

@@ -132,7 +132,7 @@ url = { version = "2", features = ["serde"] }
 uuid = { version = "1", features = ["serde", "v4", "v5", "v7"] }
 wasm-encoder = { version = "0.221", features = ["wasmparser"] }
 wasmparser = { version = "0.221" }
-wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc"] }
+wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc", "parallel-compilation"] }
 wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc"] }
 wasmtime-environ = { version = "29", default-features = false, features = ["compile", "component-model", "demangle", "gc-drc"] }
 winnow = { version = "0.7", features = ["simd"] }
@@ -267,7 +267,7 @@ url = { version = "2", features = ["serde"] }
 uuid = { version = "1", features = ["serde", "v4", "v5", "v7"] }
 wasm-encoder = { version = "0.221", features = ["wasmparser"] }
 wasmparser = { version = "0.221" }
-wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc"] }
+wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc", "parallel-compilation"] }
 wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc"] }
 wasmtime-environ = { version = "29", default-features = false, features = ["compile", "component-model", "demangle", "gc-drc"] }
 winnow = { version = "0.7", features = ["simd"] }