From b813c3464210f249ca4da42dfd8c8c58a5d4a794 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20K=C3=B6sters?= <timo@koesters.xyz>
Date: Wed, 4 Aug 2021 18:30:56 +0200
Subject: [PATCH] improvement: more auth chain caching

---
 src/database/rooms.rs |  4 ++--
 src/server_server.rs  | 17 +++++++++++------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index 9f7a6005..549aa8cf 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -87,7 +87,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<EventId, HashSet<EventId>>>,
+    pub(super) auth_chain_cache: Mutex<LruCache<Vec<EventId>, HashSet<EventId>>>,
 }
 
 impl Rooms {
@@ -2618,7 +2618,7 @@ impl Rooms {
     #[tracing::instrument(skip(self))]
     pub fn auth_chain_cache(
         &self,
-    ) -> std::sync::MutexGuard<'_, LruCache<EventId, HashSet<EventId>>> {
+    ) -> std::sync::MutexGuard<'_, LruCache<Vec<EventId>, HashSet<EventId>>> {
         self.auth_chain_cache.lock().unwrap()
     }
 }
diff --git a/src/server_server.rs b/src/server_server.rs
index 7b0bd58b..0e595d44 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -1727,9 +1727,13 @@ fn get_auth_chain(starting_events: Vec<EventId>, db: &Database) -> Result<HashSe
     let mut full_auth_chain = HashSet::new();
 
     let mut cache = db.rooms.auth_chain_cache();
-    for event_id in starting_events {
-        let auth_chain = if let Some(cached) = cache.get_mut(&event_id) {
-            cached.clone()
+    if let Some(cached) = cache.get_mut(&starting_events) {
+        return Ok(cached.clone());
+    }
+
+    for event_id in &starting_events {
+        if let Some(cached) = cache.get_mut(&[event_id.clone()][..]) {
+            full_auth_chain.extend(cached.iter().cloned());
         } else {
             drop(cache);
             let start = Instant::now();
@@ -1741,14 +1745,15 @@ fn get_auth_chain(starting_events: Vec<EventId>, db: &Database) -> Result<HashSe
 
             cache = db.rooms.auth_chain_cache();
 
-            cache.insert(event_id, auth_chain.clone());
+            cache.insert(vec![event_id.clone()], auth_chain.clone());
 
-            auth_chain
+            full_auth_chain.extend(auth_chain);
         };
 
-        full_auth_chain.extend(auth_chain);
     }
 
+    cache.insert(starting_events, full_auth_chain.clone());
+
     Ok(full_auth_chain)
 }