move auth_chain_cache to service

src/database.rs

@@ -236,8 +236,6 @@ pub(crate) struct KeyValueDatabase {
     pub(super) senderkey_pusher: Arc<dyn KvTree>,
 
     // Uncategorized trees
-    pub(super) auth_chain_cache:
-        Mutex<LruCache<Vec<ShortEventId>, Arc<HashSet<ShortEventId>>>>,
     pub(super) eventidshort_cache: Mutex<LruCache<OwnedEventId, ShortEventId>>,
     pub(super) statekeyshort_cache:
         Mutex<LruCache<(StateEventType, String), ShortStateKey>>,
@@ -466,14 +464,6 @@ impl KeyValueDatabase {
             global: builder.open_tree("global")?,
             server_signingkeys: builder.open_tree("server_signingkeys")?,
 
-            #[allow(
-                clippy::as_conversions,
-                clippy::cast_sign_loss,
-                clippy::cast_possible_truncation
-            )]
-            auth_chain_cache: Mutex::new(LruCache::new(
-                (100_000.0 * config.cache_capacity_modifier) as usize,
-            )),
             #[allow(
                 clippy::as_conversions,
                 clippy::cast_sign_loss,

src/database/key_value/rooms/auth_chain.rs

@@ -1,8 +1,7 @@
-use std::{collections::HashSet, mem::size_of, sync::Arc};
+use std::{collections::HashSet, mem::size_of};
 
 use crate::{
     database::KeyValueDatabase,
-    observability::{FoundIn, Lookup, METRICS},
     service::{self, rooms::short::ShortEventId},
     utils, Result,
 };
@@ -12,19 +11,9 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
     fn get_cached_eventid_authchain(
         &self,
         key: &[ShortEventId],
-    ) -> Result<Option<Arc<HashSet<ShortEventId>>>> {
+    ) -> Result<Option<HashSet<ShortEventId>>> {
-        let lookup = Lookup::AuthChain;
-
-        // Check RAM cache
-        if let Some(result) = self.auth_chain_cache.lock().unwrap().get_mut(key)
-        {
-            METRICS.record_lookup(lookup, FoundIn::Cache);
-            return Ok(Some(Arc::clone(result)));
-        }
-
         // We only save auth chains for single events in the db
         if key.len() == 1 {
-            // Check DB cache
             let chain = self
                 .shorteventid_authchain
                 .get(&key[0].get().to_be_bytes())?
@@ -40,28 +29,16 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
                         .collect()
                 });
 
-            if let Some(chain) = chain {
+            return Ok(chain);
-                METRICS.record_lookup(lookup, FoundIn::Database);
-                let chain = Arc::new(chain);
-
-                // Cache in RAM
-                self.auth_chain_cache
-                    .lock()
-                    .unwrap()
-                    .insert(vec![key[0]], Arc::clone(&chain));
-
-                return Ok(Some(chain));
-            }
         }
 
-        METRICS.record_lookup(lookup, FoundIn::Nothing);
         Ok(None)
     }
 
     fn cache_auth_chain(
         &self,
-        key: Vec<ShortEventId>,
+        key: &[ShortEventId],
-        auth_chain: Arc<HashSet<ShortEventId>>,
+        auth_chain: &HashSet<ShortEventId>,
     ) -> Result<()> {
         // Only persist single events in db
         if key.len() == 1 {
@@ -74,9 +51,6 @@ impl service::rooms::auth_chain::Data for KeyValueDatabase {
             )?;
         }
 
-        // Cache in RAM
-        self.auth_chain_cache.lock().unwrap().insert(key, auth_chain);
-
         Ok(())
     }
 }

src/service.rs

@@ -71,9 +71,17 @@ impl Services {
             },
             rooms: rooms::Service {
                 alias: rooms::alias::Service::new(db),
-                auth_chain: rooms::auth_chain::Service {
+                auth_chain: rooms::auth_chain::Service::new(
                     db,
-                },
+                    #[allow(
+                        clippy::as_conversions,
+                        clippy::cast_sign_loss,
+                        clippy::cast_possible_truncation
+                    )]
+                    {
+                        (100_000.0 * config.cache_capacity_modifier) as usize
+                    },
+                ),
                 directory: db,
                 edus: rooms::edus::Service {
                     read_receipt: db,

src/service/rooms/auth_chain.rs

@@ -1,28 +1,67 @@
 use std::{
     collections::{BTreeSet, HashSet},
-    sync::Arc,
+    sync::{Arc, Mutex},
 };
 
+use lru_cache::LruCache;
 use ruma::{api::client::error::ErrorKind, EventId, RoomId};
 use tracing::{debug, error, warn};
 
 use super::short::ShortEventId;
-use crate::{services, utils::debug_slice_truncated, Error, Result};
+use crate::{
+    observability::{FoundIn, Lookup, METRICS},
+    services,
+    utils::debug_slice_truncated,
+    Error, Result,
+};
 
 mod data;
 
 pub(crate) use data::Data;
 
 pub(crate) struct Service {
-    pub(crate) db: &'static dyn Data,
+    db: &'static dyn Data,
+    auth_chain_cache:
+        Mutex<LruCache<Vec<ShortEventId>, Arc<HashSet<ShortEventId>>>>,
 }
 
 impl Service {
+    pub(crate) fn new(
+        db: &'static dyn Data,
+        auth_chain_cache_size: usize,
+    ) -> Self {
+        Self {
+            db,
+            auth_chain_cache: Mutex::new(LruCache::new(auth_chain_cache_size)),
+        }
+    }
+
     pub(crate) fn get_cached_eventid_authchain(
         &self,
         key: &[ShortEventId],
     ) -> Result<Option<Arc<HashSet<ShortEventId>>>> {
-        self.db.get_cached_eventid_authchain(key)
+        let lookup = Lookup::AuthChain;
+
+        if let Some(result) = self.auth_chain_cache.lock().unwrap().get_mut(key)
+        {
+            METRICS.record_lookup(lookup, FoundIn::Cache);
+            return Ok(Some(Arc::clone(result)));
+        }
+
+        let Some(chain) = self.db.get_cached_eventid_authchain(key)? else {
+            METRICS.record_lookup(lookup, FoundIn::Nothing);
+            return Ok(None);
+        };
+
+        METRICS.record_lookup(lookup, FoundIn::Database);
+        let chain = Arc::new(chain);
+
+        self.auth_chain_cache
+            .lock()
+            .unwrap()
+            .insert(vec![key[0]], Arc::clone(&chain));
+
+        Ok(Some(chain))
     }
 
     #[tracing::instrument(skip(self))]
@@ -31,7 +70,9 @@ impl Service {
         key: Vec<ShortEventId>,
         auth_chain: Arc<HashSet<ShortEventId>>,
     ) -> Result<()> {
-        self.db.cache_auth_chain(key, auth_chain)
+        self.db.cache_auth_chain(&key, &auth_chain)?;
+        self.auth_chain_cache.lock().unwrap().insert(key, auth_chain);
+        Ok(())
     }
 
     #[tracing::instrument(

src/service/rooms/auth_chain/data.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
 
 use crate::{service::rooms::short::ShortEventId, Result};
 
@@ -6,10 +6,10 @@ pub(crate) trait Data: Send + Sync {
     fn get_cached_eventid_authchain(
         &self,
         shorteventid: &[ShortEventId],
-    ) -> Result<Option<Arc<HashSet<ShortEventId>>>>;
+    ) -> Result<Option<HashSet<ShortEventId>>>;
     fn cache_auth_chain(
         &self,
-        shorteventid: Vec<ShortEventId>,
+        shorteventid: &[ShortEventId],
-        auth_chain: Arc<HashSet<ShortEventId>>,
+        auth_chain: &HashSet<ShortEventId>,
     ) -> Result<()>;
 }