fix: simplify RwLock<Arc<HashMap<...>> as ArcSwap<HashMap<..., ahash:…

…RandomState>>

CHANGELOG.md

@@ -17,7 +17,7 @@ You should also include the user name that made the change.
 
 ## 0.4.2
 
-- Feat: Change the default hasher for hashmaps and hashsets from `FxHash` to `aHash` for better performance with string keys.
+- Feat: Change the default hasher for hashmaps and hashsets from `FxHash` to `aHash` for better performance with string keys. Use `ArcSwap` instead of `RwLock` for internal ODoH config storage.
 - Deps.
 - Refactor: Various minor improvements.
 

proxy-lib/src/doh_client/doh_client_main.rs

@@ -325,9 +325,6 @@ impl DoHClient {
     let Some(odoh_config) = self.odoh_configs.as_ref().unwrap().get(target_obj).await else {
       return Err(DohClientError::ODoHNoClientConfig);
     };
-    let Some(odoh_config) = odoh_config.as_ref() else {
-      return Err(DohClientError::ODoHNoClientConfig);
-    };
 
     // encrypt query
     let (odoh_plaintext_query, encrypted_query_body, secret) = odoh_config.encrypt_query(packet_buf)?;

proxy-lib/src/doh_client/odoh_config_store.rs

@@ -17,17 +17,14 @@ use url::Url;
 /// ODoH config store
 pub struct ODoHConfigStore {
   // inner: Arc<RwLock<HashMap<Arc<DoHTarget>, Arc<Option<ODoHConfig>>>>>,
-  inner: ArcSwap<HashMap<Arc<DoHTarget>, Arc<Option<ODoHConfig>>>>,
+  inner: ArcSwap<HashMap<Arc<DoHTarget>, Option<Arc<ODoHConfig>>>>,
   http_client: Arc<RwLock<HttpClientInner>>,
 }
 
 impl ODoHConfigStore {
   /// Create a new ODoHConfigStore
   pub async fn new(http_client: Arc<RwLock<HttpClientInner>>, targets: &[Arc<DoHTarget>]) -> Result<Self, DohClientError> {
-    let inner = targets
-      .iter()
-      .map(|target| (target.clone(), Arc::new(None as Option<ODoHConfig>)))
-      .collect::<HashMap<_, _>>();
+    let inner = targets.iter().map(|target| (target.clone(), None)).collect::<HashMap<_, _>>();
     let res = Self {
       inner: ArcSwap::new(Arc::new(inner)),
       http_client,
@@ -37,45 +34,46 @@ impl ODoHConfigStore {
   }
 
   /// Get a ODoHConfig for DoHTarget
-  pub async fn get(&self, target: &Arc<DoHTarget>) -> Option<Arc<Option<ODoHConfig>>> {
-    self.inner.load().get(target).cloned()
+  pub async fn get(&self, target: &Arc<DoHTarget>) -> Option<Arc<ODoHConfig>> {
+    self.inner.load().get(target).cloned().unwrap_or(None)
   }
 
   /// Fetch ODoHConfig from target
   pub async fn update_odoh_config_from_well_known(&self) -> Result<(), DohClientError> {
     // TODO: Add auth token when fetching config?
     // fetch public key from odoh target (/.well-known)
-    let inner = self.inner.load().clone();
+    let inner = self.inner.load();
 
-    let futures = inner.as_ref().keys().map(|target| async {
+    let futures = inner.keys().map(|target| async {
       let mut destination = Url::parse(&format!("{}://{}", target.scheme(), target.authority())).unwrap();
       destination.set_path(ODOH_CONFIG_PATH);
       let lock = self.http_client.read().await;
       debug!("Fetching ODoH config from {}", destination);
-      lock
+      let res = lock
         .get(destination)
         .header(reqwest::header::ACCEPT, "application/binary")
         .send()
-        .await
+        .await;
+      (target.clone(), res)
     });
     let joined = futures::future::join_all(futures);
-    let update_futures = joined.await.into_iter().zip(inner.iter()).map(|(res, current)| async move {
+    let update_futures = joined.await.into_iter().map(|(target, res)| async move {
       match res {
         Ok(response) => {
           if response.status() != reqwest::StatusCode::OK {
             error!("Failed to fetch ODoH config!: {:?}", response.status());
-            return (current.0.clone(), Arc::new(None as Option<ODoHConfig>));
+            return (target.clone(), None);
           }
           let Ok(body) = response.bytes().await else {
             error!("Failed to parse response body in ODoH config response");
-            return (current.0.clone(), Arc::new(None as Option<ODoHConfig>));
+            return (target.clone(), None);
           };
-          let config = ODoHConfig::new(current.0.authority(), &body).ok();
-          (current.0.clone(), Arc::new(config))
+          let config = ODoHConfig::new(target.authority(), &body).ok();
+          (target.clone(), config.map(Arc::new))
         }
         Err(e) => {
           error!("Failed to fetch ODoH config!: {:?}", e);
-          (current.0.clone(), Arc::new(None as Option<ODoHConfig>))
+          (target.clone(), None)
         }
       }
     });