From 26a8cac0d8d963992242955a1ce1398ec8358354 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Mon, 19 May 2025 18:06:33 +0200 Subject: [PATCH] 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 | 3 + Cargo.toml | 2 + crates/extension/src/extension_builder.rs | 2 +- crates/extension/src/extension_manifest.rs | 2 +- crates/extension_host/Cargo.toml | 6 + .../extension_compilation_benchmark.rs | 145 ++++++++++++++++++ crates/rope/Cargo.toml | 2 +- tooling/workspace-hack/Cargo.toml | 4 +- 8 files changed, 161 insertions(+), 5 deletions(-) create mode 100644 crates/extension_host/benches/extension_compilation_benchmark.rs diff --git a/Cargo.lock b/Cargo.lock index 31fab04d58..504cb2a573 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 65bcec73bf..cf227b8394 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/extension/src/extension_builder.rs b/crates/extension/src/extension_builder.rs index 34bf30363a..73152c667b 100644 --- a/crates/extension/src/extension_builder.rs +++ b/crates/extension/src/extension_builder.rs @@ -398,7 +398,7 @@ impl ExtensionBuilder { async fn install_wasi_sdk_if_needed(&self) -> Result { 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); }; diff --git a/crates/extension/src/extension_manifest.rs b/crates/extension/src/extension_manifest.rs index 799b30861f..6ddaee98ba 100644 --- a/crates/extension/src/extension_manifest.rs +++ b/crates/extension/src/extension_manifest.rs @@ -162,7 +162,7 @@ pub struct GrammarManifestEntry { pub path: Option, } -#[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)] diff --git a/crates/extension_host/Cargo.toml b/crates/extension_host/Cargo.toml index 5ce6e1991f..dbee29f36c 100644 --- a/crates/extension_host/Cargo.toml +++ b/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 diff --git a/crates/extension_host/benches/extension_compilation_benchmark.rs b/crates/extension_host/benches/extension_compilation_benchmark.rs new file mode 100644 index 0000000000..2837557528 --- /dev/null +++ b/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 { + 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 { + 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); diff --git a/crates/rope/Cargo.toml b/crates/rope/Cargo.toml index 5adadfd39d..d9714a23ae 100644 --- a/crates/rope/Cargo.toml +++ b/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" diff --git a/tooling/workspace-hack/Cargo.toml b/tooling/workspace-hack/Cargo.toml index f656d6a00b..ec99a7e681 100644 --- a/tooling/workspace-hack/Cargo.toml +++ b/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"] }