extension_host: Use wasmtime incremental compilation (#30948)

Ben Brandt created 

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

Change summary

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(-)

Detailed changes

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",

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"

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

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<wasmtime::Engine> = 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<wasmtime::Engine> = 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<IncrementalCompilationCache> {
+    static CACHE_STORE: LazyLock<Arc<IncrementalCompilationCache>> =
+        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<u8>, Vec<u8>>,
+}
+
+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<u8>, v: &Vec<u8>| (k.len() + v.len()).try_into().unwrap_or(u32::MAX))
+            .build();
+        Self { cache }
+    }
+}
+
+impl CacheStore for IncrementalCompilationCache {
+    fn get(&self, key: &[u8]) -> Option<Cow<[u8]>> {
+        self.cache.get(key).map(|v| v.into())
+    }
+
+    fn insert(&self, key: &[u8], value: Vec<u8>) -> 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<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 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()],
+            }],
+        }
+    }
+}

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"] }