improvement: more efficient auth chain cache

src/database/rooms.rs

@@ -19,7 +19,7 @@ use ruma::{
     },
     push::{self, Action, Tweak},
     serde::{CanonicalJsonObject, CanonicalJsonValue, Raw},
-    state_res::{self, Event, RoomVersion, StateMap},
+    state_res::{self, RoomVersion, StateMap},
     uint, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
 };
 use std::{
@@ -91,7 +91,7 @@ pub struct Rooms {
     pub(super) referencedevents: Arc<dyn Tree>,
 
     pub(super) pdu_cache: Mutex<LruCache<EventId, Arc<PduEvent>>>,
-    pub(super) auth_chain_cache: Mutex<LruCache<u64, HashSet<u64>>>,
+    pub(super) auth_chain_cache: Mutex<LruCache<Vec<u64>, HashSet<u64>>>,
     pub(super) shorteventid_cache: Mutex<LruCache<u64, EventId>>,
     pub(super) eventidshort_cache: Mutex<LruCache<EventId, u64>>,
     pub(super) statekeyshort_cache: Mutex<LruCache<(EventType, String), u64>>,
@@ -3166,7 +3166,7 @@ impl Rooms {
     }
 
     #[tracing::instrument(skip(self))]
-    pub fn auth_chain_cache(&self) -> std::sync::MutexGuard<'_, LruCache<u64, HashSet<u64>>> {
+    pub fn auth_chain_cache(&self) -> std::sync::MutexGuard<'_, LruCache<Vec<u64>, HashSet<u64>>> {
         self.auth_chain_cache.lock().unwrap()
     }
 }

src/server_server.rs

@@ -1976,27 +1976,39 @@ fn get_auth_chain(
 ) -> Result<impl Iterator<Item = EventId> + '_> {
     let mut full_auth_chain = HashSet::new();
 
-    let starting_events = starting_events
-        .iter()
-        .map(|id| {
-            db.rooms
-                .get_or_create_shorteventid(id, &db.globals)
-                .map(|s| (s, id))
-        })
-        .collect::<Result<Vec<_>>>()?;
+    const NUM_BUCKETS: usize = 100;
+
+    let mut buckets = vec![HashSet::new(); NUM_BUCKETS];
+
+    for id in starting_events {
+        let short = db.rooms.get_or_create_shorteventid(&id, &db.globals)?;
+        let bucket_id = (short % NUM_BUCKETS as u64) as usize;
+        buckets[bucket_id].insert((short, id));
+    }
 
     let mut cache = db.rooms.auth_chain_cache();
 
-    for (sevent_id, event_id) in starting_events {
-        if let Some(cached) = cache.get_mut(&sevent_id) {
+    for chunk in buckets {
+        let chunk_key = chunk.iter().map(|(short, _)| short).copied().collect();
+        if let Some(cached) = cache.get_mut(&chunk_key) {
             full_auth_chain.extend(cached.iter().cloned());
-        } else {
-            drop(cache);
-            let auth_chain = get_auth_chain_inner(&event_id, db)?;
-            cache = db.rooms.auth_chain_cache();
-            cache.insert(sevent_id, auth_chain.clone());
-            full_auth_chain.extend(auth_chain);
-        };
+            continue;
+        }
+
+        let mut chunk_cache = HashSet::new();
+        for (sevent_id, event_id) in chunk {
+            if let Some(cached) = cache.get_mut(&[sevent_id][..]) {
+                chunk_cache.extend(cached.iter().cloned());
+            } else {
+                drop(cache);
+                let auth_chain = get_auth_chain_inner(&event_id, db)?;
+                cache = db.rooms.auth_chain_cache();
+                cache.insert(vec![sevent_id], auth_chain.clone());
+                chunk_cache.extend(auth_chain);
+            };
+        }
+        cache.insert(chunk_key, chunk_cache.clone());
+        full_auth_chain.extend(chunk_cache);
     }
 
     drop(cache);