From 0b3d91c73ad3b22345cb23b532952875a5aa8bfe Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 22 Feb 2025 12:12:17 -1000
Subject: [PATCH] Respect in-flight requests in source tree and lookahead
 resolver (#11711)

## Summary

Likely not critical since these tend to run prior to resolution rather
than in parallel with it, but we _should_ respect in-flight requests
here.
---
 crates/uv-requirements/src/lookahead.rs   | 28 ++++++++----------
 crates/uv-requirements/src/source_tree.rs | 36 ++++++++++-------------
 2 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/crates/uv-requirements/src/lookahead.rs b/crates/uv-requirements/src/lookahead.rs
index 3f191ed17..c93a87b1d 100644
--- a/crates/uv-requirements/src/lookahead.rs
+++ b/crates/uv-requirements/src/lookahead.rs
@@ -148,22 +148,7 @@ impl<'a, Context: BuildContext> LookaheadResolver<'a, Context> {
         // Fetch the metadata for the distribution.
         let metadata = {
             let id = dist.version_id();
-            if let Some(archive) =
-                self.index
-                    .distributions()
-                    .get(&id)
-                    .as_deref()
-                    .and_then(|response| {
-                        if let MetadataResponse::Found(archive, ..) = response {
-                            Some(archive)
-                        } else {
-                            None
-                        }
-                    })
-            {
-                // If the metadata is already in the index, return it.
-                archive.metadata.clone()
-            } else {
+            if self.index.distributions().register(id.clone()) {
                 // Run the PEP 517 build process to extract metadata from the source distribution.
                 let archive = self
                     .database
@@ -179,6 +164,17 @@ impl<'a, Context: BuildContext> LookaheadResolver<'a, Context> {
                     .done(id, Arc::new(MetadataResponse::Found(archive)));
 
                 metadata
+            } else {
+                let response = self
+                    .index
+                    .distributions()
+                    .wait(&id)
+                    .await
+                    .expect("missing value for registered task");
+                let MetadataResponse::Found(archive) = &*response else {
+                    panic!("Failed to find metadata for: {requirement}");
+                };
+                archive.metadata.clone()
             }
         };
 
diff --git a/crates/uv-requirements/src/source_tree.rs b/crates/uv-requirements/src/source_tree.rs
index 4bdeb86ba..b97cfc32e 100644
--- a/crates/uv-requirements/src/source_tree.rs
+++ b/crates/uv-requirements/src/source_tree.rs
@@ -195,36 +195,32 @@ impl<'a, Context: BuildContext> SourceTreeResolver<'a, Context> {
         };
 
         // Fetch the metadata for the distribution.
-        //
-        // TODO(charlie): Respect in-flight fetches.
         let metadata = {
             let id = VersionId::from_url(source.url());
-            if let Some(archive) =
-                self.index
-                    .distributions()
-                    .get(&id)
-                    .as_deref()
-                    .and_then(|response| {
-                        if let MetadataResponse::Found(archive) = response {
-                            Some(archive)
-                        } else {
-                            None
-                        }
-                    })
-            {
-                // If the metadata is already in the index, return it.
-                archive.metadata.clone()
-            } else {
+            if self.index.distributions().register(id.clone()) {
                 // Run the PEP 517 build process to extract metadata from the source distribution.
                 let source = BuildableSource::Url(source);
                 let archive = self.database.build_wheel_metadata(&source, hashes).await?;
 
+                let metadata = archive.metadata.clone();
+
                 // Insert the metadata into the index.
                 self.index
                     .distributions()
-                    .done(id, Arc::new(MetadataResponse::Found(archive.clone())));
+                    .done(id, Arc::new(MetadataResponse::Found(archive)));
 
-                archive.metadata
+                metadata
+            } else {
+                let response = self
+                    .index
+                    .distributions()
+                    .wait(&id)
+                    .await
+                    .expect("missing value for registered task");
+                let MetadataResponse::Found(archive) = &*response else {
+                    panic!("Failed to find metadata for: {}", path.user_display());
+                };
+                archive.metadata.clone()
             }
         };