Merge PR #3774: fix(natives): cap rayon pool on Windows (@roboomp)

crates/pi-natives/src/lib.rs

@@ -76,6 +76,15 @@ const NAPI_TOKIO_MAX_WORKER_THREADS: usize = 4;
 #[cfg(target_os = "windows")]
 const NAPI_TOKIO_MAX_BLOCKING_THREADS: usize = 8;
 
+/// Upper bound on Rayon's global pool on Windows. Rayon is used for CPU-bound
+/// helpers (`count_tokens`, vendored `sort`), but the global pool cannot
+/// recover if its first lazy initialization fails after Windows refuses worker
+/// threads.
+#[cfg(any(target_os = "windows", test))]
+const RAYON_MAX_THREADS: usize = 8;
+#[cfg(any(target_os = "windows", test))]
+const RAYON_RESERVED_NON_RAYON_THREADS: usize = 1;
+
 /// Windows worker count we'd *like*, before checking what the OS will actually
 /// grant: the Tokio default (one per core) clamped to
 /// [`NAPI_TOKIO_MAX_WORKER_THREADS`].
@@ -86,6 +95,36 @@ fn desired_worker_threads() -> usize {
 		.clamp(1, NAPI_TOKIO_MAX_WORKER_THREADS)
 }
 
+#[cfg(target_os = "windows")]
+fn desired_rayon_threads() -> usize {
+	clamped_rayon_threads(std::thread::available_parallelism().map_or(1, |threads| threads.get()))
+}
+
+#[cfg(any(target_os = "windows", test))]
+fn clamped_rayon_threads(threads: usize) -> usize {
+	threads.clamp(1, RAYON_MAX_THREADS)
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[cfg(any(target_os = "windows", test))]
+enum RayonPoolPlan {
+	WorkerThreads(usize),
+	SkipGlobalPool,
+}
+
+#[cfg(any(target_os = "windows", test))]
+fn rayon_pool_plan(desired: usize, spawnable: usize) -> RayonPoolPlan {
+	let desired = clamped_rayon_threads(desired);
+	let workers = spawnable
+		.saturating_sub(RAYON_RESERVED_NON_RAYON_THREADS)
+		.min(desired);
+	if workers > 0 {
+		RayonPoolPlan::WorkerThreads(workers)
+	} else {
+		RayonPoolPlan::SkipGlobalPool
+	}
+}
+
 /// Probe how many worker threads Windows will let us hold alive
 /// *simultaneously*, up to `target`. Returns the count actually spawned (0 when
 /// not even one extra thread is possible).
@@ -129,6 +168,35 @@ fn probe_spawnable_workers(target: usize) -> usize {
 	spawned
 }
 
+/// Install Rayon's global pool before any `par_iter` or vendored uutils sort
+/// path can lazily initialize it with Rayon's default one-thread-per-core
+/// policy. When the probe sees fewer workers than requested, reserve capacity
+/// for native code that must still perform its own `thread::spawn` (notably
+/// vendored `sort`'s external-sort helper) and build the global pool with the
+/// remaining spawnable count; when no worker remains after that reserve, leave
+/// the global pool untouched and keep patched Rayon callsites on sequential
+/// paths.
+/// Rayon stores global initialization in a `Once`, so a failed
+/// `build_global()` call would permanently poison the process for later
+/// parallel work.
+#[cfg(target_os = "windows")]
+fn configure_rayon_pool() {
+	let desired = desired_rayon_threads();
+	let plan = rayon_pool_plan(desired, probe_spawnable_workers(desired));
+	let result = match plan {
+		RayonPoolPlan::WorkerThreads(threads) => rayon::ThreadPoolBuilder::new()
+			.num_threads(threads)
+			.build_global(),
+		RayonPoolPlan::SkipGlobalPool => {
+			pi_uutils_ctx::set_rayon_global_pool_available(false);
+			return;
+		},
+	};
+	if result.is_ok() {
+		pi_uutils_ctx::set_rayon_global_pool_available(true);
+	}
+}
+
 /// Build the custom Tokio runtime napi-rs uses on Windows, sized to what the
 /// host can actually spawn. Never panics: backs off from
 /// [`desired_worker_threads`] to whatever the probe allows, and falls back to a
@@ -198,21 +266,29 @@ fn install_native_crash_handler() {
 #[cfg(target_os = "windows")]
 static TOKIO_RUNTIME_INSTALLED: AtomicBool = AtomicBool::new(false);
 
-/// Install the bounded Tokio runtime napi-rs adopts for async exports.
+/// Install the bounded Tokio runtime napi-rs adopts for async exports and the
+/// bounded Rayon global pool used by native parallel iterators.
 ///
 /// The JS loader calls this exactly once, synchronously, right *after* `dlopen`
-/// returns and *before* any async native runs — never from `#[module_init]`.
-/// Building a multi-thread runtime eagerly spawns worker threads, and doing
-/// that during module init (while the dynamic-loader lock is held) deadlocks on
-/// some hosts: a fresh worker blocks acquiring the loader lock that the init
-/// thread still owns. napi-rs only materializes its runtime on the first async
-/// call (`RT` is a `LazyLock`) and `create_custom_tokio_runtime` merely records
-/// the runtime in a `OnceLock`, so installing it post-load is still honored.
-/// Without it napi builds its own default (one worker per CPU, spawned eagerly)
-/// which aborts the process (`os error 1455`) on a memory-constrained Windows
-/// host before any JS error can surface; [`create_windows_napi_tokio_runtime`]
-/// pre-flights the spawn instead. If no runtime can be built we leave napi-rs
-/// to its default. Idempotent.
+/// returns and *before* any async native or parallel iterator runs — never from
+/// `#[module_init]`. Building a multi-thread runtime eagerly spawns worker
+/// threads, and doing that during module init (while the dynamic-loader lock is
+/// held) deadlocks on some hosts: a fresh worker blocks acquiring the loader
+/// lock that the init thread still owns. napi-rs only materializes its runtime
+/// on the first async call (`RT` is a `LazyLock`) and
+/// `create_custom_tokio_runtime` merely records the runtime in a `OnceLock`, so
+/// installing it post-load is still honored.
+///
+/// Without the Tokio override napi builds its own default (one worker per CPU,
+/// spawned eagerly), which aborts the process (`os error 1455`) on a
+/// memory-constrained Windows host before any JS error can surface;
+/// [`create_windows_napi_tokio_runtime`] pre-flights the spawn instead. Rayon
+/// has the same one-thread-per-core lazy default, so [`configure_rayon_pool`]
+/// installs a probed global pool before `count_tokens` or vendored `sort` can
+/// trigger it across a N-API nounwind boundary. If no worker thread is
+/// spawnable, patched Rayon callsites stay sequential rather than registering a
+/// current-thread-only global pool that cannot steal work from later native
+/// calls. Idempotent.
 #[napi(js_name = "__ompInstallTokioRuntime")]
 #[allow(clippy::missing_const_for_fn, reason = "napi macro is incompatible with const fn")]
 pub fn omp_install_tokio_runtime() {
@@ -224,4 +300,39 @@ pub fn omp_install_tokio_runtime() {
 	if let Some(runtime) = create_windows_napi_tokio_runtime() {
 		create_custom_tokio_runtime(runtime);
 	}
+	#[cfg(target_os = "windows")]
+	configure_rayon_pool();
+}
+
+#[cfg(test)]
+mod tests {
+	use super::{RAYON_MAX_THREADS, RayonPoolPlan, clamped_rayon_threads, rayon_pool_plan};
+
+	#[test]
+	fn rayon_threads_are_capped_for_windows_commit_pressure() {
+		assert_eq!(clamped_rayon_threads(0), 1);
+		assert_eq!(clamped_rayon_threads(1), 1);
+		assert_eq!(clamped_rayon_threads(RAYON_MAX_THREADS + 1), RAYON_MAX_THREADS);
+	}
+
+	#[test]
+	fn rayon_uses_requested_pool_when_probe_covers_request_and_reserve() {
+		assert_eq!(rayon_pool_plan(4, 5), RayonPoolPlan::WorkerThreads(4));
+		assert_eq!(
+			rayon_pool_plan(RAYON_MAX_THREADS + 4, usize::MAX),
+			RayonPoolPlan::WorkerThreads(RAYON_MAX_THREADS)
+		);
+	}
+
+	#[test]
+	fn rayon_uses_worker_pool_after_reserving_non_rayon_capacity() {
+		assert_eq!(rayon_pool_plan(4, 3), RayonPoolPlan::WorkerThreads(2));
+		assert_eq!(rayon_pool_plan(1, 2), RayonPoolPlan::WorkerThreads(1));
+	}
+
+	#[test]
+	fn rayon_skips_global_pool_when_only_reserved_capacity_can_spawn() {
+		assert_eq!(rayon_pool_plan(4, 1), RayonPoolPlan::SkipGlobalPool);
+		assert_eq!(rayon_pool_plan(1, 0), RayonPoolPlan::SkipGlobalPool);
+	}
 }

crates/pi-natives/src/tokens.rs

@@ -17,6 +17,7 @@ use std::sync::LazyLock;
 
 use napi::bindgen_prelude::Either;
 use napi_derive::napi;
+use pi_uutils_ctx::rayon_global_pool_available;
 use rayon::prelude::*;
 use tiktoken_rs::{CoreBPE, cl100k_base, o200k_base};
 
@@ -45,9 +46,9 @@ fn encoder(encoding: Option<Encoding>) -> &'static CoreBPE {
 /// Count tokens in `input`.
 ///
 /// `input` may be a single string or an array of strings; an array returns
-/// the sum across all elements (encoded in parallel via rayon). Always
-/// returns a single token total — use this for any aggregate budget question
-/// without paying a per-element napi crossing.
+/// the sum across all elements (encoded in parallel via rayon when the global
+/// pool is available). Always returns a single token total — use this for any
+/// aggregate budget question without paying a per-element napi crossing.
 ///
 /// Uses ordinary encoding (no special-token handling), which is the right
 /// choice for measuring user/model content rather than wire-protocol tokens.
@@ -57,9 +58,13 @@ pub fn count_tokens(input: Either<String, Vec<String>>, encoding: Option<Encodin
 	let bpe = encoder(encoding);
 	match input {
 		Either::A(text) => bpe.encode_ordinary(&text).len() as u32,
-		Either::B(texts) => texts
+		Either::B(texts) if rayon_global_pool_available() => texts
 			.par_iter()
 			.map(|s| bpe.encode_ordinary(s).len() as u32)
 			.sum(),
+		Either::B(texts) => texts
+			.iter()
+			.map(|s| bpe.encode_ordinary(s).len() as u32)
+			.sum(),
 	}
 }

crates/pi-uutils-ctx/src/lib.rs

@@ -51,6 +51,8 @@ thread_local! {
 	static SCOPE_DEPTH: Cell<usize> = const { Cell::new(0) };
 }
 
+static RAYON_GLOBAL_POOL_AVAILABLE: AtomicBool = AtomicBool::new(!cfg!(target_os = "windows"));
+
 /// I/O streams, working directory, environment, and cancel flag for a single
 /// utility invocation. Grouped into one value to keep [`scope`] readable.
 pub struct ScopeIo {
@@ -124,6 +126,20 @@ pub fn is_active() -> bool {
 	SCOPE_DEPTH.with(|d| d.get() > 0)
 }
 
+/// Records whether patched native callsites may use Rayon's process-global
+/// worker pool without risking lazy initialization under Windows commit
+/// pressure.
+pub fn set_rayon_global_pool_available(available: bool) {
+	RAYON_GLOBAL_POOL_AVAILABLE.store(available, Ordering::SeqCst);
+}
+
+/// Returns whether patched native callsites may enter Rayon's process-global
+/// worker pool.
+#[must_use]
+pub fn rayon_global_pool_available() -> bool {
+	RAYON_GLOBAL_POOL_AVAILABLE.load(Ordering::SeqCst)
+}
+
 /// Returns the exit code accumulated via [`set_exit_code`] during the current
 /// scope (0 when none was set or no context is installed).
 pub fn exit_code() -> i32 {

crates/vendor/uu-sort/src/sort.rs

@@ -45,7 +45,7 @@ use custom_str_cmp::custom_str_cmp;
 use ext_sort::ext_sort;
 use foldhash::{HashMap, SharedSeed, fast::FoldHasher};
 use numeric_str_cmp::{NumInfo, NumInfoParseSettings, human_numeric_str_cmp, numeric_str_cmp};
-use pi_uutils_ctx::format_usage;
+use pi_uutils_ctx::{format_usage, rayon_global_pool_available};
 use rand::{RngExt as _, rng};
 #[cfg(not(target_os = "wasi"))]
 use rayon::slice::ParallelSliceMut;
@@ -2108,13 +2108,15 @@ fn uu_sort(matches: &ArgMatches, legacy_warnings: &[LegacyKeyWarning]) -> UResul
 			.map_or_else(|| "0".to_string(), String::from);
 		#[cfg(not(target_os = "wasi"))]
 		{
-			let num_threads = match settings.threads.parse::<usize>() {
-				Ok(0) | Err(_) => std::thread::available_parallelism().map_or(1, NonZero::get),
-				Ok(n) => n,
-			};
-			let _ = rayon::ThreadPoolBuilder::new()
-				.num_threads(num_threads)
-				.build_global();
+			if rayon_global_pool_available() {
+				let num_threads = match settings.threads.parse::<usize>() {
+					Ok(0) | Err(_) => std::thread::available_parallelism().map_or(1, NonZero::get),
+					Ok(n) => n,
+				};
+				let _ = rayon::ThreadPoolBuilder::new()
+					.num_threads(num_threads)
+					.build_global();
+			}
 		}
 	}
 
@@ -2584,15 +2586,25 @@ fn exec(
 fn sort_by<'a>(unsorted: &mut Vec<Line<'a>>, settings: &GlobalSettings, line_data: &LineData<'a>) {
 	let cmp = |a: &Line<'a>, b: &Line<'a>| compare_by(a, b, settings, line_data, line_data);
 	// WASI does not support threads, so use non-parallel sort to avoid
-	// rayon's thread pool which triggers an unreachable trap.
+	// rayon's thread pool which triggers an unreachable trap. Windows can also
+	// force sequential sort when pi-natives could not safely configure Rayon's
+	// process-global worker pool under commit pressure.
 	if settings.stable || settings.unique {
 		#[cfg(not(target_os = "wasi"))]
-		unsorted.par_sort_by(cmp);
+		if rayon_global_pool_available() {
+			unsorted.par_sort_by(cmp);
+		} else {
+			unsorted.sort_by(cmp);
+		}
 		#[cfg(target_os = "wasi")]
 		unsorted.sort_by(cmp);
 	} else {
 		#[cfg(not(target_os = "wasi"))]
-		unsorted.par_sort_unstable_by(cmp);
+		if rayon_global_pool_available() {
+			unsorted.par_sort_unstable_by(cmp);
+		} else {
+			unsorted.sort_unstable_by(cmp);
+		}
 		#[cfg(target_os = "wasi")]
 		unsorted.sort_unstable_by(cmp);
 	}

packages/natives/CHANGELOG.md

@@ -20,6 +20,9 @@
 
 - Added embedded Silver TrueType font rendering support to `renderSnapcompactPng`, featuring automatic per-glyph fallback for missing bitmap characters and anti-aliased scaling for East Asian wide code points.
 - Added the `snapcompactSupportedChars` function to check font capability for specific characters.
+### Fixed
+
+- Fixed `pi-natives` aborting on Windows under low commit charge when Rayon's global pool lazily spawned one worker per logical CPU for `count_tokens` or vendored `sort`. The post-load native runtime installer now configures Rayon's global pool through the same Windows thread-spawn probe used for Tokio and keeps patched Rayon callsites sequential if no worker remains after reserving capacity for native helper threads ([#3770](https://github.com/can1357/oh-my-pi/issues/3770)).
 
 ## [16.2.5] - 2026-06-28