diff --git a/crates/pi-natives/src/lib.rs b/crates/pi-natives/src/lib.rs index 97d30648f..24ddeca0f 100644 --- a/crates/pi-natives/src/lib.rs +++ b/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); + } } diff --git a/crates/pi-natives/src/tokens.rs b/crates/pi-natives/src/tokens.rs index 0b524c169..3b02c4cdf 100644 --- a/crates/pi-natives/src/tokens.rs +++ b/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) -> &'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>, encoding: Option 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(), } } diff --git a/crates/pi-uutils-ctx/src/lib.rs b/crates/pi-uutils-ctx/src/lib.rs index 1009e2ca5..05174f86e 100644 --- a/crates/pi-uutils-ctx/src/lib.rs +++ b/crates/pi-uutils-ctx/src/lib.rs @@ -51,6 +51,8 @@ thread_local! { static SCOPE_DEPTH: Cell = 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 { diff --git a/crates/vendor/uu-sort/src/sort.rs b/crates/vendor/uu-sort/src/sort.rs index fdd82ee18..2ba83e946 100644 --- a/crates/vendor/uu-sort/src/sort.rs +++ b/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::() { - 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::() { + 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>, 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); } diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 658aa220b..617c2571a 100644 --- a/packages/natives/CHANGELOG.md +++ b/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