From e92b38cfb964222d976ee3ff2e70b62f8c1c02a9 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Wed, 24 Apr 2024 09:53:44 -0500
Subject: [PATCH] Only perform fetches of credentials for a realm once (#3237)

Closes https://github.com/astral-sh/uv/issues/3205

Tested with

`RUST_LOG=uv=trace cargo run -- pip install -r
scripts/requirements/trio.in --index-url
https://oauth2accesstoken@us-central1-python.pkg.dev/zb-test-project-421213/pypyi/simple/
--no-cache --keyring-provider subprocess -vv --reinstall 2>&1 | grep
keyring`

On `main` you can see a dozen keyring attempts at once. Here, the other
requests wait for the first attempt and only a single keyring call is
performed.
---
 Cargo.lock                       |  1 +
 crates/uv-auth/Cargo.toml        |  3 ++-
 crates/uv-auth/src/cache.rs      |  4 +++
 crates/uv-auth/src/middleware.rs | 42 +++++++++++++++++++++++++++++---
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5ab9e8b8d..2d576f94b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4492,6 +4492,7 @@ dependencies = [
  "futures",
  "http",
  "insta",
+ "once-map",
  "once_cell",
  "reqwest",
  "reqwest-middleware",
diff --git a/crates/uv-auth/Cargo.toml b/crates/uv-auth/Cargo.toml
index a98255ee5..d60f6f151 100644
--- a/crates/uv-auth/Cargo.toml
+++ b/crates/uv-auth/Cargo.toml
@@ -10,11 +10,12 @@ base64 = { workspace = true }
 futures = { workspace = true }
 http = { workspace = true }
 once_cell = { workspace = true }
+once-map = { workspace = true }
 reqwest = { workspace = true }
 reqwest-middleware = { workspace = true }
 rust-netrc = { workspace = true }
-tracing = { workspace = true }
 tokio = { workspace = true }
+tracing = { workspace = true }
 url = { workspace = true }
 urlencoding = { workspace = true }
 
diff --git a/crates/uv-auth/src/cache.rs b/crates/uv-auth/src/cache.rs
index fe3f0a0d2..758d058a9 100644
--- a/crates/uv-auth/src/cache.rs
+++ b/crates/uv-auth/src/cache.rs
@@ -4,12 +4,15 @@ use std::{collections::HashMap, sync::Mutex};
 use crate::credentials::{Credentials, Username};
 use crate::Realm;
 
+use once_map::OnceMap;
 use tracing::trace;
 use url::Url;
 
 pub struct CredentialsCache {
     /// A cache per realm and username
     realms: Mutex<HashMap<(Realm, Username), Arc<Credentials>>>,
+    /// A cache tracking the result of fetches from external services
+    pub(crate) fetches: OnceMap<(Realm, Username), Option<Arc<Credentials>>>,
     /// A cache per URL, uses a trie for efficient prefix queries.
     urls: Mutex<UrlTrie>,
 }
@@ -25,6 +28,7 @@ impl CredentialsCache {
     pub fn new() -> Self {
         Self {
             realms: Mutex::new(HashMap::new()),
+            fetches: OnceMap::default(),
             urls: Mutex::new(UrlTrie::new()),
         }
     }
diff --git a/crates/uv-auth/src/middleware.rs b/crates/uv-auth/src/middleware.rs
index 92b11d985..facc0dd41 100644
--- a/crates/uv-auth/src/middleware.rs
+++ b/crates/uv-auth/src/middleware.rs
@@ -165,7 +165,7 @@ impl Middleware for AuthMiddleware {
                 .await
             {
                 request = credentials.authenticate(request);
-                Some(Arc::new(credentials))
+                Some(credentials)
             } else {
                 // If we don't find a password, we'll still attempt the request with the existing credentials
                 Some(credentials)
@@ -244,7 +244,7 @@ impl Middleware for AuthMiddleware {
             retry_request = credentials.authenticate(retry_request);
             trace!("Retrying request for {url} with {credentials:?}");
             return self
-                .complete_request(Some(Arc::new(credentials)), retry_request, extensions, next)
+                .complete_request(Some(credentials), retry_request, extensions, next)
                 .await;
         }
 
@@ -300,9 +300,37 @@ impl AuthMiddleware {
         &self,
         credentials: Option<&Credentials>,
         url: &Url,
-    ) -> Option<Credentials> {
+    ) -> Option<Arc<Credentials>> {
+        // Fetches can be expensive, so we will only run them _once_ per realm and username combination
+        // All other requests for the same realm will wait until the first one completes
+        let key = (
+            Realm::from(url),
+            Username::from(
+                credentials
+                    .map(|credentials| credentials.username().unwrap_or_default().to_string()),
+            ),
+        );
+
+        if !self.cache().fetches.register(key.clone()) {
+            let credentials = Arc::<_>::unwrap_or_clone(
+                self.cache()
+                    .fetches
+                    .wait(&key)
+                    .await
+                    .expect("The key must exist after register is called"),
+            );
+
+            if credentials.is_some() {
+                trace!("Using credentials from previous fetch for {url}");
+            } else {
+                trace!("Skipping fetch of credentails for {url}, previous attempt failed");
+            };
+
+            return credentials;
+        }
+
         // Netrc support based on: <https://github.com/gribouille/netrc>.
-        if let Some(credentials) = self.netrc.as_ref().and_then(|netrc| {
+        let credentials = if let Some(credentials) = self.netrc.as_ref().and_then(|netrc| {
             trace!("Checking netrc for credentials for {url}");
             Credentials::from_netrc(
                 netrc,
@@ -336,6 +364,12 @@ impl AuthMiddleware {
         } else {
             None
         }
+        .map(Arc::new);
+
+        // Register the fetch for this key
+        self.cache().fetches.done(key.clone(), credentials.clone());
+
+        credentials
     }
 }