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