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