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
 23    let mut group = c.benchmark_group("load");
 24
 25    let mut manifest = manifest();
 26    let wasm_bytes = wasm_bytes(&cx, &mut manifest);
 27    let manifest = Arc::new(manifest);
 28    let extensions_dir = TempTree::new(json!({
 29        "installed": {},
 30        "work": {}
 31    }));
 32    let wasm_host = wasm_host(&cx, &extensions_dir);
 33
 34    group.bench_function(BenchmarkId::from_parameter(1), |b| {
 35        b.iter_batched(
 36            || wasm_bytes.clone(),
 37            |wasm_bytes| {
 38                let _extension = cx
 39                    .executor()
 40                    .block(wasm_host.load_extension(wasm_bytes, &manifest, cx.executor()))
 41                    .unwrap();
 42            },
 43            BatchSize::SmallInput,
 44        );
 45    });
 46}
 47
 48fn init() -> TestAppContext {
 49    const SEED: u64 = 9999;
 50    let dispatcher = TestDispatcher::new(StdRng::seed_from_u64(SEED));
 51    let cx = TestAppContext::build(dispatcher, 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.executor()
 71        .block(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);