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.
This commit is contained in:
Zanie Blue 2024-04-24 09:53:44 -05:00 committed by GitHub
parent 116d47ed03
commit e92b38cfb9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 45 additions and 5 deletions

1
Cargo.lock generated
View File

@ -4492,6 +4492,7 @@ dependencies = [
"futures", "futures",
"http", "http",
"insta", "insta",
"once-map",
"once_cell", "once_cell",
"reqwest", "reqwest",
"reqwest-middleware", "reqwest-middleware",

View File

@ -10,11 +10,12 @@ base64 = { workspace = true }
futures = { workspace = true } futures = { workspace = true }
http = { workspace = true } http = { workspace = true }
once_cell = { workspace = true } once_cell = { workspace = true }
once-map = { workspace = true }
reqwest = { workspace = true } reqwest = { workspace = true }
reqwest-middleware = { workspace = true } reqwest-middleware = { workspace = true }
rust-netrc = { workspace = true } rust-netrc = { workspace = true }
tracing = { workspace = true }
tokio = { workspace = true } tokio = { workspace = true }
tracing = { workspace = true }
url = { workspace = true } url = { workspace = true }
urlencoding = { workspace = true } urlencoding = { workspace = true }

View File

@ -4,12 +4,15 @@ use std::{collections::HashMap, sync::Mutex};
use crate::credentials::{Credentials, Username}; use crate::credentials::{Credentials, Username};
use crate::Realm; use crate::Realm;
use once_map::OnceMap;
use tracing::trace; use tracing::trace;
use url::Url; use url::Url;
pub struct CredentialsCache { pub struct CredentialsCache {
/// A cache per realm and username /// A cache per realm and username
realms: Mutex<HashMap<(Realm, Username), Arc<Credentials>>>, 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. /// A cache per URL, uses a trie for efficient prefix queries.
urls: Mutex<UrlTrie>, urls: Mutex<UrlTrie>,
} }
@ -25,6 +28,7 @@ impl CredentialsCache {
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
realms: Mutex::new(HashMap::new()), realms: Mutex::new(HashMap::new()),
fetches: OnceMap::default(),
urls: Mutex::new(UrlTrie::new()), urls: Mutex::new(UrlTrie::new()),
} }
} }

View File

@ -165,7 +165,7 @@ impl Middleware for AuthMiddleware {
.await .await
{ {
request = credentials.authenticate(request); request = credentials.authenticate(request);
Some(Arc::new(credentials)) Some(credentials)
} else { } else {
// If we don't find a password, we'll still attempt the request with the existing credentials // If we don't find a password, we'll still attempt the request with the existing credentials
Some(credentials) Some(credentials)
@ -244,7 +244,7 @@ impl Middleware for AuthMiddleware {
retry_request = credentials.authenticate(retry_request); retry_request = credentials.authenticate(retry_request);
trace!("Retrying request for {url} with {credentials:?}"); trace!("Retrying request for {url} with {credentials:?}");
return self return self
.complete_request(Some(Arc::new(credentials)), retry_request, extensions, next) .complete_request(Some(credentials), retry_request, extensions, next)
.await; .await;
} }
@ -300,9 +300,37 @@ impl AuthMiddleware {
&self, &self,
credentials: Option<&Credentials>, credentials: Option<&Credentials>,
url: &Url, 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>. // 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}"); trace!("Checking netrc for credentials for {url}");
Credentials::from_netrc( Credentials::from_netrc(
netrc, netrc,
@ -336,6 +364,12 @@ impl AuthMiddleware {
} else { } else {
None None
} }
.map(Arc::new);
// Register the fetch for this key
self.cache().fetches.done(key.clone(), credentials.clone());
credentials
} }
} }