chore: URI caching now avoids hash collisions and reduces lock contention

Signed-off-by: Dmitry Dygalo <[email protected]>

Author: Stranger6667

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Changed
 
+- `referencing`: URI caching now avoids hash collisions and reduces lock contention.
 - Update `fluent-uri` to `0.4.1`.
 - Bump MSRV to `1.83.0`.
 - Drop the `Send + Sync` bounds from `Retrieve`/`AsyncRetrieve` on `wasm32`.

crates/jsonschema-referencing/Cargo.toml

@@ -25,6 +25,7 @@ futures = { version = "0.3.31", optional = true }
 parking_lot = "0.12.3"
 percent-encoding = "2.3.1"
 serde_json.workspace = true
+hashbrown = "0.16"
 
 [dev-dependencies]
 benchmark = { path = "../benchmark/" }

crates/jsonschema-referencing/src/cache.rs

@@ -1,30 +1,29 @@
-use core::hash::{BuildHasherDefault, Hash, Hasher};
-use std::{
-    collections::{hash_map::Entry, HashMap},
-    sync::Arc,
-};
+use std::sync::Arc;
 
-use ahash::AHasher;
 use fluent_uri::Uri;
-use parking_lot::RwLock;
+use hashbrown::hash_map::{EntryRef, HashMap};
+use parking_lot::{RwLock, RwLockUpgradableReadGuard};
 
-use crate::{hasher::BuildNoHashHasher, uri, Error};
+use crate::{uri, Error};
+
+type CacheBucket = HashMap<String, Arc<Uri<String>>>;
+type CacheMap = HashMap<String, CacheBucket>;
 
 #[derive(Debug, Clone)]
 pub(crate) struct UriCache {
-    cache: HashMap<u64, Arc<Uri<String>>, BuildNoHashHasher>,
+    cache: CacheMap,
 }
 
 impl UriCache {
     pub(crate) fn new() -> Self {
         Self {
-            cache: HashMap::with_hasher(BuildHasherDefault::default()),
+            cache: HashMap::new(),
         }
     }
 
     pub(crate) fn with_capacity(capacity: usize) -> Self {
         Self {
-            cache: HashMap::with_capacity_and_hasher(capacity, BuildHasherDefault::default()),
+            cache: HashMap::with_capacity(capacity),
         }
     }
 
@@ -33,17 +32,31 @@ impl UriCache {
         base: &Uri<&str>,
         uri: impl AsRef<str>,
     ) -> Result<Arc<Uri<String>>, Error> {
-        let mut hasher = AHasher::default();
-        (base.as_str(), uri.as_ref()).hash(&mut hasher);
-        let hash = hasher.finish();
-
-        Ok(match self.cache.entry(hash) {
-            Entry::Occupied(entry) => Arc::clone(entry.get()),
-            Entry::Vacant(entry) => {
-                let new = Arc::new(uri::resolve_against(base, uri.as_ref())?);
-                Arc::clone(entry.insert(new))
+        let base_str = base.as_str();
+        let reference = uri.as_ref();
+
+        let resolved = match self.cache.entry_ref(base_str) {
+            EntryRef::Occupied(mut entry) => {
+                if let Some(cached) = entry.get().get(reference) {
+                    return Ok(Arc::clone(cached));
+                }
+
+                let resolved = Arc::new(uri::resolve_against(base, reference)?);
+                entry
+                    .get_mut()
+                    .insert(reference.to_owned(), Arc::clone(&resolved));
+                resolved
             }
-        })
+            EntryRef::Vacant(entry) => {
+                let resolved = Arc::new(uri::resolve_against(base, reference)?);
+                let mut inner = HashMap::with_capacity(1);
+                inner.insert(reference.to_owned(), Arc::clone(&resolved));
+                entry.insert(inner);
+                resolved
+            }
+        };
+
+        Ok(resolved)
     }
 
     pub(crate) fn into_shared(self) -> SharedUriCache {
@@ -56,7 +69,7 @@ impl UriCache {
 /// A dedicated type for URI resolution caching.
 #[derive(Debug)]
 pub(crate) struct SharedUriCache {
-    cache: RwLock<HashMap<u64, Arc<Uri<String>>, BuildNoHashHasher>>,
+    cache: RwLock<CacheMap>,
 }
 
 impl Clone for SharedUriCache {
@@ -66,7 +79,15 @@ impl Clone for SharedUriCache {
                 self.cache
                     .read()
                     .iter()
-                    .map(|(k, v)| (*k, Arc::clone(v)))
+                    .map(|(base, entries)| {
+                        (
+                            base.clone(),
+                            entries
+                                .iter()
+                                .map(|(reference, value)| (reference.clone(), Arc::clone(value)))
+                                .collect(),
+                        )
+                    })
                     .collect(),
             ),
         }
@@ -79,17 +100,45 @@ impl SharedUriCache {
         base: &Uri<&str>,
         uri: impl AsRef<str>,
     ) -> Result<Arc<Uri<String>>, Error> {
-        let mut hasher = AHasher::default();
-        (base.as_str(), uri.as_ref()).hash(&mut hasher);
-        let hash = hasher.finish();
+        let base_str = base.as_str();
+        let reference = uri.as_ref();
 
-        if let Some(cached) = self.cache.read().get(&hash).cloned() {
-            return Ok(cached);
+        if let Some(cached) = self
+            .cache
+            .read()
+            .get(base_str)
+            .and_then(|inner| inner.get(reference))
+        {
+            return Ok(Arc::clone(cached));
         }
 
-        let new = Arc::new(uri::resolve_against(base, uri.as_ref())?);
-        self.cache.write().insert(hash, Arc::clone(&new));
-        Ok(new)
+        let cache = self.cache.upgradable_read();
+        if let Some(inner) = cache.get(base_str).and_then(|inner| inner.get(reference)) {
+            return Ok(Arc::clone(inner));
+        }
+
+        let resolved = Arc::new(uri::resolve_against(base, reference)?);
+
+        let mut cache = RwLockUpgradableReadGuard::upgrade(cache);
+        let inserted = match cache.entry_ref(base_str) {
+            EntryRef::Occupied(mut entry) => {
+                if let Some(existing) = entry.get().get(reference) {
+                    return Ok(Arc::clone(existing));
+                }
+                entry
+                    .get_mut()
+                    .insert(reference.to_owned(), Arc::clone(&resolved));
+                resolved
+            }
+            EntryRef::Vacant(entry) => {
+                let mut inner = HashMap::with_capacity(1);
+                inner.insert(reference.to_owned(), Arc::clone(&resolved));
+                entry.insert(inner);
+                resolved
+            }
+        };
+
+        Ok(inserted)
     }
 
     pub(crate) fn into_local(self) -> UriCache {

crates/jsonschema-referencing/src/hasher.rs

@@ -4,7 +4,6 @@
 mod anchors;
 mod cache;
 mod error;
-mod hasher;
 mod list;
 pub mod meta;
 mod registry;

crates/jsonschema-referencing/src/lib.rs

@@ -1,18 +1,17 @@
 use std::{
-    collections::{hash_map::Entry, HashSet, VecDeque},
-    hash::{Hash, Hasher},
+    collections::{hash_map::Entry, VecDeque},
+    num::NonZeroUsize,
     pin::Pin,
     sync::{Arc, LazyLock},
 };
 
-use ahash::{AHashMap, AHashSet, AHasher};
+use ahash::{AHashMap, AHashSet};
 use fluent_uri::Uri;
 use serde_json::Value;
 
 use crate::{
     anchors::{AnchorKey, AnchorKeyRef},
     cache::{SharedUriCache, UriCache},
-    hasher::BuildNoHashHasher,
     list::List,
     meta::{self, metas_for_draft},
     resource::{unescape_segment, InnerResourcePtr, JsonSchemaResource},
@@ -616,9 +615,27 @@ fn process_meta_schemas(
     Ok(())
 }
 
+#[derive(Hash, Eq, PartialEq)]
+struct ReferenceKey {
+    base_ptr: NonZeroUsize,
+    reference: String,
+}
+
+impl ReferenceKey {
+    fn new(base: &Arc<Uri<String>>, reference: &str) -> Self {
+        Self {
+            base_ptr: NonZeroUsize::new(Arc::as_ptr(base) as usize)
+                .expect("Arc pointer should never be null"),
+            reference: reference.to_owned(),
+        }
+    }
+}
+
+type ReferenceTracker = AHashSet<ReferenceKey>;
+
 struct ProcessingState {
     queue: VecDeque<(Arc<Uri<String>>, InnerResourcePtr)>,
-    seen: HashSet<u64, BuildNoHashHasher>,
+    seen: ReferenceTracker,
     external: AHashSet<(String, Uri<String>)>,
     scratch: String,
     refers_metaschemas: bool,
@@ -628,7 +645,7 @@ impl ProcessingState {
     fn new() -> Self {
         Self {
             queue: VecDeque::with_capacity(32),
-            seen: HashSet::with_hasher(BuildNoHashHasher::default()),
+            seen: ReferenceTracker::new(),
             external: AHashSet::new(),
             scratch: String::new(),
             refers_metaschemas: false,
@@ -879,10 +896,10 @@ async fn process_resources_async(
 }
 
 fn collect_external_resources(
-    base: &Uri<String>,
+    base: &Arc<Uri<String>>,
     contents: &Value,
     collected: &mut AHashSet<(String, Uri<String>)>,
-    seen: &mut HashSet<u64, BuildNoHashHasher>,
+    seen: &mut ReferenceTracker,
     resolution_cache: &mut UriCache,
     scratch: &mut String,
     refers_metaschemas: &mut bool,
@@ -903,10 +920,7 @@ fn collect_external_resources(
                     *refers_metaschemas = true;
                 }
             } else if $reference != "#" {
-                let mut hasher = AHasher::default();
-                (base.as_str(), $reference).hash(&mut hasher);
-                let hash = hasher.finish();
-                if seen.insert(hash) {
+                if mark_reference(seen, base, $reference) {
                     // Handle local references separately as they may have nested references to external resources
                     if $reference.starts_with('#') {
                         if let Some(referenced) =
@@ -924,7 +938,7 @@ fn collect_external_resources(
                         }
                     } else {
                         let resolved = if base.has_fragment() {
-                            let mut base_without_fragment = base.clone();
+                            let mut base_without_fragment = base.as_ref().clone();
                             base_without_fragment.set_fragment(None);
 
                             let (path, fragment) = match $reference.split_once('#') {
@@ -951,7 +965,9 @@ fn collect_external_resources(
                             }
                             resolved
                         } else {
-                            (*resolution_cache.resolve_against(&base.borrow(), $reference)?).clone()
+                            (*resolution_cache
+                                .resolve_against(&base.borrow(), $reference)?)
+                            .clone()
                         };
 
                         collected.insert(($reference.to_string(), resolved));
@@ -986,6 +1002,10 @@ fn collect_external_resources(
     Ok(())
 }
 
+fn mark_reference(seen: &mut ReferenceTracker, base: &Arc<Uri<String>>, reference: &str) -> bool {
+    seen.insert(ReferenceKey::new(base, reference))
+}
+
 /// Look up a value by a JSON Pointer.
 ///
 /// **NOTE**: A slightly faster version of pointer resolution based on `Value::pointer` from `serde_json`.