diff --git a/bench/pm-bench-pcap.sh b/bench/pm-bench-pcap.sh index c8e869d78e..cbadda3e96 100755 --- a/bench/pm-bench-pcap.sh +++ b/bench/pm-bench-pcap.sh @@ -19,12 +19,11 @@ export BUN_INSTALL_CACHE_DIR="$BUN_CACHE" mkdir -p "$PCAP_DIR" -# Project must already be cloned by pm-bench-phases.sh (this script is meant -# to run after it, reusing the clone). PROJECT_DIR="$BENCH_DIR/$PROJECT" +mkdir -p "$BENCH_DIR" if [ ! -d "$PROJECT_DIR" ]; then - echo "missing $PROJECT_DIR — run pm-bench-phases.sh first" >&2 - exit 1 + echo "cloning $PROJECT into $PROJECT_DIR" + git clone --depth=1 "https://github.com/ant-design/${PROJECT}.git" "$PROJECT_DIR" fi # Extract hostname (strip scheme + any path) for DNS lookup. diff --git a/crates/ruborist/Cargo.toml b/crates/ruborist/Cargo.toml index 4f63f3ec7f..49b3b53273 100644 --- a/crates/ruborist/Cargo.toml +++ b/crates/ruborist/Cargo.toml @@ -7,57 +7,72 @@ license = "MIT" # tombi: format.rules.table-keys-order.disabled = true [dependencies] -anyhow = { workspace = true } -bytes = "1" -pathdiff = { workspace = true } -deno_semver = "0.7" -dirs = { workspace = true } -futures = "0.3" -parking_lot = { workspace = true } -petgraph = "0.6" -serde = { version = "1.0", features = ["derive"] } -serde_json = { workspace = true } -simd-json = "0.17" -thiserror = { workspace = true } -tokio-fs-ext = { workspace = true } -tokio-retry = "0.3" -tracing = { workspace = true } +anyhow = { workspace = true } +bytes = "1" +pathdiff = { workspace = true } +crossbeam-queue = "0.3" +dashmap = "6" +deno_semver = "0.7" +dirs = { workspace = true } +futures = "0.3" +parking_lot = { workspace = true } +petgraph = "0.6" +serde = { version = "1.0", features = ["derive"] } +serde_json = { workspace = true } +simd-json = "0.17" +thiserror = { workspace = true } +tokio-fs-ext = { workspace = true } +tokio-retry = "0.3" +tracing = { workspace = true } # HTTP client - works on both native and WASM -reqwest = { version = "0.12", default-features = false, features = [ +# +# We use `rustls-tls-native-roots-no-provider` (instead of the default +# `rustls-tls-native-roots`) so reqwest doesn't pin the crypto provider +# to `ring`. We then build our own `rustls::ClientConfig` with +# `aws-lc-rs` (BoringSSL's crypto primitives) in `service/http.rs` and +# pass it via `Client::use_preconfigured_tls`. Measured on CI against +# npmjs.org: `ring` took 78 ms p50 / 154 ms max from CCS → first +# AppData per handshake (CPU-bound; 128 parallel handshakes serialise +# across 4 cores). aws-lc-rs's per-handshake CPU is roughly 3× faster. +reqwest = { version = "0.12", default-features = false, features = [ "http2", "json", - "rustls-tls", - "rustls-tls-native-roots", + "rustls-tls-native-roots-no-provider", "socks" ] } # Shared URL-hash + staging helpers for native-only resolvers (git / http) -sha2 = { workspace = true, optional = true } -tempfile = { workspace = true, optional = true } +sha2 = { workspace = true, optional = true } +tempfile = { workspace = true, optional = true } # Git clone (gated on `native-git`) -gix = { workspace = true, features = [ +gix = { workspace = true, features = [ "blocking-http-transport-reqwest-rust-tls", "blocking-network-client", ], optional = true } # HTTP tarball extraction (gated on `http-tarball`) # Uses libdeflater to match pm's extractor and avoid pulling a second gzip # implementation (flate2) into the workspace. -libdeflater = { version = "1.25", optional = true } -tar = { version = "0.4", optional = true } +libdeflater = { version = "1.25", optional = true } +tar = { version = "0.4", optional = true } # Async runtime (works on both native and WASM with forked tokio) [dependencies.tokio] workspace = true [dev-dependencies] -tokio = { workspace = true, features = ["macros", "rt"] } +tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } # Native-only dependencies (not compiled for WASM) [target.'cfg(not(target_arch = "wasm32"))'.dependencies] -libc = "0.2" +libc = "0.2" +# Custom rustls config with aws-lc-rs crypto provider — see service/http.rs. +# wasm32 falls back to reqwest's default rustls + ring (transitive via reqwest). +rustls = { version = "0.23", default-features = false, features = ["aws-lc-rs", "logging", "std", "tls12"] } +rustls-native-certs = "0.8" # WASM-only dependencies [target.'cfg(target_arch = "wasm32")'.dependencies] -js-sys = "0.3" # For time in WASM +js-sys = "0.3" # For time in WASM +wasm-bindgen-futures = "0.4" # spawn_local replaces tokio::spawn on wasm [features] default = [] diff --git a/crates/ruborist/src/lib.rs b/crates/ruborist/src/lib.rs index 890d6c2798..2896da21b4 100644 --- a/crates/ruborist/src/lib.rs +++ b/crates/ruborist/src/lib.rs @@ -22,11 +22,13 @@ //! }).await?; //! ``` +pub mod maybe_send; pub mod model; pub mod resolver; pub mod service; pub mod spec; pub mod traits; +pub mod util; // ============================================================================ // Re-exports: only items actually used by consumers @@ -99,8 +101,3 @@ pub mod http { #[cfg(feature = "http-tarball")] pub use crate::resolver::http::{file_cache_slot, http_cache_slot}; } - -/// Utility functions. -pub mod util { - pub use crate::model::util::{PackageNameStr, parse_package_spec, read_package_json}; -} diff --git a/crates/ruborist/src/maybe_send.rs b/crates/ruborist/src/maybe_send.rs new file mode 100644 index 0000000000..309f901ea0 --- /dev/null +++ b/crates/ruborist/src/maybe_send.rs @@ -0,0 +1,27 @@ +//! `Send`/`Sync` shim that is `Send`/`Sync` on native targets and +//! no-op on `wasm32`. Lets the resolver share a single trait surface +//! across multi-thread tokio (where futures must be `Send` to be +//! `tokio::spawn`-able) and wasm-bindgen (where `JsFuture` is `!Send` +//! and tasks run via `wasm_bindgen_futures::spawn_local`). + +/// `Send` on native, no-op on wasm32. +#[cfg(not(target_arch = "wasm32"))] +pub trait MaybeSend: Send {} +#[cfg(not(target_arch = "wasm32"))] +impl MaybeSend for T {} + +#[cfg(target_arch = "wasm32")] +pub trait MaybeSend {} +#[cfg(target_arch = "wasm32")] +impl MaybeSend for T {} + +/// `Sync` on native, no-op on wasm32. +#[cfg(not(target_arch = "wasm32"))] +pub trait MaybeSync: Sync {} +#[cfg(not(target_arch = "wasm32"))] +impl MaybeSync for T {} + +#[cfg(target_arch = "wasm32")] +pub trait MaybeSync {} +#[cfg(target_arch = "wasm32")] +impl MaybeSync for T {} diff --git a/crates/ruborist/src/resolver/builder.rs b/crates/ruborist/src/resolver/builder.rs index b03db0fc43..9cb09ea96f 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. @@ -458,13 +458,16 @@ async fn process_file_dep( /// /// # Returns /// The result of processing (reused, created, or skipped) -pub async fn process_dependency( +pub async fn process_dependency( graph: &mut DependencyGraph, registry: &R, node_index: NodeIndex, edge_info: &DependencyEdgeInfo, config: &BuildDepsConfig, -) -> Result> { +) -> Result> +where + R::Error: crate::maybe_send::MaybeSend, +{ tracing::debug!( "Processing dependency {}@{} from [{:?}]", edge_info.name, @@ -707,11 +710,14 @@ 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, -) -> Result<(), ResolveError> { +) -> Result<(), ResolveError> +where + R::Error: crate::maybe_send::MaybeSend, +{ let config = BuildDepsConfig::default().with_peer_deps(peer_deps); build_deps_with_config(graph, registry, config, &NoopReceiver).await } @@ -729,12 +735,15 @@ 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, receiver: &E, -) -> Result<(), ResolveError> { +) -> Result<(), ResolveError> +where + R::Error: crate::maybe_send::MaybeSend, +{ let config = BuildDepsConfig::default().with_peer_deps(peer_deps); build_deps_with_config(graph, registry, config, receiver).await } @@ -759,12 +768,15 @@ 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, receiver: &E, -) -> Result<(), ResolveError> { +) -> Result<(), ResolveError> +where + R::Error: crate::maybe_send::MaybeSend, +{ tracing::debug!( "Starting dependency tree build, peer_deps: {:?}, concurrency: {}, skip_preload: {}", config.peer_deps, @@ -786,12 +798,14 @@ 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, receiver: &E, -) { +) where + R::Error: crate::maybe_send::MaybeSend, +{ if config.skip_preload { return; } @@ -804,18 +818,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 Arc for the worker pool. Each worker holds an + // Arc so resolution runs as an independent `tokio::spawn` task on + // the multi-thread runtime. UnifiedRegistry's Clone is shallow + // (Arc-based), so this clone is cheap. + let registry_arc = std::sync::Arc::new(registry.clone()); + let stats = preload_manifests( initial_deps, - registry, + registry_arc, preload_config, receiver, |_name, _manifest| { @@ -829,21 +847,20 @@ 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()); } /// Run the BFS traversal phase to build the dependency tree. -async fn run_bfs_phase( +async fn run_bfs_phase( graph: &mut DependencyGraph, registry: &R, config: &BuildDepsConfig, receiver: &E, -) -> Result<(), ResolveError> { +) -> Result<(), ResolveError> +where + R::Error: crate::maybe_send::MaybeSend, +{ let start = tokio::time::Instant::now(); let mut current_level = vec![graph.root_index]; @@ -961,10 +978,13 @@ 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> { +) -> Result> +where + R::Error: crate::maybe_send::MaybeSend, +{ resolve_with_options(pkg, registry, PeerDeps::Include, &NoopReceiver).await } @@ -975,12 +995,15 @@ 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, receiver: &E, -) -> Result> { +) -> Result> +where + R::Error: crate::maybe_send::MaybeSend, +{ // Create graph with root node let mut graph = DependencyGraph::from_package_json(PathBuf::from("."), pkg.clone()); diff --git a/crates/ruborist/src/resolver/preload.rs b/crates/ruborist/src/resolver/preload.rs index 9fb5aff67b..8b092869b5 100644 --- a/crates/ruborist/src/resolver/preload.rs +++ b/crates/ruborist/src/resolver/preload.rs @@ -1,20 +1,39 @@ -//! Parallel manifest preloading for dependency resolution. +//! Parallel manifest preloading via multi-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::spawn` workers pull work items from a shared +//! lock-free `SegQueue`. Each spawned task is independent on tokio's +//! multi-thread runtime, so when one worker is parsing manifest JSON +//! (CPU-bound, simd_json), other workers can still drive their network +//! IO. 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, plus all parses serialised in that task's polling. +//! +//! Wasm32 fallback uses `wasm_bindgen_futures::spawn_local` since +//! `JsFuture` is `!Send`. Workers still run independently on the JS +//! event loop; the queue + Notify + mpsc termination story is identical. -use std::collections::{HashSet, VecDeque}; +use std::collections::HashSet; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use futures::stream::{FuturesUnordered, StreamExt}; +use crossbeam_queue::SegQueue; +use dashmap::DashSet; +use tokio::sync::{Notify, mpsc}; +use crate::maybe_send::{MaybeSend, MaybeSync}; 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 +/// Default concurrency limit for manifest fetching. +/// +/// Number of long-lived `tokio::spawn` workers. Each processes one +/// `resolve_package` at a time on tokio's multi-thread runtime. pub const DEFAULT_CONCURRENCY: usize = 64; /// A dependency spec: (name, version_spec) @@ -71,79 +90,179 @@ 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 multi-thread worker pool. +/// +/// `registry` is moved in and shared across workers via `Arc`. Each +/// worker is a long-lived spawned task that pulls work items from a +/// shared lock-free `SegQueue` until both the queue and the in-flight +/// counter reach zero, signalling end-of-phase. Each spawned task is +/// independent on tokio's multi-thread runtime — IO and CPU work +/// (simd_json parse) parallelise across cores. +/// +/// `receiver` and `on_manifest` run on the main task only — they do +/// not need to be `Send`/`Sync`. pub async fn preload_manifests( initial_deps: Vec, - registry: &R, + registry: Arc, config: PreloadConfig, receiver: &E, mut on_manifest: F, ) -> PreloadStats where - R: RegistryClient, + R: RegistryClient + MaybeSend + MaybeSync + 'static, + R::Error: MaybeSend, 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; - - tracing::debug!( - "Preload: {} initial deps, concurrency={}", - pending.len(), - 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)); - } - }; + // Shared work queue and dedup set (lock-free, multi-thread safe). + let pending: Arc> = Arc::new(SegQueue::new()); + let processed: Arc> = Arc::new(DashSet::new()); - let Some((name, spec)) = item else { - break; - }; + // Counters for global termination — when `dispatched == completed` + // and the queue is empty, the phase is done. + let dispatched: Arc = Arc::new(AtomicUsize::new(0)); + let completed: Arc = Arc::new(AtomicUsize::new(0)); + let shutdown: Arc = Arc::new(AtomicBool::new(false)); - if !started { - receiver.on_event(BuildEvent::PreloadStart { count: 1 }); - started = true; - } else { - receiver.on_event(BuildEvent::PreloadQueued { count: 1 }); - } + // Wake-up signal for workers parked on an empty queue. + let notify: Arc = Arc::new(Notify::new()); - receiver.on_event(BuildEvent::PreloadFetching { name: &name }); + // Result channel — workers send completions, main task drains. + let (result_tx, mut result_rx) = mpsc::unbounded_channel::>(); - 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) - }); - in_flight += 1; + // Seed the queue with initial deps (deduped via DashSet). + for dep in initial_deps { + let key = format!("{}@{}", dep.0, dep.1); + if processed.insert(key) { + pending.push(dep); + dispatched.fetch_add(1, Ordering::Relaxed); } + } + let initial_count = dispatched.load(Ordering::Relaxed); + let concurrency = config.concurrency.max(1); - if in_flight == 0 { - break; - } + tracing::debug!( + "Preload: {} initial deps, concurrency={}, mode=mt-worker-pool", + initial_count, + concurrency + ); - let Some((name, result, elapsed_ms)) = futures.next().await else { - break; + // 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; + } + + for _ in 0..concurrency { + let pending = Arc::clone(&pending); + let processed = Arc::clone(&processed); + let dispatched = Arc::clone(&dispatched); + let completed = Arc::clone(&completed); + let shutdown = Arc::clone(&shutdown); + let notify = Arc::clone(¬ify); + let result_tx = result_tx.clone(); + let config_for_worker = config.clone(); + let registry = Arc::clone(®istry); + + let worker = async move { + loop { + // Try fetching work first — fast path when queue is hot. + if let Some((name, spec)) = pending.pop() { + 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.insert(key) { + pending.push(tdep); + count += 1; + } + } + if count > 0 { + dispatched.fetch_add(count, Ordering::Release); + // 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.fetch_add(1, Ordering::AcqRel) + 1; + + // After completion, check global done condition. The + // `Acquire` on dispatched pairs with the `Release` in + // the producer add above, so we won't miss late + // transitives queued by a sibling worker. + if done == dispatched.load(Ordering::Acquire) && pending.is_empty() { + shutdown.store(true, Ordering::Release); + notify.notify_waiters(); + break; + } + continue; + } + + // 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.load(Ordering::Acquire) { + break; + } + let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if !pending.is_empty() { + continue; + } + if shutdown.load(Ordering::Acquire) { + break; + } + if completed.load(Ordering::Acquire) == dispatched.load(Ordering::Acquire) { + shutdown.store(true, Ordering::Release); + notify.notify_waiters(); + break; + } + notified.await; + } }; - in_flight -= 1; + #[cfg(not(target_arch = "wasm32"))] + tokio::spawn(worker); + #[cfg(target_arch = "wasm32")] + wasm_bindgen_futures::spawn_local(worker); + } + // Drop the original sender so when all worker clones drop on exit, the + // result channel closes and the main loop terminates. + drop(result_tx); + receiver.on_event(BuildEvent::PreloadStart { + count: initial_count, + }); + + // Main task: drain completions, run user callbacks. Receiver/callback + // run on this task only — they don't need to be Send/Sync. + 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; @@ -153,6 +272,10 @@ where } stats.total_request_ms += elapsed_ms; + if new_added > 0 { + receiver.on_event(BuildEvent::PreloadQueued { count: new_added }); + } + match result { Ok(resolved) => { stats.success_count += 1; @@ -163,11 +286,8 @@ where version: &resolved.version, current: stats.success_count, }); - - // Send PackageResolved event for pipeline downloading receiver.on_event(BuildEvent::PackageResolved((&*resolved.manifest).into())); - pending.extend(extract_transitive_deps(&resolved.manifest, &config)); on_manifest(&name, resolved.manifest); } Err(e) => { @@ -177,7 +297,14 @@ where } } - stats.total_processed = processed.len(); + // Snapshot the dedup set's size for the historical observable. + stats.total_processed = { + let mut set: HashSet = HashSet::with_capacity(processed.len()); + for entry in processed.iter() { + set.insert(entry.key().clone()); + } + set.len() + }; receiver.on_event(BuildEvent::PreloadComplete { success: stats.success_count, @@ -208,9 +335,8 @@ mod tests { use crate::model::manifest::CoreVersionManifest; use crate::traits::progress::NoopReceiver; use crate::traits::registry::mock::MockRegistryClient; - use std::cell::RefCell; use std::collections::HashMap; - use std::rc::Rc; + use std::sync::Mutex; fn manifest(name: &str, version: &str) -> CoreVersionManifest { CoreVersionManifest { @@ -237,30 +363,30 @@ mod tests { } } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_preload_single() { let mut registry = MockRegistryClient::new(); registry.add_package("lodash", "4.17.21", manifest("lodash", "4.17.21")); - let cache: Rc>>> = Default::default(); - let cache_clone = Rc::clone(&cache); + let cache: Arc>>> = Default::default(); + let cache_clone = Arc::clone(&cache); let stats = preload_manifests( vec![("lodash".into(), "^4.17.0".into())], - ®istry, + Arc::new(registry), PreloadConfig::default(), &NoopReceiver, - |name, m| { - cache_clone.borrow_mut().insert(name.into(), m); + move |name, m| { + cache_clone.lock().unwrap().insert(name.into(), m); }, ) .await; assert_eq!(stats.success_count, 1); - assert!(cache.borrow().contains_key("lodash")); + assert!(cache.lock().unwrap().contains_key("lodash")); } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_preload_transitive() { let mut registry = MockRegistryClient::new(); registry.add_package( @@ -270,44 +396,44 @@ mod tests { ); registry.add_package("b", "1.0.0", manifest("b", "1.0.0")); - let cache: Rc>>> = Default::default(); - let cache_clone = Rc::clone(&cache); + let cache: Arc>>> = Default::default(); + let cache_clone = Arc::clone(&cache); let stats = preload_manifests( vec![("a".into(), "^1.0.0".into())], - ®istry, + Arc::new(registry), PreloadConfig::default(), &NoopReceiver, - |name, m| { - cache_clone.borrow_mut().insert(name.into(), m); + move |name, m| { + cache_clone.lock().unwrap().insert(name.into(), m); }, ) .await; assert_eq!(stats.success_count, 2); - assert!(cache.borrow().contains_key("a")); - assert!(cache.borrow().contains_key("b")); + assert!(cache.lock().unwrap().contains_key("a")); + assert!(cache.lock().unwrap().contains_key("b")); } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_preload_missing() { let registry = MockRegistryClient::new(); - let cache: Rc>>> = Default::default(); - let cache_clone = Rc::clone(&cache); + let cache: Arc>>> = Default::default(); + let cache_clone = Arc::clone(&cache); let stats = preload_manifests( vec![("nonexistent".into(), "^1.0.0".into())], - ®istry, + Arc::new(registry), PreloadConfig::default(), &NoopReceiver, - |name, m| { - cache_clone.borrow_mut().insert(name.into(), m); + move |name, m| { + cache_clone.lock().unwrap().insert(name.into(), m); }, ) .await; assert_eq!(stats.failed_count, 1); - assert!(cache.borrow().is_empty()); + assert!(cache.lock().unwrap().is_empty()); } #[test] diff --git a/crates/ruborist/src/resolver/registry.rs b/crates/ruborist/src/resolver/registry.rs index 6edeb9069c..19bbfc9ffe 100644 --- a/crates/ruborist/src/resolver/registry.rs +++ b/crates/ruborist/src/resolver/registry.rs @@ -114,11 +114,14 @@ impl std::error::Error for ResolveError { /// let resolved = resolve_package(®istry, "lodash", "^4.0.0").await?; /// println!("Resolved to {}@{}", resolved.name, resolved.version); /// ``` -pub async fn resolve_package( +pub async fn resolve_package( registry: &R, name: &str, spec: &str, -) -> Result> { +) -> Result> +where + R::Error: crate::maybe_send::MaybeSend, +{ // Normalize spec first to handle npm: alias and workspace: prefix // This ensures correct behavior even if RegistryClient::resolve_package is overridden // e.g., "wrap-ansi-cjs" + "npm:wrap-ansi@^7.0.0" -> fetch "wrap-ansi" @ "^7.0.0" @@ -167,12 +170,15 @@ pub fn resolve_from_manifest( /// /// For optional dependencies, returns `Ok(None)` on resolution failure /// instead of propagating the error. -pub async fn resolve_registry_dep( +pub async fn resolve_registry_dep( registry: &R, name: &str, spec: &str, edge_type: &EdgeType, -) -> Result, ResolveError> { +) -> Result, ResolveError> +where + R::Error: crate::maybe_send::MaybeSend, +{ match resolve_package(registry, name, spec).await { Ok(resolved) => Ok(Some(resolved)), Err(e) => { diff --git a/crates/ruborist/src/service/http.rs b/crates/ruborist/src/service/http.rs index b959020973..9179f51b1a 100644 --- a/crates/ruborist/src/service/http.rs +++ b/crates/ruborist/src/service/http.rs @@ -93,6 +93,49 @@ pub(crate) fn get_client() -> Result<&'static reqwest::Client> { HTTP_CLIENT.as_ref().map_err(|e| anyhow!("{e}")) } +/// Build a `rustls::ClientConfig` using the `aws-lc-rs` crypto provider +/// instead of reqwest's default `ring`. Measured on CI (4-core runner) +/// against npmjs.org, ring's per-TLS-handshake client-side CPU cost +/// (ECDHE key derivation + cert verification + Finished MAC) serialised +/// across 128 parallel handshakes into a 154 ms "CCS → first AppData" +/// span — the HTTP requests couldn't fire until all TLS crypto drained +/// through 4 async workers. aws-lc-rs uses BoringSSL's assembly-optimised +/// primitives and is roughly 3× faster at handshake work. +#[cfg(not(target_arch = "wasm32"))] +fn build_rustls_config() -> Result { + // Install aws-lc-rs as the default for any other rustls consumer in + // the process. Idempotent — only the first call per process wins. + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + // Load OS root certs. On Linux this hits /etc/ssl/certs/, + // on macOS it queries the Security framework keychain. Called once + // per process via `HTTP_CLIENT`'s `LazyLock`. + let roots = rustls_native_certs::load_native_certs(); + let mut root_store = rustls::RootCertStore::empty(); + for cert in roots.certs { + // Best-effort: skip any cert rustls refuses (same tolerance + // native-tls shows). A hard fail here would brick every + // request on a box with one bad root in its trust store. + let _ = root_store.add(cert); + } + if !roots.errors.is_empty() { + tracing::debug!( + "rustls-native-certs reported {} non-fatal load issues", + roots.errors.len() + ); + } + + let config = rustls::ClientConfig::builder_with_provider(std::sync::Arc::new( + rustls::crypto::aws_lc_rs::default_provider(), + )) + .with_safe_default_protocol_versions() + .map_err(|e| anyhow!("rustls safe_default_protocol_versions: {e}"))? + .with_root_certificates(root_store) + .with_no_client_auth(); + + Ok(config) +} + /// Create a [`reqwest::ClientBuilder`] with TLS, DNS caching, and proxy /// from environment variables. /// @@ -112,8 +155,9 @@ pub fn client_builder() -> Result { let builder = { use crate::service::dns::shared_resolver; + let tls_config = build_rustls_config()?; let mut builder = builder - .use_rustls_tls() + .use_preconfigured_tls(tls_config) .no_proxy() .dns_resolver(shared_resolver()); diff --git a/crates/ruborist/src/service/registry.rs b/crates/ruborist/src/service/registry.rs index 772cac3f79..52c2c22201 100644 --- a/crates/ruborist/src/service/registry.rs +++ b/crates/ruborist/src/service/registry.rs @@ -68,6 +68,13 @@ pub struct UnifiedRegistry { registry_url: String, cache: Arc, supports_semver: bool, + /// Dedupes concurrent `resolve_full_manifest` fetches for the same + /// package name. First caller hits the network and stores the result; + /// other callers wait on `Notify` and read the shared `Arc`. Built on + /// `DashMap` + `tokio::sync::Notify` so the fast path (cache hit) is + /// lock-free. + #[cfg(not(target_arch = "wasm32"))] + inflight: Arc>, } /// Builder for `UnifiedRegistry`. @@ -137,6 +144,8 @@ impl UnifiedRegistryBuilder { registry_url, cache, supports_semver, + #[cfg(not(target_arch = "wasm32"))] + inflight: Arc::new(crate::util::oncemap::OnceMap::new()), } } } @@ -153,6 +162,8 @@ impl Clone for UnifiedRegistry { registry_url: self.registry_url.clone(), cache: Arc::clone(&self.cache), supports_semver: self.supports_semver, + #[cfg(not(target_arch = "wasm32"))] + inflight: Arc::clone(&self.inflight), } } } @@ -161,8 +172,13 @@ impl Clone for UnifiedRegistry { /// /// Separates the 200 (full data) and 304 (use cache) cases at the type level, /// so callers can pattern-match instead of string-matching error messages. -/// Transient return value, immediately destructured — Box not needed. +/// +/// `Clone` is required so multiple `resolve_full_manifest` callers that +/// coalesce through `OnceMap` can each take an owned copy of the shared +/// `Arc`. `FullManifest` clones cheaply via +/// `Arc<[u8]>` for its raw bytes, and the other fields are small. #[allow(clippy::large_enum_variant)] +#[derive(Clone)] enum FullManifestResult { /// Fresh manifest fetched from the network (HTTP 200). Full(FullManifest), @@ -202,13 +218,50 @@ impl UnifiedRegistry { /// 4. On 304: use disk cache data /// 5. On 200: update memory + disk cache async fn resolve_full_manifest(&self, name: &str) -> Result { - // 1. Check memory cache first + // 1. Fast path: memory cache hit (lock-free read). if let Some(manifest) = self.cache.get_full_manifest(name) { tracing::debug!("Memory cache hit for full manifest: {}", name); return Ok(FullManifestResult::Full(manifest)); } - // 2. Load etag from disk cache (if available) + // 2. Coalesce concurrent callers for the same name via OnceMap. + // First caller runs the fetch closure; others await the shared + // result on the OnceMap's `Notify` and clone the cached value. + // This avoids 50 packages all triggering 50 simultaneous fetches + // of `react` when only 1 is needed. + #[cfg(not(target_arch = "wasm32"))] + { + let shared = self + .inflight + .get_or_init(name.to_string(), || async { + self.fetch_full_manifest_network(name).await.ok() + }) + .await; + + match shared { + Some(arc) => Ok((*arc).clone()), + None => { + // OnceMap clears the key on None, so the next caller + // retries the fetch. Retry once here with a fresh error + // so we surface a useful message to this caller. + self.fetch_full_manifest_network(name).await + } + } + } + #[cfg(target_arch = "wasm32")] + { + self.fetch_full_manifest_network(name).await + } + } + + /// Perform the actual network fetch + cache update. Separated from + /// `resolve_full_manifest` so the OnceMap closure can invoke it + /// without re-entering the dedup layer. + async fn fetch_full_manifest_network( + &self, + name: &str, + ) -> Result { + // 1. Load etag from disk cache (if available) let disk_versions = self.cache.get_versions_from_disk(name).await; let etag = disk_versions.as_ref().and_then(|v| v.etag.clone()); diff --git a/crates/ruborist/src/traits/registry.rs b/crates/ruborist/src/traits/registry.rs index 0760ff9b19..efde5da0b0 100644 --- a/crates/ruborist/src/traits/registry.rs +++ b/crates/ruborist/src/traits/registry.rs @@ -140,7 +140,7 @@ pub trait RegistryClient { fn fetch_full_manifest( &self, name: &str, - ) -> impl Future>; + ) -> impl Future> + crate::maybe_send::MaybeSend; /// Fetch specific version manifest from registry. /// @@ -153,7 +153,11 @@ pub trait RegistryClient { &self, name: &str, spec: &str, - ) -> impl Future, Self::Error>> { + ) -> impl Future, Self::Error>> + crate::maybe_send::MaybeSend + where + Self: crate::maybe_send::MaybeSync, + Self::Error: crate::maybe_send::MaybeSend, + { async move { let manifest = self.fetch_full_manifest(name).await?; let version_list: Vec = manifest.versions.clone(); @@ -193,7 +197,11 @@ pub trait RegistryClient { &self, name: &str, spec: &str, - ) -> impl Future> { + ) -> impl Future> + crate::maybe_send::MaybeSend + where + Self: crate::maybe_send::MaybeSync, + Self::Error: crate::maybe_send::MaybeSend, + { async move { // Normalize spec (handles npm: alias and workspace: prefix) let (fetch_name, fetch_spec) = normalize_spec(name, spec); @@ -267,7 +275,11 @@ pub trait RegistryClient { fn fetch_versions_info( &self, name: &str, - ) -> impl Future> { + ) -> impl Future> + crate::maybe_send::MaybeSend + where + Self: crate::maybe_send::MaybeSync, + Self::Error: crate::maybe_send::MaybeSend, + { async move { let manifest = self.fetch_full_manifest(name).await?; Ok(VersionsInfo { @@ -314,12 +326,41 @@ pub trait RegistryClient { } } +/// Convenience bound for registries usable by the parallel preload worker pool. +/// +/// `Clone + 'static` lets callers share one owned handle across N spawned +/// workers (`Arc` on native, `Rc` on wasm). `MaybeSend + MaybeSync` +/// expand to `Send + Sync` on native (so the spawned worker future can be +/// scheduled across tokio's multi-thread runtime) and to no-op shims on +/// wasm (where `wasm_bindgen_futures::spawn_local` runs everything on the +/// JS event loop and `JsFuture` is `!Send`). Both `UnifiedRegistry` (real +/// client; `Clone` is shallow Arc-based) and `MockRegistryClient` (test +/// client) satisfy this trivially. +pub trait PreloadRegistry: + RegistryClient + Clone + crate::maybe_send::MaybeSend + crate::maybe_send::MaybeSync + 'static +where + Self::Error: crate::maybe_send::MaybeSend, +{ +} + +impl PreloadRegistry for T +where + T: RegistryClient + + Clone + + crate::maybe_send::MaybeSend + + crate::maybe_send::MaybeSync + + 'static, + T::Error: crate::maybe_send::MaybeSend, +{ +} + /// 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 +368,7 @@ pub mod mock { } /// Mock registry client that returns predefined packages. + #[derive(Clone)] pub struct MockRegistryClient { packages: HashMap, } diff --git a/crates/ruborist/src/util/mod.rs b/crates/ruborist/src/util/mod.rs new file mode 100644 index 0000000000..dfa8bb3234 --- /dev/null +++ b/crates/ruborist/src/util/mod.rs @@ -0,0 +1,5 @@ +//! Utility modules and re-exports. + +pub use crate::model::util::{PackageNameStr, parse_package_spec, read_package_json}; + +pub mod oncemap; diff --git a/crates/ruborist/src/util/oncemap.rs b/crates/ruborist/src/util/oncemap.rs new file mode 100644 index 0000000000..c189f61cc9 --- /dev/null +++ b/crates/ruborist/src/util/oncemap.rs @@ -0,0 +1,465 @@ +#![allow(dead_code)] +//! OnceMap: A concurrent map that ensures each key's work is done exactly once. +//! +//! Based on the pattern from uv package manager: +//! +//! +//! When multiple tasks request the same resource concurrently: +//! - First caller executes the work +//! - Other callers wait and receive the shared result +//! - No duplicate work is performed +//! +//! # Typical Use Case in Package Manager +//! +//! ```text +//! Dependency Graph +//! +//! ┌─────────┐ +//! │ app │ +//! └────┬────┘ +//! ┌──────┼──────┐ +//! ▼ ▼ ▼ +//! ┌───────┐ ┌───┐ ┌───────┐ +//! │ lib-a │ │...│ │ lib-z │ +//! └───┬───┘ └───┘ └───┬───┘ +//! │ │ +//! └───────┬───────┘ +//! ▼ +//! ┌─────────┐ +//! │ react │ ◄── requested by multiple deps +//! └─────────┘ +//! +//! Without OnceMap: +//! lib-a fetches react ──► network request +//! lib-z fetches react ──► network request (duplicate!) +//! +//! With OnceMap: +//! lib-a fetches react ──► network request +//! lib-z waits on lib-a ──► shares result (no duplicate!) +//! ``` + +use dashmap::{DashMap, mapref::entry::Entry}; +use std::future::Future; +use std::hash::Hash; +use std::sync::Arc; +use tokio::sync::Notify; + +/// The state of a value in the OnceMap. +enum Value { + /// Work is in progress, waiters can subscribe to the notify. + Waiting(Arc), + /// Work is complete, result is available. + Done(Arc), +} + +/// A concurrent map that ensures each key's work is done exactly once. +/// +/// # Example +/// ```ignore +/// let map: OnceMap> = OnceMap::new(); +/// +/// // Multiple tasks can call get_or_init concurrently +/// // Only one will actually fetch, others will wait +/// let result = map.get_or_init("react", || async { +/// fetch_package("react").await +/// }).await; +/// ``` +pub struct OnceMap { + map: DashMap>, + /// Waiters for keys that don't exist yet. + /// When get_or_init/complete finishes a key, it notifies and removes the entry. + waiters: DashMap>, +} + +impl Default for OnceMap +where + K: Eq + Hash + Clone, +{ + fn default() -> Self { + Self::new() + } +} + +impl OnceMap +where + K: Eq + Hash + Clone, +{ + /// Create a new empty OnceMap. + pub fn new() -> Self { + Self { + map: DashMap::new(), + waiters: DashMap::new(), + } + } + + /// Register a key as pending (Waiting state) without starting work. + /// Returns the Notify handle if newly registered, None if already exists. + pub fn register(&self, key: K) -> Option> { + use dashmap::mapref::entry::Entry; + match self.map.entry(key) { + Entry::Occupied(_) => None, + Entry::Vacant(vacant) => { + let notify = Arc::new(Notify::new()); + vacant.insert(Value::Waiting(Arc::clone(¬ify))); + Some(notify) + } + } + } + + /// Wait for a key to complete. + /// + /// - If the key exists and is Done, returns immediately. + /// - If the key exists and is Waiting, waits for it to complete. + /// - If the key doesn't exist yet, waits for it to be created and completed. + /// + /// This enables parent-ordering in the clone pipeline: a child package + /// can wait for its parent's clone to finish even if the parent's + /// `get_or_init` hasn't been called yet. + pub async fn wait_if_pending(&self, key: &K) { + // Fast path: key already exists + if let Some(entry) = self.map.get(key) { + match entry.value() { + Value::Done(_) => return, + Value::Waiting(notify) => { + let notify = Arc::clone(notify); + let notified = notify.notified(); + drop(entry); + + if let Some(entry) = self.map.get(key) + && matches!(entry.value(), Value::Done(_)) + { + return; + } + + notified.await; + return; + } + } + } + + // Key doesn't exist yet — register a creation waiter. + // get_or_init/complete will notify this when the key reaches Done (or fails). + let creation_notify = self + .waiters + .entry(key.clone()) + .or_insert_with(|| Arc::new(Notify::new())) + .clone(); + + // Register notified BEFORE double-checking to prevent missed notifications. + let notified = creation_notify.notified(); + + // Double-check: key might have been created and completed between + // the initial map.get and registering here. + if let Some(entry) = self.map.get(key) + && matches!(entry.value(), Value::Done(_)) + { + return; + } + + // Wait for the key to be created and completed + notified.await; + } + + /// Complete a pre-registered key with a value. + /// + /// This is used in conjunction with `register()` for the pre-registration pattern: + /// 1. Call `register(key)` synchronously to claim the key + /// 2. Do async work + /// 3. Call `complete(key, value, notify)` to store result and notify waiters + /// + /// The notify handle must be the one returned from `register()`. + pub fn complete(&self, key: K, value: Option, notify: Arc) { + match value { + Some(v) => { + self.map.insert(key.clone(), Value::Done(Arc::new(v))); + } + None => { + self.map.remove(&key); + } + } + self.notify_all(&key, ¬ify); + } + + /// Notify both the work-notify and any creation waiters for a key. + fn notify_all(&self, key: &K, notify: &Notify) { + notify.notify_waiters(); + if let Some((_, cn)) = self.waiters.remove(key) { + cn.notify_waiters(); + } + } + + /// Get or initialize a value for the given key. + /// + /// If the key doesn't exist, the provided async closure will be called to compute the value. + /// If another task is already computing the value, this will wait for it to complete. + /// + /// Returns `Some(Arc)` if the computation succeeded, `None` if it failed. + pub async fn get_or_init(&self, key: K, init: F) -> Option> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // Fast path: check if already done + if let Some(entry) = self.map.get(&key) + && let Value::Done(result) = entry.value() + { + return Some(Arc::clone(result)); + } + + // Try to register as the worker + let notify = Arc::new(Notify::new()); + let entry = self.map.entry(key.clone()); + + match entry { + Entry::Occupied(occupied) => { + // Someone else is working on it or it's done + match occupied.get() { + Value::Done(result) => { + return Some(Arc::clone(result)); + } + Value::Waiting(existing_notify) => { + let existing_notify = Arc::clone(existing_notify); + + // Register the waiter BEFORE releasing the lock. + // This prevents a race condition where the worker completes + // and calls notify_waiters() between dropping the lock and + // registering the waiter - which would cause us to miss the + // notification and wait forever. + let notified = existing_notify.notified(); + + // Now safe to release the lock + drop(occupied); + + // Double-check: the value might have been inserted between + // releasing the lock and here. If so, return it directly. + if let Some(entry) = self.map.get(&key) + && let Value::Done(result) = entry.value() + { + return Some(Arc::clone(result)); + } + + // Wait for the worker to complete + notified.await; + + // Check the result + if let Some(entry) = self.map.get(&key) + && let Value::Done(result) = entry.value() + { + return Some(Arc::clone(result)); + } + return None; + } + } + } + Entry::Vacant(vacant) => { + // We are the worker + vacant.insert(Value::Waiting(Arc::clone(¬ify))); + } + } + + // Execute the work + let result = init().await; + + // Update the map with the result + let arc_value = result.map(|v| { + let arc = Arc::new(v); + self.map.insert(key.clone(), Value::Done(Arc::clone(&arc))); + arc + }); + if arc_value.is_none() { + self.map.remove(&key); + } + self.notify_all(&key, ¬ify); + arc_value + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + #[tokio::test] + async fn test_single_execution() { + let map: OnceMap = OnceMap::new(); + let call_count = Arc::new(AtomicUsize::new(0)); + + let call_count_clone = Arc::clone(&call_count); + let result = map + .get_or_init("key".to_string(), || async move { + call_count_clone.fetch_add(1, Ordering::SeqCst); + Some(42) + }) + .await; + + assert_eq!(*result.unwrap(), 42); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + // Second call should return cached result + let call_count_clone = Arc::clone(&call_count); + let result = map + .get_or_init("key".to_string(), || async move { + call_count_clone.fetch_add(1, Ordering::SeqCst); + Some(100) + }) + .await; + + assert_eq!(*result.unwrap(), 42); // Still 42, not 100 + assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, not called again + } + + #[tokio::test] + async fn test_concurrent_requests() { + let map = Arc::new(OnceMap::::new()); + let call_count = Arc::new(AtomicUsize::new(0)); + + let mut handles = vec![]; + + // Spawn 10 concurrent tasks requesting the same key + for _ in 0..10 { + let map = Arc::clone(&map); + let call_count = Arc::clone(&call_count); + + handles.push(tokio::spawn(async move { + map.get_or_init("shared_key".to_string(), || async move { + // Simulate slow work + tokio::time::sleep(Duration::from_millis(50)).await; + call_count.fetch_add(1, Ordering::SeqCst); + Some(42) + }) + .await + })); + } + + // Wait for all tasks + let results: Vec<_> = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + + // All should get the same result + for result in results { + assert_eq!(*result.unwrap(), 42); + } + + // Work should only be done once + assert_eq!(call_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_failed_work_allows_retry() { + let map: OnceMap = OnceMap::new(); + let attempt = Arc::new(AtomicUsize::new(0)); + + // First attempt fails + let attempt_clone = Arc::clone(&attempt); + let result = map + .get_or_init("key".to_string(), || async move { + attempt_clone.fetch_add(1, Ordering::SeqCst); + None // Fail + }) + .await; + assert!(result.is_none()); + + // Second attempt should be allowed and succeed + let attempt_clone = Arc::clone(&attempt); + let result = map + .get_or_init("key".to_string(), || async move { + attempt_clone.fetch_add(1, Ordering::SeqCst); + Some(42) + }) + .await; + assert_eq!(*result.unwrap(), 42); + assert_eq!(attempt.load(Ordering::SeqCst), 2); + } + + /// Test that waiters don't miss notifications due to race conditions. + /// + /// This test verifies the fix for a race condition where a waiter could + /// miss the notification if the worker completed between the waiter + /// releasing the lock and registering for notifications. + /// + /// Timeline of the race condition (before fix): + /// ```text + /// Worker Waiter + /// ────── ────── + /// 1. register key + /// insert Waiting(notify) + /// start work... + /// 2. sees Waiting state + /// clone notify + /// drop(lock) ← release lock + /// + /// ┌─────────────────┐ + /// │ DANGER WINDOW │ + /// │ (not listening │ + /// │ for notify) │ + /// └─────────────────┘ + /// + /// 3. work done! + /// insert Done(value) + /// notify_waiters() ← sent! + /// (but no one listening) + /// + /// 4. notified = notify.notified() + /// ← too late! already notified + /// + /// 5. notified.await + /// ← waits forever, deadlock! + /// ``` + /// + /// The fix: register `notified()` BEFORE releasing the lock, then + /// double-check if the value was inserted. + #[tokio::test] + async fn test_no_missed_notifications() { + use tokio::sync::Barrier; + + let map = Arc::new(OnceMap::::new()); + let barrier = Arc::new(Barrier::new(2)); + + // Spawn the worker task + let map_clone = Arc::clone(&map); + let barrier_clone = Arc::clone(&barrier); + let worker = tokio::spawn(async move { + map_clone + .get_or_init("key".to_string(), || async move { + // Wait for waiter to be ready + barrier_clone.wait().await; + // Small delay to let waiter release lock + tokio::time::sleep(Duration::from_millis(5)).await; + Some(42) + }) + .await + }); + + // Spawn waiter task + let map_clone = Arc::clone(&map); + let barrier_clone = Arc::clone(&barrier); + let waiter = tokio::spawn(async move { + // Small delay to ensure worker registers first + tokio::time::sleep(Duration::from_millis(1)).await; + + // Signal worker to proceed + barrier_clone.wait().await; + + // This should not hang - if it does, the race condition exists + map_clone + .get_or_init("key".to_string(), || async move { + panic!("Waiter should not execute work"); + }) + .await + }); + + // Use timeout to detect deadlock + let timeout = Duration::from_secs(2); + let results = tokio::time::timeout(timeout, futures::future::join(worker, waiter)) + .await + .expect("Test timed out - possible deadlock due to missed notification"); + + // Both should get the same result + assert_eq!(*results.0.unwrap().unwrap(), 42); + assert_eq!(*results.1.unwrap().unwrap(), 42); + } +}