From ac348eecdfca2af2ec3cab0056faf7991c1f8ae8 Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Thu, 19 Dec 2024 15:47:01 +0100
Subject: [PATCH] Batch prefetch per fork (#10029)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Previously, the batch prefetcher was part of the solver loop, used
across forks. This would lead to each preference in a fork being counted
as a tried version, so that after 5 forks with the identical version, we
would start batch prefetching. The reported numbers of tried versions
are also reported. By tracking the batch prefetcher on the fork the
numbers are corrected.

An alternative would be tracking the actually tried versions, but that
would mean more overhead in the top level solver loop when the current
heuristic works.

In `ecosystem/transformers`:

```
$ hyperfine --runs 10 --prepare "rm -f uv.lock" "../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z" "uv lock --exclude-newer 2024-08-08T00:00:00Z"
Benchmark 1: ../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z
  Time (mean ± σ):     386.2 ms ±   6.1 ms    [User: 396.0 ms, System: 144.5 ms]
  Range (min … max):   378.5 ms … 397.9 ms    10 runs

Benchmark 2: uv lock --exclude-newer 2024-08-08T00:00:00Z
  Time (mean ± σ):     422.0 ms ±   5.5 ms    [User: 459.6 ms, System: 190.3 ms]
  Range (min … max):   415.0 ms … 430.5 ms    10 runs

Summary
  ../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z ran
    1.09 ± 0.02 times faster than uv lock --exclude-newer 2024-08-08T00:00:00Z
```
---
 .../src/resolver/batch_prefetch.rs            |  1 +
 crates/uv-resolver/src/resolver/mod.rs        | 21 ++++++++++++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs
index f1f9e7a19..b8e6e9533 100644
--- a/crates/uv-resolver/src/resolver/batch_prefetch.rs
+++ b/crates/uv-resolver/src/resolver/batch_prefetch.rs
@@ -39,6 +39,7 @@ enum BatchPrefetchStrategy {
 /// have to fetch the metadata for a lot of versions.
 ///
 /// Note that these all heuristics that could totally prefetch lots of irrelevant versions.
+#[derive(Clone)]
 pub(crate) struct BatchPrefetcher {
     // Internal types.
     tried_versions: FxHashMap<PackageName, usize>,
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
index 1bbe9f26b..57e8cffa3 100644
--- a/crates/uv-resolver/src/resolver/mod.rs
+++ b/crates/uv-resolver/src/resolver/mod.rs
@@ -310,12 +310,17 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
 
         let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone()));
         let pubgrub = State::init(root.clone(), MIN_VERSION.clone());
-        let mut prefetcher = BatchPrefetcher::new(
+        let prefetcher = BatchPrefetcher::new(
             self.capabilities.clone(),
             self.index.clone(),
             request_sink.clone(),
         );
-        let state = ForkState::new(pubgrub, self.env.clone(), self.python_requirement.clone());
+        let state = ForkState::new(
+            pubgrub,
+            self.env.clone(),
+            self.python_requirement.clone(),
+            prefetcher,
+        );
         let mut preferences = self.preferences.clone();
         let mut forked_states = self.env.initial_forked_states(state);
         let mut resolutions = vec![];
@@ -393,7 +398,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                         else {
                             // All packages have been assigned, the fork has been successfully resolved
                             if tracing::enabled!(Level::DEBUG) {
-                                prefetcher.log_tried_versions();
+                                state.prefetcher.log_tried_versions();
                             }
                             debug!(
                                 "{} resolution took {:.3}s",
@@ -472,7 +477,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 // (idempotent due to caching).
                 self.request_package(next_package, url, index, &request_sink)?;
 
-                prefetcher.version_tried(next_package);
+                state.prefetcher.version_tried(next_package);
 
                 let term_intersection = state
                     .pubgrub
@@ -545,7 +550,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
 
                 // Only consider registry packages for prefetch.
                 if url.is_none() {
-                    prefetcher.prefetch_batches(
+                    state.prefetcher.prefetch_batches(
                         next_package,
                         index,
                         &version,
@@ -2259,6 +2264,10 @@ pub(crate) struct ForkState {
     /// solution that omits Python 3.8 support.
     python_requirement: PythonRequirement,
     conflict_tracker: ConflictTracker,
+    /// Prefetch package versions for packages with many rejected versions.
+    ///
+    /// Tracked on the fork state to avoid counting each identical version between forks as new try.
+    prefetcher: BatchPrefetcher,
 }
 
 impl ForkState {
@@ -2266,6 +2275,7 @@ impl ForkState {
         pubgrub: State<UvDependencyProvider>,
         env: ResolverEnvironment,
         python_requirement: PythonRequirement,
+        prefetcher: BatchPrefetcher,
     ) -> Self {
         Self {
             initial: None,
@@ -2279,6 +2289,7 @@ impl ForkState {
             env,
             python_requirement,
             conflict_tracker: ConflictTracker::default(),
+            prefetcher,
         }
     }