diff --git a/crates/pm/src/util/user_config.rs b/crates/pm/src/util/user_config.rs index 34ee45a340..d3618e4a85 100644 --- a/crates/pm/src/util/user_config.rs +++ b/crates/pm/src/util/user_config.rs @@ -132,9 +132,10 @@ pub fn get_install_scope() -> InstallScope { INSTALL_SCOPE.get().copied().unwrap_or_default() } -// Manifest fetch concurrency configuration +// Manifest fetch concurrency configuration. Used by both preload (worker +// pool size) and downloader (Semaphore for tarball fetches). static MANIFESTS_CONCURRENCY_LIMIT: LazyLock> = - LazyLock::new(|| ConfigValue::new("manifests-concurrency-limit", 64)); + LazyLock::new(|| ConfigValue::new("manifests-concurrency-limit", 128)); pub fn set_manifests_concurrency_limit(value: Option) { MANIFESTS_CONCURRENCY_LIMIT.set(value); diff --git a/crates/ruborist/src/resolver/builder.rs b/crates/ruborist/src/resolver/builder.rs index b03db0fc43..5e6ffb1ffa 100644 --- a/crates/ruborist/src/resolver/builder.rs +++ b/crates/ruborist/src/resolver/builder.rs @@ -34,7 +34,7 @@ use crate::resolver::preload::{PreloadConfig, preload_manifests}; use crate::resolver::registry::{ResolveError, resolve_registry_dep}; use crate::spec::{Catalogs, PackageSpec, Protocol}; use crate::traits::progress::{BuildEvent, EventReceiver, NoopReceiver}; -use crate::traits::registry::{RegistryClient, ResolvedPackage}; +use crate::traits::registry::{PreloadRegistry, RegistryClient, ResolvedPackage}; /// Dispatch a git/github spec to the real `gix`-backed resolver when the /// `native-git` feature is enabled, otherwise error with a hint. @@ -707,7 +707,7 @@ pub async fn process_dependency( /// // Add initial dependency edges to root... /// build_deps(&mut graph, ®istry, PeerDeps::Include).await?; /// ``` -pub async fn build_deps( +pub async fn build_deps( graph: &mut DependencyGraph, registry: &R, peer_deps: PeerDeps, @@ -729,7 +729,7 @@ pub async fn build_deps( /// * `registry` - Registry client for fetching packages /// * `peer_deps` - How to handle peer dependencies /// * `receiver` - Event receiver for handling build events -pub async fn build_deps_with_receiver( +pub async fn build_deps_with_receiver( graph: &mut DependencyGraph, registry: &R, peer_deps: PeerDeps, @@ -759,7 +759,7 @@ pub async fn build_deps_with_receiver( /// /// build_deps_with_config(&mut graph, ®istry, config, &receiver).await?; /// ``` -pub async fn build_deps_with_config( +pub async fn build_deps_with_config( graph: &mut DependencyGraph, registry: &R, config: BuildDepsConfig, @@ -786,7 +786,7 @@ pub async fn build_deps_with_config( } /// Run the preload phase to warm up the cache with manifests. -async fn run_preload_phase( +async fn run_preload_phase( graph: &DependencyGraph, registry: &R, config: &BuildDepsConfig, @@ -804,18 +804,22 @@ async fn run_preload_phase( } tracing::debug!("Preload phase: {} initial dependencies", initial_deps.len()); - receiver.on_event(BuildEvent::PreloadStart { - count: initial_deps.len(), - }); + // PreloadStart / PreloadComplete are emitted inside `preload_manifests`. let preload_config = PreloadConfig { peer_deps: config.peer_deps, concurrency: config.concurrency, }; + // Wrap registry in Rc for the worker pool. Each worker holds a + // Rc so resolution runs as an independent `spawn_local` task on + // the current thread. UnifiedRegistry's Clone is shallow (Arc-based), + // so this clone is cheap. + let registry_rc = std::rc::Rc::new(registry.clone()); + let stats = preload_manifests( initial_deps, - registry, + registry_rc, preload_config, receiver, |_name, _manifest| { @@ -829,10 +833,6 @@ async fn run_preload_phase( stats.success_count, stats.failed_count ); - receiver.on_event(BuildEvent::PreloadComplete { - success: stats.success_count, - failed: stats.failed_count, - }); tracing::debug!("Preload phase: {:?}", start.elapsed()); } @@ -961,7 +961,7 @@ use std::path::Path; /// let pkg: PackageJson = serde_json::from_str(&pkg_content)?; /// let lock = resolve(&pkg, ®istry).await?; /// ``` -pub async fn resolve( +pub async fn resolve( pkg: &PackageJson, registry: &R, ) -> Result> { @@ -975,7 +975,7 @@ pub async fn resolve( /// * `registry` - Registry client for fetching packages /// * `peer_deps` - How to handle peer dependencies /// * `receiver` - Event receiver for progress tracking -pub async fn resolve_with_options( +pub async fn resolve_with_options( pkg: &PackageJson, registry: &R, peer_deps: PeerDeps, diff --git a/crates/ruborist/src/resolver/preload.rs b/crates/ruborist/src/resolver/preload.rs index 9fb5aff67b..090171f122 100644 --- a/crates/ruborist/src/resolver/preload.rs +++ b/crates/ruborist/src/resolver/preload.rs @@ -1,21 +1,38 @@ -//! Parallel manifest preloading for dependency resolution. +//! Parallel manifest preloading via single-thread worker pool. //! -//! Uses FuturesUnordered for true streaming concurrency: when a package resolves, -//! its transitive dependencies are immediately added to the queue. - +//! N long-lived `tokio::task::spawn_local` workers pull work items from a +//! shared `VecDeque`. Replaces the prior `FuturesUnordered` design where +//! the main task owned all preload futures and polled them cooperatively +//! — every per-future `await` continuation (cache check / OnceMap / +//! RetryIf / reqwest send / bytes / parse / transitive dispatch) ran on +//! the same single task, saturating it. CI ant-design preload sustained +//! avg_conc ~55-60 even when the standalone manifest-bench (same reqwest +//! stack, no resolver overhead) hit 90+ at the same nominal cap. Splitting +//! the futures into N independent `spawn_local` tasks lets tokio schedule +//! each future's continuation independently while still running all +//! workers on the current OS thread — no `Send`/`Sync` bound on the +//! registry trait, wasm-compatible. + +use std::cell::{Cell, RefCell}; use std::collections::{HashSet, VecDeque}; +use std::rc::Rc; use std::sync::Arc; -use futures::stream::{FuturesUnordered, StreamExt}; +use tokio::sync::{Notify, mpsc}; +use tokio::task::LocalSet; use crate::model::manifest::CoreVersionManifest; use crate::model::node::PeerDeps; +use crate::resolver::registry::ResolveError; use crate::resolver::registry::resolve_package; use crate::traits::progress::{BuildEvent, EventReceiver}; -use crate::traits::registry::RegistryClient; +use crate::traits::registry::{RegistryClient, ResolvedPackage}; -/// Default concurrency limit for manifest fetching -pub const DEFAULT_CONCURRENCY: usize = 64; +/// Default concurrency limit for manifest fetching. +/// +/// Number of long-lived `spawn_local` workers. Each processes one +/// `resolve_package` at a time on the current thread. +pub const DEFAULT_CONCURRENCY: usize = 128; /// A dependency spec: (name, version_spec) pub type Dep = (String, String); @@ -71,113 +88,223 @@ fn extract_transitive_deps(manifest: &CoreVersionManifest, config: &PreloadConfi deps } -/// Preload all package manifests in parallel with streaming concurrency. +/// Result message sent from worker to main task: name, resolve result, +/// per-request elapsed ms, count of new transitives queued by this worker. +type Completion = (String, Result>, u64, usize); + +/// Preload all package manifests in parallel via a single-thread worker pool. +/// +/// Spawns N long-lived `spawn_local` workers on a `LocalSet` that all share +/// a `Rc>>` work queue and a `Rc>` +/// dedup set. Workers pull, call `resolve_package`, push transitives, and +/// send completions back via an unbounded mpsc channel which the main task +/// drains. +/// +/// `registry` is moved in and shared across workers via `Rc` (single-thread +/// reference count, no `Send` required). `receiver` and `on_manifest` are +/// borrowed for the lifetime of the call. pub async fn preload_manifests( initial_deps: Vec, - registry: &R, + registry: Rc, config: PreloadConfig, receiver: &E, mut on_manifest: F, ) -> PreloadStats where - R: RegistryClient, + R: RegistryClient + 'static, E: EventReceiver, F: FnMut(&str, Arc), { let mut stats = PreloadStats::default(); - let mut processed: HashSet = HashSet::new(); - let mut pending: VecDeque = initial_deps.into(); - let concurrency = config.concurrency; + + // Shared single-thread state. `Rc>` is fine here — + // every worker is `spawn_local`-ed onto the same OS thread, so no + // synchronisation primitives are needed. + let pending: Rc>> = Rc::new(RefCell::new(VecDeque::new())); + let processed: Rc>> = Rc::new(RefCell::new(HashSet::new())); + + // Counters for global termination — when `dispatched == completed` + // and the queue is empty, the phase is done. + let dispatched: Rc> = Rc::new(Cell::new(0)); + let completed: Rc> = Rc::new(Cell::new(0)); + let shutdown: Rc> = Rc::new(Cell::new(false)); + + // Wake-up signal for workers parked on an empty queue. + let notify: Rc = Rc::new(Notify::new()); + + // Result channel — workers send completions, main task drains. + let (result_tx, mut result_rx) = mpsc::unbounded_channel::>(); + + // Seed the queue with initial deps (deduped via the shared set). + { + let mut p = pending.borrow_mut(); + let mut s = processed.borrow_mut(); + for dep in initial_deps { + let key = format!("{}@{}", dep.0, dep.1); + if s.insert(key) { + p.push_back(dep); + dispatched.set(dispatched.get() + 1); + } + } + } + let initial_count = dispatched.get(); + let concurrency = config.concurrency.max(1); tracing::debug!( - "Preload: {} initial deps, concurrency={}", - pending.len(), + "Preload: {} initial deps, concurrency={}, mode=spawn_local-pool", + initial_count, concurrency ); - let mut futures = FuturesUnordered::new(); - let mut in_flight = 0usize; - let mut started = false; - - loop { - // Fill up to concurrency limit - while in_flight < concurrency { - let item = loop { - let Some((name, spec)) = pending.pop_front() else { - break None; - }; - let key = format!("{}@{}", name, spec); - if !processed.contains(&key) { - processed.insert(key); - break Some((name, spec)); - } - }; + // Short-circuit if nothing was seeded. + if initial_count == 0 { + receiver.on_event(BuildEvent::PreloadStart { count: 0 }); + receiver.on_event(BuildEvent::PreloadComplete { + success: 0, + failed: 0, + }); + return stats; + } - let Some((name, spec)) = item else { - break; - }; + // Run all worker spawns + the main drain loop inside a LocalSet so + // `spawn_local` is valid and worker futures can borrow `&R`. + let local = LocalSet::new(); + + for _ in 0..concurrency { + let pending = Rc::clone(&pending); + let processed = Rc::clone(&processed); + let dispatched = Rc::clone(&dispatched); + let completed = Rc::clone(&completed); + let shutdown = Rc::clone(&shutdown); + let notify = Rc::clone(¬ify); + let result_tx = result_tx.clone(); + let config_for_worker = config.clone(); + let registry = Rc::clone(®istry); + + local.spawn_local(async move { + loop { + // Try fetching work first — fast path when queue is hot. + let work = pending.borrow_mut().pop_front(); + if let Some((name, spec)) = work { + let start = tokio::time::Instant::now(); + let result = resolve_package(&*registry, &name, &spec).await; + let elapsed_ms = start.elapsed().as_millis() as u64; + + let new_added = if let Ok(resolved) = &result { + let mut count = 0usize; + for tdep in extract_transitive_deps(&resolved.manifest, &config_for_worker) + { + let key = format!("{}@{}", tdep.0, tdep.1); + if processed.borrow_mut().insert(key) { + pending.borrow_mut().push_back(tdep); + count += 1; + } + } + if count > 0 { + dispatched.set(dispatched.get() + count); + // Wake parked workers so they pick up the new + // work before checking the termination condition. + notify.notify_waiters(); + } + count + } else { + 0 + }; + + if result_tx + .send((name, result, elapsed_ms, new_added)) + .is_err() + { + // Main task dropped the receiver — done. + break; + } + let done = completed.get() + 1; + completed.set(done); + + // After completion, check global done condition. + if done == dispatched.get() && pending.borrow().is_empty() { + shutdown.set(true); + notify.notify_waiters(); + break; + } + continue; + } - if !started { - receiver.on_event(BuildEvent::PreloadStart { count: 1 }); - started = true; - } else { - receiver.on_event(BuildEvent::PreloadQueued { count: 1 }); + // Queue empty — register interest before re-checking, then + // park. The Notify+enable() pattern guarantees we won't + // miss a `notify_waiters()` racing with our park. + if shutdown.get() { + break; + } + let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if !pending.borrow().is_empty() { + continue; + } + if shutdown.get() { + break; + } + if completed.get() == dispatched.get() { + shutdown.set(true); + notify.notify_waiters(); + break; + } + notified.await; } - - receiver.on_event(BuildEvent::PreloadFetching { name: &name }); - - futures.push(async move { - let start = tokio::time::Instant::now(); - let result = resolve_package(registry, &name, &spec).await; - let elapsed_ms = start.elapsed().as_millis() as u64; - (name, result, elapsed_ms) + }); + } + // Drop the original sender so when all worker clones drop on exit, the + // result channel closes and the main loop terminates. + drop(result_tx); + + // Main drain loop — runs inside the LocalSet alongside the workers. + local + .run_until(async { + receiver.on_event(BuildEvent::PreloadStart { + count: initial_count, }); - in_flight += 1; - } - - if in_flight == 0 { - break; - } - - let Some((name, result, elapsed_ms)) = futures.next().await else { - break; - }; - in_flight -= 1; - - if stats.success_count == 0 && stats.failed_count == 0 { - stats.min_request_ms = elapsed_ms; - stats.max_request_ms = elapsed_ms; - } else { - stats.min_request_ms = stats.min_request_ms.min(elapsed_ms); - stats.max_request_ms = stats.max_request_ms.max(elapsed_ms); - } - stats.total_request_ms += elapsed_ms; - - match result { - Ok(resolved) => { - stats.success_count += 1; - tracing::debug!("Preloaded {}@{}", name, resolved.version); - receiver.on_event(BuildEvent::PreloadProgress { - name: &name, - version: &resolved.version, - current: stats.success_count, - }); + while let Some((name, result, elapsed_ms, new_added)) = result_rx.recv().await { + if stats.success_count == 0 && stats.failed_count == 0 { + stats.min_request_ms = elapsed_ms; + stats.max_request_ms = elapsed_ms; + } else { + stats.min_request_ms = stats.min_request_ms.min(elapsed_ms); + stats.max_request_ms = stats.max_request_ms.max(elapsed_ms); + } + stats.total_request_ms += elapsed_ms; - // Send PackageResolved event for pipeline downloading - receiver.on_event(BuildEvent::PackageResolved((&*resolved.manifest).into())); + // Grow the progress bar length as transitives are discovered. + if new_added > 0 { + receiver.on_event(BuildEvent::PreloadQueued { count: new_added }); + } - pending.extend(extract_transitive_deps(&resolved.manifest, &config)); - on_manifest(&name, resolved.manifest); - } - Err(e) => { - stats.failed_count += 1; - tracing::debug!("Failed to preload {}: {}", name, e); + match result { + Ok(resolved) => { + stats.success_count += 1; + tracing::debug!("Preloaded {}@{}", name, resolved.version); + + receiver.on_event(BuildEvent::PreloadProgress { + name: &name, + version: &resolved.version, + current: stats.success_count, + }); + receiver + .on_event(BuildEvent::PackageResolved((&*resolved.manifest).into())); + + on_manifest(&name, resolved.manifest); + } + Err(e) => { + stats.failed_count += 1; + tracing::debug!("Failed to preload {}: {}", name, e); + } + } } - } - } + }) + .await; - stats.total_processed = processed.len(); + stats.total_processed = processed.borrow().len(); receiver.on_event(BuildEvent::PreloadComplete { success: stats.success_count, @@ -247,7 +374,7 @@ mod tests { let stats = preload_manifests( vec![("lodash".into(), "^4.17.0".into())], - ®istry, + Rc::new(registry), PreloadConfig::default(), &NoopReceiver, |name, m| { @@ -275,7 +402,7 @@ mod tests { let stats = preload_manifests( vec![("a".into(), "^1.0.0".into())], - ®istry, + Rc::new(registry), PreloadConfig::default(), &NoopReceiver, |name, m| { @@ -297,7 +424,7 @@ mod tests { let stats = preload_manifests( vec![("nonexistent".into(), "^1.0.0".into())], - ®istry, + Rc::new(registry), PreloadConfig::default(), &NoopReceiver, |name, m| { diff --git a/crates/ruborist/src/traits/registry.rs b/crates/ruborist/src/traits/registry.rs index 0760ff9b19..c2f318ba4a 100644 --- a/crates/ruborist/src/traits/registry.rs +++ b/crates/ruborist/src/traits/registry.rs @@ -314,12 +314,26 @@ pub trait RegistryClient { } } +/// Convenience bound for registries usable by the parallel preload worker pool. +/// +/// `Clone + 'static` lets callers wrap the registry in `Rc` and share one +/// owned handle across N `spawn_local` workers; without `'static` the +/// spawned worker future couldn't outlive any borrow, and without `Clone` +/// the caller couldn't hand the worker pool an owned copy. Both +/// `UnifiedRegistry` (real client, `Clone` is `Arc`-based and shallow) and +/// `MockRegistryClient` (test client) satisfy this trivially, so the bound +/// is invisible at every real call site — the alias just keeps the bound +/// list short on the resolver entry-point signatures. +pub trait PreloadRegistry: RegistryClient + Clone + 'static {} +impl PreloadRegistry for T {} + /// A simple in-memory registry client for testing. #[cfg(test)] pub mod mock { use super::*; /// Internal package data for mock registry. + #[derive(Clone)] struct MockPackage { name: String, dist_tags: HashMap, @@ -327,6 +341,7 @@ pub mod mock { } /// Mock registry client that returns predefined packages. + #[derive(Clone)] pub struct MockRegistryClient { packages: HashMap, }