From 14a9008e54def81e46955cebeb443c196813be71 Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Tue, 7 Jan 2025 15:12:18 +0100
Subject: [PATCH] Avoid overcounting versions in batch prefetcher (#10350)

---
 .../src/resolver/batch_prefetch.rs            | 24 +++++++++++--------
 crates/uv-resolver/src/resolver/mod.rs        |  4 ++--
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs
index dd63fc60f..6d35a2f97 100644
--- a/crates/uv-resolver/src/resolver/batch_prefetch.rs
+++ b/crates/uv-resolver/src/resolver/batch_prefetch.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 use itertools::Itertools;
 use pubgrub::{Range, Ranges, Term};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 use tokio::sync::mpsc::Sender;
 use tracing::{debug, trace};
 
@@ -43,7 +43,7 @@ enum BatchPrefetchStrategy {
 #[derive(Clone)]
 pub(crate) struct BatchPrefetcher {
     // Types to determine whether we need to prefetch.
-    tried_versions: FxHashMap<PackageName, usize>,
+    tried_versions: FxHashMap<PackageName, FxHashSet<Version>>,
     last_prefetch: FxHashMap<PackageName, usize>,
     // Types to execute the prefetch.
     prefetch_runner: BatchPrefetcherRunner,
@@ -142,7 +142,7 @@ impl BatchPrefetcher {
     }
 
     /// Each time we tried a version for a package, we register that here.
-    pub(crate) fn version_tried(&mut self, package: &PubGrubPackage) {
+    pub(crate) fn version_tried(&mut self, package: &PubGrubPackage, version: &Version) {
         // Only track base packages, no virtual packages from extras.
         let PubGrubPackageInner::Package {
             name,
@@ -153,7 +153,10 @@ impl BatchPrefetcher {
         else {
             return;
         };
-        *self.tried_versions.entry(name.clone()).or_default() += 1;
+        self.tried_versions
+            .entry(name.clone())
+            .or_default()
+            .insert(version.clone());
     }
 
     /// After 5, 10, 20, 40 tried versions, prefetch that many versions to start early but not
@@ -170,7 +173,7 @@ impl BatchPrefetcher {
             return (0, false);
         };
 
-        let num_tried = self.tried_versions.get(name).copied().unwrap_or_default();
+        let num_tried = self.tried_versions.get(name).map_or(0, FxHashSet::len);
         let previous_prefetch = self.last_prefetch.get(name).copied().unwrap_or_default();
         let do_prefetch = (num_tried >= 5 && previous_prefetch < 5)
             || (num_tried >= 10 && previous_prefetch < 10)
@@ -180,12 +183,13 @@ impl BatchPrefetcher {
     }
 
     /// Log stats about how many versions we tried.
-    ///
-    /// Note that they may be inflated when we count the same version repeatedly during
-    /// backtracking.
     pub(crate) fn log_tried_versions(&self) {
-        let total_versions: usize = self.tried_versions.values().sum();
-        let mut tried_versions: Vec<_> = self.tried_versions.iter().collect();
+        let total_versions: usize = self.tried_versions.values().map(FxHashSet::len).sum();
+        let mut tried_versions: Vec<_> = self
+            .tried_versions
+            .iter()
+            .map(|(name, versions)| (name, versions.len()))
+            .collect();
         tried_versions.sort_by(|(p1, c1), (p2, c2)| {
             c1.cmp(c2)
                 .reverse()
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
index dd71b0b0d..2b3e4f696 100644
--- a/crates/uv-resolver/src/resolver/mod.rs
+++ b/crates/uv-resolver/src/resolver/mod.rs
@@ -476,8 +476,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 // (idempotent due to caching).
                 self.request_package(next_package, url, index, &request_sink)?;
 
-                state.prefetcher.version_tried(next_package);
-
                 let version = if let Some(version) = state.initial_version.take() {
                     // If we just forked based on platform support, we can skip version selection,
                     // since the fork operation itself already selected the appropriate version for
@@ -571,6 +569,8 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     version
                 };
 
+                state.prefetcher.version_tried(next_package, &version);
+
                 self.on_progress(next_package, &version);
 
                 if !state