From a6ad8fb342cc4dd3647821ccfef860cb89a651de Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Sat, 21 Jun 2025 17:44:30 +0200
Subject: [PATCH] [ty] Add multithreaded benchmark (#18822)

---
 crates/ruff_benchmark/benches/ty_walltime.rs | 45 +++++++++++---------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/crates/ruff_benchmark/benches/ty_walltime.rs b/crates/ruff_benchmark/benches/ty_walltime.rs
index 39029f4505..cb5026f01b 100644
--- a/crates/ruff_benchmark/benches/ty_walltime.rs
+++ b/crates/ruff_benchmark/benches/ty_walltime.rs
@@ -204,8 +204,8 @@ static SYMPY: std::sync::LazyLock<Benchmark<'static>> = std::sync::LazyLock::new
     )
 });
 
-#[bench(args=[&*ALTAIR, &*FREQTRADE, &*PYDANTIC], sample_size=2, sample_count=3)]
-fn small(bencher: Bencher, benchmark: &Benchmark) {
+#[track_caller]
+fn run_single_threaded(bencher: Bencher, benchmark: &Benchmark) {
     bencher
         .with_inputs(|| benchmark.setup_iteration())
         .bench_local_refs(|db| {
@@ -213,41 +213,46 @@ fn small(bencher: Bencher, benchmark: &Benchmark) {
         });
 }
 
+#[bench(args=[&*ALTAIR, &*FREQTRADE, &*PYDANTIC], sample_size=2, sample_count=3)]
+fn small(bencher: Bencher, benchmark: &Benchmark) {
+    run_single_threaded(bencher, benchmark);
+}
+
 #[bench(args=[&*COLOUR_SCIENCE, &*PANDAS], sample_size=1, sample_count=3)]
 fn medium(bencher: Bencher, benchmark: &Benchmark) {
-    bencher
-        .with_inputs(|| benchmark.setup_iteration())
-        .bench_local_refs(|db| {
-            check_project(db, benchmark.max_diagnostics);
-        });
+    run_single_threaded(bencher, benchmark);
 }
 
 #[bench(args=[&*SYMPY], sample_size=1, sample_count=2)]
 fn large(bencher: Bencher, benchmark: &Benchmark) {
+    run_single_threaded(bencher, benchmark);
+}
+
+#[bench(args=[&*PYDANTIC], sample_size=3, sample_count=3)]
+fn multithreaded(bencher: Bencher, benchmark: &Benchmark) {
+    let thread_pool = ThreadPoolBuilder::new().build().unwrap();
+
     bencher
         .with_inputs(|| benchmark.setup_iteration())
-        .bench_local_refs(|db| {
-            check_project(db, benchmark.max_diagnostics);
+        .bench_local_values(|db| {
+            thread_pool.install(|| {
+                check_project(&db, benchmark.max_diagnostics);
+                db
+            })
         });
 }
 
 fn main() {
-    let filter =
-        std::env::var("TY_LOG").unwrap_or("ty_walltime=info,ruff_benchmark=info".to_string());
-
-    let _logging = setup_logging_with_filter(&filter).expect("Filter to be valid");
-
-    // Disable multithreading for now due to
-    // https://github.com/salsa-rs/salsa/issues/918.
-    //
-    // Salsa has a fast-path for the first db when looking up ingredients.
-    // It seems that this fast-path becomes extremely slow for all db's other
-    // than the first one, especially when using multithreading (10x slower than the first run).
     ThreadPoolBuilder::new()
         .num_threads(1)
         .use_current_thread()
         .build_global()
         .unwrap();
 
+    let filter =
+        std::env::var("TY_LOG").unwrap_or("ty_walltime=info,ruff_benchmark=info".to_string());
+
+    let _logging = setup_logging_with_filter(&filter).expect("Filter to be valid");
+
     divan::main();
 }