1use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
2
3use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
4use extension::{
5 ExtensionCapability, ExtensionHostProxy, ExtensionLibraryKind, ExtensionManifest,
6 LanguageServerManifestEntry, LibManifestEntry, SchemaVersion,
7 extension_builder::{CompileExtensionOptions, ExtensionBuilder},
8};
9use extension_host::wasm_host::WasmHost;
10use fs::{Fs, RealFs};
11use gpui::{TestAppContext, TestDispatcher};
12use http_client::{FakeHttpClient, Response};
13use node_runtime::NodeRuntime;
14use rand::{SeedableRng, rngs::StdRng};
15use reqwest_client::ReqwestClient;
16use serde_json::json;
17use settings::SettingsStore;
18use util::test::TempTree;
19
20fn extension_benchmarks(c: &mut Criterion) {
21 let cx = init();
22 cx.update(gpui_tokio::init);
23
24 let mut group = c.benchmark_group("load");
25
26 let mut manifest = manifest();
27 let wasm_bytes = wasm_bytes(
28 &cx,
29 &mut manifest,
30 Arc::new(RealFs::new(None, cx.executor())),
31 );
32 let manifest = Arc::new(manifest);
33 let extensions_dir = TempTree::new(json!({
34 "installed": {},
35 "work": {}
36 }));
37 let wasm_host = wasm_host(&cx, &extensions_dir);
38
39 group.bench_function(BenchmarkId::from_parameter(1), |b| {
40 b.iter_batched(
41 || wasm_bytes.clone(),
42 |wasm_bytes| {
43 let _extension = cx
44 .executor()
45 .block(wasm_host.load_extension(wasm_bytes, &manifest, &cx.to_async()))
46 .unwrap();
47 },
48 BatchSize::SmallInput,
49 );
50 });
51}
52
53fn init() -> TestAppContext {
54 const SEED: u64 = 9999;
55 let dispatcher = TestDispatcher::new(StdRng::seed_from_u64(SEED));
56 let cx = TestAppContext::build(dispatcher, None);
57 cx.executor().allow_parking();
58 cx.update(|cx| {
59 let store = SettingsStore::test(cx);
60 cx.set_global(store);
61 release_channel::init(semver::Version::new(0, 0, 0), cx);
62 });
63
64 cx
65}
66
67fn wasm_bytes(cx: &TestAppContext, manifest: &mut ExtensionManifest, fs: Arc<dyn Fs>) -> Vec<u8> {
68 let extension_builder = extension_builder();
69 let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
70 .parent()
71 .unwrap()
72 .parent()
73 .unwrap()
74 .join("extensions/test-extension");
75 cx.executor()
76 .block(extension_builder.compile_extension(
77 &path,
78 manifest,
79 CompileExtensionOptions { release: true },
80 fs,
81 ))
82 .unwrap();
83 std::fs::read(path.join("extension.wasm")).unwrap()
84}
85
86fn extension_builder() -> ExtensionBuilder {
87 let user_agent = format!(
88 "Zed Extension CLI/{} ({}; {})",
89 env!("CARGO_PKG_VERSION"),
90 std::env::consts::OS,
91 std::env::consts::ARCH
92 );
93 let http_client = Arc::new(ReqwestClient::user_agent(&user_agent).unwrap());
94 // Local dir so that we don't have to download it on every run
95 let build_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("benches/.build");
96 ExtensionBuilder::new(http_client, build_dir)
97}
98
99fn wasm_host(cx: &TestAppContext, extensions_dir: &TempTree) -> Arc<WasmHost> {
100 let http_client = FakeHttpClient::create(async |_| {
101 Ok(Response::builder().status(404).body("not found".into())?)
102 });
103 let extensions_dir = extensions_dir.path().canonicalize().unwrap();
104 let work_dir = extensions_dir.join("work");
105 let fs = Arc::new(RealFs::new(None, cx.executor()));
106
107 cx.update(|cx| {
108 WasmHost::new(
109 fs,
110 http_client,
111 NodeRuntime::unavailable(),
112 Arc::new(ExtensionHostProxy::new()),
113 work_dir,
114 cx,
115 )
116 })
117}
118
119fn manifest() -> ExtensionManifest {
120 ExtensionManifest {
121 id: "test-extension".into(),
122 name: "Test Extension".into(),
123 version: "0.1.0".into(),
124 schema_version: SchemaVersion(1),
125 description: Some("An extension for use in tests.".into()),
126 authors: Vec::new(),
127 repository: None,
128 themes: Default::default(),
129 icon_themes: Vec::new(),
130 lib: LibManifestEntry {
131 kind: Some(ExtensionLibraryKind::Rust),
132 version: Some(semver::Version::new(0, 1, 0)),
133 },
134 languages: Vec::new(),
135 grammars: BTreeMap::default(),
136 language_servers: [("gleam".into(), LanguageServerManifestEntry::default())]
137 .into_iter()
138 .collect(),
139 context_servers: BTreeMap::default(),
140 agent_servers: BTreeMap::default(),
141 slash_commands: BTreeMap::default(),
142 snippets: None,
143 capabilities: vec![ExtensionCapability::ProcessExec(
144 extension::ProcessExecCapability {
145 command: "echo".into(),
146 args: vec!["hello!".into()],
147 },
148 )],
149 debug_adapters: Default::default(),
150 debug_locators: Default::default(),
151 }
152}
153
154criterion_group!(benches, extension_benchmarks);
155criterion_main!(benches);