Initialize rayon outside of benchmark

astral-sh · Aug 23, 2024 · 5401274 · 5401274
1 parent 39fb0a2
commit 5401274
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/ruff_benchmark/Cargo.toml b/crates/ruff_benchmark/Cargo.toml
@@ -37,13 +37,14 @@ name = "red_knot"
 harness = false
 
 [dependencies]
+codspeed-criterion-compat = { workspace = true, default-features = false, optional = true }
+criterion = { workspace = true, default-features = false }
 once_cell = { workspace = true }
+rayon = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 url = { workspace = true }
 ureq = { workspace = true }
-criterion = { workspace = true, default-features = false }
-codspeed-criterion-compat = { workspace = true, default-features = false, optional = true }
 
 [dev-dependencies]
 ruff_db = { workspace = true }

diff --git a/crates/ruff_benchmark/benches/red_knot.rs b/crates/ruff_benchmark/benches/red_knot.rs
@@ -1,5 +1,6 @@
 #![allow(clippy::disallowed_names)]
 
+use rayon::ThreadPoolBuilder;
 use red_knot_python_semantic::PythonVersion;
 use red_knot_workspace::db::RootDatabase;
 use red_knot_workspace::watch::{ChangeEvent, ChangedKind};
@@ -112,7 +113,20 @@ fn setup_case() -> Case {
     }
 }
 
+static RAYON_INITIALIZED: std::sync::Once = std::sync::Once::new();
+
+fn setup_rayon() {
+    RAYON_INITIALIZED.call_once(|| {
+        ThreadPoolBuilder::new()
+            .num_threads(1)
+            .build_global()
+            .unwrap()
+    })
+}
+
 fn benchmark_incremental(criterion: &mut Criterion) {
+    setup_rayon();
+
     criterion.bench_function("red_knot_check_file[incremental]", |b| {
         b.iter_batched_ref(
             || {
@@ -149,6 +163,8 @@ fn benchmark_incremental(criterion: &mut Criterion) {
 }
 
 fn benchmark_cold(criterion: &mut Criterion) {
+    setup_rayon();
+
     criterion.bench_function("red_knot_check_file[cold]", |b| {
         b.iter_batched_ref(
             setup_case,