From 4ece4a635fb7f525aaafc5fd4f3926f04e4172a8 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Wed, 21 May 2025 10:12:16 +0200 Subject: [PATCH] extension_host: Use wasmtime incremental compilation (#30948) Builds on top of https://github.com/zed-industries/zed/pull/30942 This turns on incremental compilation and decreases extension compilation times by up to another 41% Putting us at roughly 92% improved extension load times from what is in the app today. Because we only have a static engine, I can't reset the cache between every run. So technically the benchmarks are always running with a warmed cache. So the first extension we load will take the 8.8ms, and then any subsequent extensions will be closer to the measured time in this benchmark. This is also measuring the entire load process, not just the compilation. However, since this is the loading we likely think of when thinking about extensions, I felt it was likely more helpful to see the impact on the overall time. This works because our extensions are largely the same Wasm bytecode (SDK code + std lib functions etc) with minor changes in the trait impl. The more different that extensions implementation is, there will be less benefit, however, there will always be a large part of every extension that is always the same across extensions, so this should be a speedup regardless. I used `moka` to provide a bound to the cache. We could use a bare `DashMap`, however if there was some issue this could lead to a memory leak. `moka` has some slight overhead, but makes sure that we don't go over 32mb while using an LRU-style mechanism for deciding which compilation artifacts to keep. I measured our current extensions to take roughly 512kb in the cache. Which means with a cap of 32mb, we can keep roughly 64 *completely novel* extensions with no overlap. Since our extensions will have more overlap than this though, we can actually keep much more in the cache without having to worry about it. #### Before: ``` load/1 time: [8.8301 ms 8.8616 ms 8.8931 ms] change: [-0.1880% +0.3221% +0.8679%] (p = 0.23 > 0.05) No change in performance detected. ``` #### After: ``` load/1 time: [5.1575 ms 5.1726 ms 5.1876 ms] change: [-41.894% -41.628% -41.350%] (p = 0.00 < 0.05) Performance has improved. ``` Release Notes: - N/A --- Cargo.lock | 60 +++++++++ Cargo.toml | 2 + crates/extension_host/Cargo.toml | 1 + crates/extension_host/src/wasm_host.rs | 163 +++++++++++++++++++++++-- tooling/workspace-hack/Cargo.toml | 14 ++- 5 files changed, 224 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4a0d3189443a9319da5ed5ad7e712a1b78467ed7..bbf67873670a18511ed5d3f1cb3f03f14ca835a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3636,9 +3636,12 @@ dependencies = [ "gimli", "hashbrown 0.14.5", "log", + "postcard", "regalloc2", "rustc-hash 2.1.1", "serde", + "serde_derive", + "sha2", "smallvec", "target-lexicon 0.13.2", ] @@ -5154,6 +5157,7 @@ dependencies = [ "language_extension", "log", "lsp", + "moka", "node_runtime", "parking_lot", "paths", @@ -5910,6 +5914,20 @@ dependencies = [ "thread_local", ] +[[package]] +name = "generator" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d18470a76cb7f8ff746cf1f7470914f900252ec36bbc40b569d74b1258446827" +dependencies = [ + "cc", + "cfg-if", + "libc", + "log", + "rustversion", + "windows 0.61.1", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -9348,6 +9366,19 @@ dependencies = [ "logos-codegen", ] +[[package]] +name = "loom" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" +dependencies = [ + "cfg-if", + "generator", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + [[package]] name = "loop9" version = "0.1.5" @@ -9846,6 +9877,25 @@ dependencies = [ "workspace-hack", ] +[[package]] +name = "moka" +version = "0.12.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9321642ca94a4282428e6ea4af8cc2ca4eac48ac7a6a4ea8f33f76d0ce70926" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "loom", + "parking_lot", + "portable-atomic", + "rustc_version", + "smallvec", + "tagptr", + "thiserror 1.0.69", + "uuid", +] + [[package]] name = "msvc_spectre_libs" version = "0.1.3" @@ -12782,6 +12832,7 @@ dependencies = [ "hashbrown 0.15.3", "log", "rustc-hash 2.1.1", + "serde", "smallvec", ] @@ -15433,6 +15484,12 @@ dependencies = [ "slotmap", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "take-until" version = "0.2.0" @@ -19128,7 +19185,9 @@ dependencies = [ "core-foundation 0.9.4", "core-foundation-sys", "coreaudio-sys", + "cranelift-codegen", "crc32fast", + "crossbeam-epoch", "crossbeam-utils", "crypto-common", "deranged", @@ -19204,6 +19263,7 @@ dependencies = [ "rand 0.9.1", "rand_chacha 0.3.1", "rand_core 0.6.4", + "regalloc2", "regex", "regex-automata 0.4.9", "regex-syntax 0.8.5", diff --git a/Cargo.toml b/Cargo.toml index cf227b83942ca10db28b4db8f77e71eb529993c9..5b70db35fdc0db6bfcebe57f0f242f1d7ff8b1fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -476,6 +476,7 @@ lsp-types = { git = "https://github.com/zed-industries/lsp-types", rev = "c9c189 markup5ever_rcdom = "0.3.0" metal = "0.29" mlua = { version = "0.10", features = ["lua54", "vendored", "async", "send"] } +moka = { version = "0.12.10", features = ["sync"] } naga = { version = "25.0", features = ["wgsl-in"] } nanoid = "0.4" nbformat = { git = "https://github.com/ConradIrwin/runtimed", rev = "7130c804216b6914355d15d0b91ea91f6babd734" } @@ -609,6 +610,7 @@ wasmtime = { version = "29", default-features = false, features = [ "runtime", "cranelift", "component-model", + "incremental-cache", "parallel-compilation", ] } wasmtime-wasi = "29" diff --git a/crates/extension_host/Cargo.toml b/crates/extension_host/Cargo.toml index dbee29f36cbe1ad0d7bd059260ece0a7959e114b..68cbd6a4a3def3d8dbd2482231ee2445e976a11e 100644 --- a/crates/extension_host/Cargo.toml +++ b/crates/extension_host/Cargo.toml @@ -31,6 +31,7 @@ http_client.workspace = true language.workspace = true log.workspace = true lsp.workspace = true +moka.workspace = true node_runtime.workspace = true paths.workspace = true project.workspace = true diff --git a/crates/extension_host/src/wasm_host.rs b/crates/extension_host/src/wasm_host.rs index 26d0a073e71adf05876328c15844da4c9f5b22ec..1aafd15092f89276c235f7dc834570c2f20c05d4 100644 --- a/crates/extension_host/src/wasm_host.rs +++ b/crates/extension_host/src/wasm_host.rs @@ -22,15 +22,18 @@ use gpui::{App, AsyncApp, BackgroundExecutor, Task}; use http_client::HttpClient; use language::LanguageName; use lsp::LanguageServerName; +use moka::sync::Cache; use node_runtime::NodeRuntime; use release_channel::ReleaseChannel; use semantic_version::SemanticVersion; +use std::borrow::Cow; +use std::sync::LazyLock; use std::{ path::{Path, PathBuf}, - sync::{Arc, OnceLock}, + sync::Arc, }; use wasmtime::{ - Engine, Store, + CacheStore, Engine, Store, component::{Component, ResourceTable}, }; use wasmtime_wasi::{self as wasi, WasiView}; @@ -411,16 +414,23 @@ type ExtensionCall = Box< >; fn wasm_engine() -> wasmtime::Engine { - static WASM_ENGINE: OnceLock = OnceLock::new(); - - WASM_ENGINE - .get_or_init(|| { - let mut config = wasmtime::Config::new(); - config.wasm_component_model(true); - config.async_support(true); - wasmtime::Engine::new(&config).unwrap() - }) - .clone() + static WASM_ENGINE: LazyLock = LazyLock::new(|| { + let mut config = wasmtime::Config::new(); + config.wasm_component_model(true); + config.async_support(true); + config + .enable_incremental_compilation(cache_store()) + .unwrap(); + wasmtime::Engine::new(&config).unwrap() + }); + + WASM_ENGINE.clone() +} + +fn cache_store() -> Arc { + static CACHE_STORE: LazyLock> = + LazyLock::new(|| Arc::new(IncrementalCompilationCache::new())); + CACHE_STORE.clone() } impl WasmHost { @@ -667,3 +677,132 @@ impl wasi::WasiView for WasmState { &mut self.ctx } } + +/// Wrapper around a mini-moka bounded cache for storing incremental compilation artifacts. +/// Since wasm modules have many similar elements, this can save us a lot of work at the +/// cost of a small memory footprint. However, we don't want this to be unbounded, so we use +/// a LFU/LRU cache to evict less used cache entries. +#[derive(Debug)] +struct IncrementalCompilationCache { + cache: Cache, Vec>, +} + +impl IncrementalCompilationCache { + fn new() -> Self { + let cache = Cache::builder() + // Cap this at 32 MB for now. Our extensions turn into roughly 512kb in the cache, + // which means we could store 64 completely novel extensions in the cache, but in + // practice we will more than that, which is more than enough for our use case. + .max_capacity(32 * 1024 * 1024) + .weigher(|k: &Vec, v: &Vec| (k.len() + v.len()).try_into().unwrap_or(u32::MAX)) + .build(); + Self { cache } + } +} + +impl CacheStore for IncrementalCompilationCache { + fn get(&self, key: &[u8]) -> Option> { + self.cache.get(key).map(|v| v.into()) + } + + fn insert(&self, key: &[u8], value: Vec) -> bool { + self.cache.insert(key.to_vec(), value); + true + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use extension::{ + ExtensionCapability, ExtensionLibraryKind, LanguageServerManifestEntry, LibManifestEntry, + SchemaVersion, + extension_builder::{CompileExtensionOptions, ExtensionBuilder}, + }; + use gpui::TestAppContext; + use reqwest_client::ReqwestClient; + + use super::*; + + #[gpui::test] + fn test_cache_size_for_test_extension(cx: &TestAppContext) { + let cache_store = cache_store(); + let engine = wasm_engine(); + let wasm_bytes = wasm_bytes(cx, &mut manifest()); + + Component::new(&engine, wasm_bytes).unwrap(); + + cache_store.cache.run_pending_tasks(); + let size: usize = cache_store + .cache + .iter() + .map(|(k, v)| k.len() + v.len()) + .sum(); + // If this assertion fails, it means extensions got larger and we may want to + // reconsider our cache size. + assert!(size < 512 * 1024); + } + + 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 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()], + }], + } + } +} diff --git a/tooling/workspace-hack/Cargo.toml b/tooling/workspace-hack/Cargo.toml index ec99a7e681823fea09e46108c13f2b3a7811d1fb..273ffa72f6d9f5459277d96a48a98e979d629a65 100644 --- a/tooling/workspace-hack/Cargo.toml +++ b/tooling/workspace-hack/Cargo.toml @@ -44,7 +44,9 @@ chrono = { version = "0.4", features = ["serde"] } clap = { version = "4", features = ["cargo", "derive", "string", "wrap_help"] } clap_builder = { version = "4", default-features = false, features = ["cargo", "color", "std", "string", "suggestions", "usage", "wrap_help"] } concurrent-queue = { version = "2" } +cranelift-codegen = { version = "0.116", default-features = false, features = ["host-arch", "incremental-cache", "std", "timing", "unwind"] } crc32fast = { version = "1" } +crossbeam-epoch = { version = "0.9" } crossbeam-utils = { version = "0.8" } deranged = { version = "0.4", default-features = false, features = ["powerfmt", "serde", "std"] } digest = { version = "0.10", features = ["mac", "oid", "std"] } @@ -95,6 +97,7 @@ prost-types = { version = "0.9" } rand-c38e5c1d305a1b54 = { package = "rand", version = "0.8", features = ["small_rng"] } rand_chacha = { version = "0.3" } rand_core = { version = "0.6", default-features = false, features = ["std"] } +regalloc2 = { version = "0.11", features = ["checker", "enable-serde"] } regex = { version = "1" } regex-automata = { version = "0.4" } regex-syntax = { version = "0.8" } @@ -132,8 +135,8 @@ 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", "parallel-compilation"] } -wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc"] } +wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc", "incremental-cache", "parallel-compilation"] } +wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc", "incremental-cache"] } wasmtime-environ = { version = "29", default-features = false, features = ["compile", "component-model", "demangle", "gc-drc"] } winnow = { version = "0.7", features = ["simd"] } @@ -168,7 +171,9 @@ chrono = { version = "0.4", features = ["serde"] } clap = { version = "4", features = ["cargo", "derive", "string", "wrap_help"] } clap_builder = { version = "4", default-features = false, features = ["cargo", "color", "std", "string", "suggestions", "usage", "wrap_help"] } concurrent-queue = { version = "2" } +cranelift-codegen = { version = "0.116", default-features = false, features = ["host-arch", "incremental-cache", "std", "timing", "unwind"] } crc32fast = { version = "1" } +crossbeam-epoch = { version = "0.9" } crossbeam-utils = { version = "0.8" } deranged = { version = "0.4", default-features = false, features = ["powerfmt", "serde", "std"] } digest = { version = "0.10", features = ["mac", "oid", "std"] } @@ -224,6 +229,7 @@ quote = { version = "1" } rand-c38e5c1d305a1b54 = { package = "rand", version = "0.8", features = ["small_rng"] } rand_chacha = { version = "0.3" } rand_core = { version = "0.6", default-features = false, features = ["std"] } +regalloc2 = { version = "0.11", features = ["checker", "enable-serde"] } regex = { version = "1" } regex-automata = { version = "0.4" } regex-syntax = { version = "0.8" } @@ -267,8 +273,8 @@ 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", "parallel-compilation"] } -wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc"] } +wasmtime = { version = "29", default-features = false, features = ["async", "component-model", "cranelift", "demangle", "gc-drc", "incremental-cache", "parallel-compilation"] } +wasmtime-cranelift = { version = "29", default-features = false, features = ["component-model", "gc-drc", "incremental-cache"] } wasmtime-environ = { version = "29", default-features = false, features = ["compile", "component-model", "demangle", "gc-drc"] } winnow = { version = "0.7", features = ["simd"] }