mirror of https://github.com/astral-sh/uv
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:
parent
116d47ed03
commit
e92b38cfb9
|
|
@ -4492,6 +4492,7 @@ dependencies = [
|
||||||
"futures",
|
"futures",
|
||||||
"http",
|
"http",
|
||||||
"insta",
|
"insta",
|
||||||
|
"once-map",
|
||||||
"once_cell",
|
"once_cell",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
"reqwest-middleware",
|
"reqwest-middleware",
|
||||||
|
|
|
||||||
|
|
@ -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 }
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue