improvement: auth chain cache

src/database/abstraction/rocksdb.rs

@@ -5,14 +5,14 @@ use std::{future::Future, pin::Pin, sync::Arc};
 
 use super::{DatabaseEngine, Tree};
 
-use std::{collections::BTreeMap, sync::RwLock};
+use std::{collections::HashMap, sync::RwLock};
 
 pub struct Engine(rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>);
 
 pub struct RocksDbEngineTree<'a> {
     db: Arc<Engine>,
     name: &'a str,
-    watchers: RwLock<BTreeMap<Vec<u8>, Vec<tokio::sync::oneshot::Sender<()>>>>,
+    watchers: RwLock<HashMap<Vec<u8>, Vec<tokio::sync::oneshot::Sender<()>>>>,
 }
 
 impl DatabaseEngine for Engine {
@@ -58,7 +58,7 @@ impl DatabaseEngine for Engine {
         Ok(Arc::new(RocksDbEngineTree {
             name,
             db: Arc::clone(self),
-            watchers: RwLock::new(BTreeMap::new()),
+            watchers: RwLock::new(HashMap::new()),
         }))
     }
 }

src/database/abstraction/sqlite.rs

@@ -7,7 +7,7 @@ use log::debug;
 use parking_lot::{Mutex, MutexGuard, RwLock};
 use rusqlite::{params, Connection, DatabaseName::Main, OptionalExtension};
 use std::{
-    collections::BTreeMap,
+    collections::HashMap,
     future::Future,
     ops::Deref,
     path::{Path, PathBuf},
@@ -206,7 +206,7 @@ impl DatabaseEngine for Engine {
         Ok(Arc::new(SqliteTable {
             engine: Arc::clone(self),
             name: name.to_owned(),
-            watchers: RwLock::new(BTreeMap::new()),
+            watchers: RwLock::new(HashMap::new()),
         }))
     }
 
@@ -266,7 +266,7 @@ impl Engine {
 pub struct SqliteTable {
     engine: Arc<Engine>,
     name: String,
-    watchers: RwLock<BTreeMap<Vec<u8>, Vec<Sender<()>>>>,
+    watchers: RwLock<HashMap<Vec<u8>, Vec<Sender<()>>>>,
 }
 
 type TupleOfBytes = (Vec<u8>, Vec<u8>);

src/database/globals.rs

@@ -41,12 +41,12 @@ pub struct Globals {
     dns_resolver: TokioAsyncResolver,
     jwt_decoding_key: Option<jsonwebtoken::DecodingKey<'static>>,
     pub(super) server_signingkeys: Arc<dyn Tree>,
-    pub bad_event_ratelimiter: Arc<RwLock<BTreeMap<EventId, RateLimitState>>>,
-    pub bad_signature_ratelimiter: Arc<RwLock<BTreeMap<Vec<String>, RateLimitState>>>,
-    pub servername_ratelimiter: Arc<RwLock<BTreeMap<Box<ServerName>, Arc<Semaphore>>>>,
-    pub sync_receivers: RwLock<BTreeMap<(UserId, Box<DeviceId>), SyncHandle>>,
-    pub roomid_mutex: RwLock<BTreeMap<RoomId, Arc<Mutex<()>>>>,
-    pub roomid_mutex_federation: RwLock<BTreeMap<RoomId, Arc<Mutex<()>>>>, // this lock will be held longer
+    pub bad_event_ratelimiter: Arc<RwLock<HashMap<EventId, RateLimitState>>>,
+    pub bad_signature_ratelimiter: Arc<RwLock<HashMap<Vec<String>, RateLimitState>>>,
+    pub servername_ratelimiter: Arc<RwLock<HashMap<Box<ServerName>, Arc<Semaphore>>>>,
+    pub sync_receivers: RwLock<HashMap<(UserId, Box<DeviceId>), SyncHandle>>,
+    pub roomid_mutex: RwLock<HashMap<RoomId, Arc<Mutex<()>>>>,
+    pub roomid_mutex_federation: RwLock<HashMap<RoomId, Arc<Mutex<()>>>>, // this lock will be held longer
     pub rotate: RotationHandler,
 }
 
@@ -196,12 +196,12 @@ impl Globals {
             tls_name_override,
             server_signingkeys,
             jwt_decoding_key,
-            bad_event_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())),
-            bad_signature_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())),
-            servername_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())),
-            roomid_mutex: RwLock::new(BTreeMap::new()),
-            roomid_mutex_federation: RwLock::new(BTreeMap::new()),
-            sync_receivers: RwLock::new(BTreeMap::new()),
+            bad_event_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
+            bad_signature_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
+            servername_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
+            roomid_mutex: RwLock::new(HashMap::new()),
+            roomid_mutex_federation: RwLock::new(HashMap::new()),
+            sync_receivers: RwLock::new(HashMap::new()),
             rotate: RotationHandler::new(),
         };
 

src/database/rooms.rs

@@ -25,7 +25,7 @@ use std::{
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     convert::{TryFrom, TryInto},
     mem,
-    sync::{Arc, RwLock},
+    sync::{Arc, Mutex},
 };
 
 use super::{abstraction::Tree, admin::AdminCommand, pusher};
@@ -84,7 +84,8 @@ pub struct Rooms {
     /// RoomId + EventId -> Parent PDU EventId.
     pub(super) prevevent_parent: Arc<dyn Tree>,
 
-    pub(super) pdu_cache: RwLock<LruCache<EventId, Arc<PduEvent>>>,
+    pub(super) pdu_cache: Mutex<LruCache<EventId, Arc<PduEvent>>>,
+    pub(super) auth_chain_cache: Mutex<LruCache<EventId, HashSet<EventId>>>,
 }
 
 impl Rooms {
@@ -109,7 +110,7 @@ impl Rooms {
     pub fn state_full(
         &self,
         shortstatehash: u64,
-    ) -> Result<BTreeMap<(EventType, String), Arc<PduEvent>>> {
+    ) -> Result<HashMap<(EventType, String), Arc<PduEvent>>> {
         let state = self
             .stateid_shorteventid
             .scan_prefix(shortstatehash.to_be_bytes().to_vec())
@@ -282,7 +283,7 @@ impl Rooms {
     pub fn force_state(
         &self,
         room_id: &RoomId,
-        state: BTreeMap<(EventType, String), EventId>,
+        state: HashMap<(EventType, String), EventId>,
         db: &Database,
     ) -> Result<()> {
         let state_hash = self.calculate_hash(
@@ -402,11 +403,11 @@ impl Rooms {
     pub fn room_state_full(
         &self,
         room_id: &RoomId,
-    ) -> Result<BTreeMap<(EventType, String), Arc<PduEvent>>> {
+    ) -> Result<HashMap<(EventType, String), Arc<PduEvent>>> {
         if let Some(current_shortstatehash) = self.current_shortstatehash(room_id)? {
             self.state_full(current_shortstatehash)
         } else {
-            Ok(BTreeMap::new())
+            Ok(HashMap::new())
         }
     }
 
@@ -542,7 +543,7 @@ impl Rooms {
     ///
     /// Checks the `eventid_outlierpdu` Tree if not found in the timeline.
     pub fn get_pdu(&self, event_id: &EventId) -> Result<Option<Arc<PduEvent>>> {
-        if let Some(p) = self.pdu_cache.write().unwrap().get_mut(&event_id) {
+        if let Some(p) = self.pdu_cache.lock().unwrap().get_mut(&event_id) {
             return Ok(Some(Arc::clone(p)));
         }
 
@@ -568,7 +569,7 @@ impl Rooms {
             .transpose()?
         {
             self.pdu_cache
-                .write()
+                .lock()
                 .unwrap()
                 .insert(event_id.clone(), Arc::clone(&pdu));
             Ok(Some(pdu))
@@ -2520,4 +2521,10 @@ impl Rooms {
 
         Ok(self.userroomid_leftstate.get(&userroom_id)?.is_some())
     }
+
+    pub fn auth_chain_cache(
+        &self,
+    ) -> std::sync::MutexGuard<'_, LruCache<EventId, HashSet<EventId>>> {
+        self.auth_chain_cache.lock().unwrap()
+    }
 }