# HG changeset patch
# User Raphaël Gomès <rgomes@octobus.net>
# Date 1679595015 -3600
# Node ID 14b57943ae6df0ede484dfa3ea5af6957cae8a5f
# Parent  9c5e743e400cf284375ce044b5de7d9e7cdf2ce7
rust: fix thread cap (for real this time)

Both e2f8ed37201c and c52435820bbd failed to put a *default* ceiling on
the number of threads used by Rayon to prevent a contention issue.

Calling `rayon::available_parallelism()` creates the global threadpool,
which made our whole dance useless last time.

diff -r 9c5e743e400c -r 14b57943ae6d rust/hg-core/src/dirstate_tree/status.rs
--- a/rust/hg-core/src/dirstate_tree/status.rs	Wed Mar 22 17:18:32 2023 +0000
+++ b/rust/hg-core/src/dirstate_tree/status.rs	Thu Mar 23 19:10:15 2023 +0100
@@ -47,16 +47,10 @@
     options: StatusOptions,
 ) -> Result<(DirstateStatus<'dirstate>, Vec<PatternFileWarning>), StatusError>
 {
-    // Force the global rayon threadpool to not exceed 16 concurrent threads.
-    // This is a stop-gap measure until we figure out why using more than 16
-    // threads makes `status` slower for each additional thread.
-    // We use `ok()` in case the global threadpool has already been
-    // instantiated in `rhg` or some other caller.
-    // TODO find the underlying cause and fix it, then remove this.
-    rayon::ThreadPoolBuilder::new()
-        .num_threads(16.min(rayon::current_num_threads()))
-        .build_global()
-        .ok();
+    // Also cap for a Python caller of this function, but don't complain if
+    // the global threadpool has already been set since this code path is also
+    // being used by `rhg`, which calls this early.
+    let _ = crate::utils::cap_default_rayon_threads();
 
     let (ignore_fn, warnings, patterns_changed): (IgnoreFnType, _, _) =
         if options.list_ignored || options.list_unknown {
diff -r 9c5e743e400c -r 14b57943ae6d rust/hg-core/src/utils.rs
--- a/rust/hg-core/src/utils.rs	Wed Mar 22 17:18:32 2023 +0000
+++ b/rust/hg-core/src/utils.rs	Thu Mar 23 19:10:15 2023 +0100
@@ -498,3 +498,35 @@
         Err(e) => Some(Err(e)),
     })
 }
+
+/// Force the global rayon threadpool to not exceed 16 concurrent threads
+/// unless the user has specified a value.
+/// This is a stop-gap measure until we figure out why using more than 16
+/// threads makes `status` slower for each additional thread.
+///
+/// TODO find the underlying cause and fix it, then remove this.
+///
+/// # Errors
+///
+/// Returns an error if the global threadpool has already been initialized if
+/// we try to initialize it.
+pub fn cap_default_rayon_threads() -> Result<(), rayon::ThreadPoolBuildError> {
+    const THREAD_CAP: usize = 16;
+
+    if std::env::var("RAYON_NUM_THREADS").is_err() {
+        let available_parallelism = std::thread::available_parallelism()
+            .map(usize::from)
+            .unwrap_or(1);
+        let new_thread_count = THREAD_CAP.min(available_parallelism);
+        let res = rayon::ThreadPoolBuilder::new()
+            .num_threads(new_thread_count)
+            .build_global();
+        if res.is_ok() {
+            log::trace!(
+                "Capped the rayon threadpool to {new_thread_count} threads",
+            );
+        }
+        return res;
+    }
+    Ok(())
+}
diff -r 9c5e743e400c -r 14b57943ae6d rust/rhg/src/main.rs
--- a/rust/rhg/src/main.rs	Wed Mar 22 17:18:32 2023 +0000
+++ b/rust/rhg/src/main.rs	Thu Mar 23 19:10:15 2023 +0100
@@ -140,6 +140,13 @@
 
     env_logger::init();
 
+    // Make sure nothing in a future version of `rhg` sets the global
+    // threadpool before we can cap default threads. (This is also called
+    // in core because Python uses the same code path, we're adding a
+    // redundant check.)
+    hg::utils::cap_default_rayon_threads()
+        .expect("Rayon threadpool already initialized");
+
     let early_args = EarlyArgs::parse(&argv);
 
     let initial_current_dir = early_args.cwd.map(|cwd| {