From a07adf72dea1bf0f826c315dec569e22c4366c02 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Wed, 24 Apr 2024 10:17:16 -0500
Subject: [PATCH] Use read-write locks instead of mutexes in authentication
 handling (#3210)

- Use `RwLock` for `KeyringProvider` cache
- Use `RwLock` for `CredentialsCache`
---
 crates/uv-auth/src/cache.rs   | 18 +++++++++---------
 crates/uv-auth/src/keyring.rs | 21 ++++++++++-----------
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/crates/uv-auth/src/cache.rs b/crates/uv-auth/src/cache.rs
index 758d058a9051..386e05b68126 100644
--- a/crates/uv-auth/src/cache.rs
+++ b/crates/uv-auth/src/cache.rs
@@ -1,5 +1,5 @@
 use std::sync::Arc;
-use std::{collections::HashMap, sync::Mutex};
+use std::{collections::HashMap, sync::RwLock};
 
 use crate::credentials::{Credentials, Username};
 use crate::Realm;
@@ -10,11 +10,11 @@ use url::Url;
 
 pub struct CredentialsCache {
     /// A cache per realm and username
-    realms: Mutex<HashMap<(Realm, Username), Arc<Credentials>>>,
+    realms: RwLock<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>,
+    urls: RwLock<UrlTrie>,
 }
 
 impl Default for CredentialsCache {
@@ -27,15 +27,15 @@ impl CredentialsCache {
     /// Create a new cache.
     pub fn new() -> Self {
         Self {
-            realms: Mutex::new(HashMap::new()),
             fetches: OnceMap::default(),
-            urls: Mutex::new(UrlTrie::new()),
+            realms: RwLock::new(HashMap::new()),
+            urls: RwLock::new(UrlTrie::new()),
         }
     }
 
     /// Return the credentials that should be used for a realm and username, if any.
     pub(crate) fn get_realm(&self, realm: Realm, username: Username) -> Option<Arc<Credentials>> {
-        let realms = self.realms.lock().unwrap();
+        let realms = self.realms.read().unwrap();
         let name = if let Some(username) = username.as_deref() {
             format!("{username}@{realm}")
         } else {
@@ -65,7 +65,7 @@ impl CredentialsCache {
     /// cached credentials have a username equal to the provided one — otherwise `None` is returned.
     /// If multiple usernames are used per URL, the realm cache should be queried instead.
     pub(crate) fn get_url(&self, url: &Url, username: Username) -> Option<Arc<Credentials>> {
-        let urls = self.urls.lock().unwrap();
+        let urls = self.urls.read().unwrap();
         let credentials = urls.get(url);
         if let Some(credentials) = credentials {
             if username.is_none() || username.as_deref() == credentials.username() {
@@ -100,7 +100,7 @@ impl CredentialsCache {
         self.insert_realm((Realm::from(url), Username::none()), credentials.clone());
 
         // Insert an entry for the URL
-        let mut urls = self.urls.lock().unwrap();
+        let mut urls = self.urls.write().unwrap();
         urls.insert(url.clone(), credentials.clone());
     }
 
@@ -117,7 +117,7 @@ impl CredentialsCache {
             return None;
         }
 
-        let mut realms = self.realms.lock().unwrap();
+        let mut realms = self.realms.write().unwrap();
 
         // Always replace existing entries if we have a password
         if credentials.password().is_some() {
diff --git a/crates/uv-auth/src/keyring.rs b/crates/uv-auth/src/keyring.rs
index 773224ff1bc2..bdcc76888245 100644
--- a/crates/uv-auth/src/keyring.rs
+++ b/crates/uv-auth/src/keyring.rs
@@ -1,4 +1,5 @@
-use std::{collections::HashSet, sync::Mutex};
+use std::collections::HashSet;
+use std::sync::RwLock;
 
 use tokio::process::Command;
 use tracing::{debug, instrument, warn};
@@ -13,7 +14,7 @@ use crate::credentials::Credentials;
 #[derive(Debug)]
 pub struct KeyringProvider {
     /// Tracks host and username pairs with no credentials to avoid repeated queries.
-    cache: Mutex<HashSet<(String, String)>>,
+    cache: RwLock<HashSet<(String, String)>>,
     backend: KeyringProviderBackend,
 }
 
@@ -29,7 +30,7 @@ impl KeyringProvider {
     /// Create a new [`KeyringProvider::Subprocess`].
     pub fn subprocess() -> Self {
         Self {
-            cache: Mutex::new(HashSet::new()),
+            cache: RwLock::new(HashSet::new()),
             backend: KeyringProviderBackend::Subprocess,
         }
     }
@@ -65,7 +66,7 @@ impl KeyringProvider {
         //      use-cases.
         let key = (host.to_string(), username.to_string());
         {
-            let cache = self.cache.lock().unwrap();
+            let cache = self.cache.read().unwrap();
 
             if cache.contains(&key) {
                 debug!(
@@ -95,11 +96,9 @@ impl KeyringProvider {
             };
         }
 
-        {
-            let mut cache = self.cache.lock().unwrap();
-            if password.is_none() {
-                cache.insert(key);
-            }
+        if password.is_none() {
+            let mut cache = self.cache.write().unwrap();
+            cache.insert(key);
         }
 
         password.map(|password| Credentials::new(Some(username.to_string()), Some(password)))
@@ -148,7 +147,7 @@ impl KeyringProvider {
         use std::collections::HashMap;
 
         Self {
-            cache: Mutex::new(HashSet::new()),
+            cache: RwLock::new(HashSet::new()),
             backend: KeyringProviderBackend::Dummy(HashMap::from_iter(
                 iter.into_iter()
                     .map(|((service, username), password)| ((service.into(), username), password)),
@@ -162,7 +161,7 @@ impl KeyringProvider {
         use std::collections::HashMap;
 
         Self {
-            cache: Mutex::new(HashSet::new()),
+            cache: RwLock::new(HashSet::new()),
             backend: KeyringProviderBackend::Dummy(HashMap::new()),
         }
     }