next.js/turbopack/crates/turbopack-tracing/benches/node_file_trace.rs

use std::{fs, path::PathBuf};

use criterion::{Bencher, BenchmarkId, Criterion};
use regex::Regex;
use turbo_rcstr::{RcStr, rcstr};
use turbo_tasks::{
    Effects, OperationVc, TurboTasks, Vc, mark_top_level_task,
    read_strongly_consistent_and_apply_effects, take_effects,
    unmark_top_level_task_may_leak_eventually_consistent_state,
};
use turbo_tasks_backend::{BackendOptions, TurboTasksBackend, noop_backing_storage};
use turbo_tasks_fs::{DiskFileSystem, FileSystem, NullFileSystem};
use turbopack::{
    ModuleAssetContext, emit_assets_into_dir_operation,
    module_options::{EcmascriptOptionsContext, ModuleOptionsContext},
};
use turbopack_core::{
    compile_time_info::CompileTimeInfo,
    context::AssetContext,
    environment::{Environment, ExecutionEnvironment, NodeJsEnvironment},
    file_source::FileSource,
    ident::Layer,
    rebase::RebasedAsset,
    reference::all_assets_from_entry,
    reference_type::ReferenceType,
};
use turbopack_resolve::resolve_options_context::ResolveOptionsContext;

#[turbo_tasks::function(operation, root)]
async fn extract_effects_operation(op: OperationVc<()>) -> anyhow::Result<Vc<Effects>> {
    let _ = op.resolve().strongly_consistent().await?;
    Ok(take_effects(op).await?.cell())
}

// TODO this should move to the `node-file-trace` crate
pub fn benchmark(c: &mut Criterion) {
    let bench_filter = Regex::new(r"(empty|simple|dynamic-in-package|react|whatwg-url|axios|azure-cosmos|cowsay|env-var|fast-glob)\.js$").unwrap();

    let tests_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests");
    let tests_dir = tests_root.join("node-file-trace/integration");

    let mut group = c.benchmark_group("node-file-trace");
    group.sample_size(10);

    let results = fs::read_dir(tests_dir).unwrap();
    for result in results {
        let entry = result.unwrap();
        if entry.file_type().unwrap().is_file() {
            let name = entry.file_name().into_string().unwrap();
            if !bench_filter.is_match(&name) {
                continue;
            }

            let input = format!("node-file-trace/integration/{name}");
            let tests_root = tests_root.to_string_lossy().to_string();

            let bench_input = BenchInput { tests_root, input };

            group.bench_with_input(
                BenchmarkId::new("emit", &bench_input.input),
                &bench_input,
                bench_emit,
            );
        }
    }

    group.finish();
}

struct BenchInput {
    tests_root: String,
    input: String,
}

fn bench_emit(b: &mut Bencher, bench_input: &BenchInput) {
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    b.to_async(rt).iter(move || {
        let tt = TurboTasks::new(TurboTasksBackend::new(
            BackendOptions::default(),
            noop_backing_storage(),
        ));
        let tests_root: RcStr = bench_input.tests_root.clone().into();
        let input: RcStr = bench_input.input.clone().into();
        async move {
            tt.run_once(async move {
                // This benchmark performs eventually-consistent Vc reads at the top level, which
                // is fine for a throughput benchmark but trips the top-level-task assertion under
                // debug-assertions otherwise. Unmark the top-level task for the reads, then re-mark
                // it before `Effects::apply`, which *requires* a top-level task.
                unmark_top_level_task_may_leak_eventually_consistent_state();
                let input_fs = DiskFileSystem::new(rcstr!("tests"), Vc::cell(tests_root.clone()));
                let input = input_fs.root().await?.join(&input)?;

                let input_dir = input.parent().parent();
                let output_fs: Vc<NullFileSystem> = NullFileSystem.cell();
                let output_dir = output_fs.root().owned().await?;

                let source = FileSource::new(input);
                let compile_time_info = CompileTimeInfo::builder(
                    Environment::new(ExecutionEnvironment::NodeJsLambda(
                        NodeJsEnvironment::default().resolved_cell(),
                    ))
                    .to_resolved()
                    .await?,
                )
                .cell()
                .await?;
                let module_asset_context = ModuleAssetContext::new(
                    Default::default(),
                    compile_time_info,
                    ModuleOptionsContext {
                        ecmascript: EcmascriptOptionsContext {
                            enable_types: true,
                            ..Default::default()
                        },
                        ..Default::default()
                    }
                    .cell(),
                    ResolveOptionsContext {
                        emulate_environment: Some(
                            compile_time_info.environment().to_resolved().await?,
                        ),
                        collect_affecting_sources: true,
                        ..Default::default()
                    }
                    .cell(),
                    Layer::new(rcstr!("node_file_trace")),
                );
                let module = module_asset_context
                    .process(Vc::upcast(source), ReferenceType::Undefined)
                    .module();
                let rebased = RebasedAsset::new(module, input_dir, output_dir.clone());
                let assets = all_assets_from_entry(Vc::upcast(rebased))
                    .to_resolved()
                    .await?;

                // `read_strongly_consistent_and_apply_effects` applies effects, which *requires* a
                // top-level task. Restore the mark we cleared above for the eventually-consistent
                // reads.
                mark_top_level_task();
                read_strongly_consistent_and_apply_effects(
                    extract_effects_operation(emit_assets_into_dir_operation(assets, output_dir)),
                    |e| e,
                )
                .await?;

                anyhow::Ok(())
            })
            .await
            .unwrap();
        }
    })
}