[uv] OnceMap: Rust Pattern for Running Concurrent Work Exactly Once
Design · How to coordinate parallel tasks and deduplicate expensive operations in async Rust
uv, the Rust-based Python package manager from Astral, is 10-100x faster than pip. As part of the speedup, uv aggressively parallelizes dependency resolution, wheel downloads, and source builds.
When running uv pip install scipy pandas, both packages depend on numpy, causing two threads to independently fetch numpy metadata. Without coordination, this triggers redundant requests. OnceMap ensures only the first caller fetches the data while others wait for the result.
Parallelization creates a coordination challenge. When multiple tasks need the same resource, how do you ensure the work happens exactly once? uv’s solution is OnceMap - a lightweight concurrent memoization primitive that powers deduplication across the resolver and installer.
Core Data Structure
The OnceMap struct wraps DashMap, a concurrent hash map. Unlike a standard HashMap with a single lock, DashMap shards data into segments, each with its own lock, enabling parallel access without contention.
Value enum is a crucial element. An entry tracks state beyond simple presence: in progress, completed, or not started. The Waiting variant holds a Tokio Notify, a lightweight synchronization primitive that lets tasks sleep and wake efficiently.
Three-Method Protocol
OnceMap uses a three-method protocol: register, done, and wait.
Method 1: Register Intent
register returns a boolean contract. If true, you must perform the work and call done. If false, another task is already handling it.
Method 2: Signal Completion
done replaces the Waiting sentinel with the result and wakes all pending tasks using notify_waiters().
Method 3: Wait for Results
Race Condition
This design addresses a race condition discovered during development (Issue #3724).
The original implementation had a flaw:
Task A calls
register()and starts work.Task B calls
register(), sees it’s in progress, and callswait().Task B retrieves the
Notifyhandle.Task A finishes and calls
done()(signaling waiters).Task B calls
notify.notified().await
Notify signals are not queued, so Task B misses the notification and hangs forever.
The fix is to register as a waiter before checking the map again as implemented in wait() above.
pin!(notify.notified()) registers Task B to receive signals immediately. Even if Task A calls notify_waiters() during the second get(key) check, Task B will catch it. The pin! macro keeps the future at a fixed memory address. Rust futures are movable values by default. Since Notified registers its own address in Tokio’s wait queue, moving it afterward would leave a dangling pointer. pin! guarantees the future stays in place.
Note that the race condition fix is entirely in wait().
OnceMap Usage: Download Preparer
OnceMap coordinates wheel downloads and metadata fetching. Here is how the installer uses it.
astral-sh/uv:preparer.rs#L109-L211
One task performs the work while concurrent requests wait for the shared result.
Design Choices
Why not a standard cache? A cache stores results but doesn’t track in-flight work, so it can’t prevent duplicate requests.
Why
DashMap+Notify?DashMapminimizes contention with fine-grained locking;Notifyefficiently handles async synchronization.Why clone on retrieval? Returning references requires holding locks across async boundaries, risking deadlocks.
Arc<V>makes cloning efficient.
Major Contributions
PR #544 Keep track of in flight unzips using
OnceMapby @konstin (konsti)PR #3627 Parallelize resolver by @ibraheemdev (Ibraheem Ahmed)
PR #3987 Avoid race condition in
OnceMapby @ibraheemdev (Ibraheem Ahmed)
Great article and writeup. We will be including this piece in our upcoming newsletter issue.
Thanks for putting this into writing.
Is there a race condition in the register function as well? After getting the current Vacant entry, and before inserting a new Waiting value, could another task also insert its own Waiting value, given that the "check for vacancy then register intent" logic doesn't seem to be atomic?
Is this right, or could someone please explain how that part works? Maybe it's a fail-safe scenario which leads to doing the work twice or more, rather than exactly once?