From 107f4521e0dd622d49f8aad9d42010b62a1398b7 Mon Sep 17 00:00:00 2001
From: Charles Hall <charles@computer.surgery>
Date: Tue, 8 Oct 2024 15:41:17 -0700
Subject: [PATCH] move appservice_in_room_cache to service

---
 src/database.rs                             |  5 +--
 src/database/key_value/rooms/state_cache.rs | 44 +++++----------------
 src/service/rooms/state_cache.rs            | 30 +++++++++++++-
 3 files changed, 39 insertions(+), 40 deletions(-)

diff --git a/src/database.rs b/src/database.rs
index bf22a0a3..66025b90 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -4,7 +4,7 @@ use std::{
     io::Write,
     mem::size_of,
     path::Path,
-    sync::{Arc, Mutex, RwLock},
+    sync::{Arc, Mutex},
 };
 
 use ruma::{
@@ -233,8 +233,6 @@ pub(crate) struct KeyValueDatabase {
     pub(super) senderkey_pusher: Arc<dyn KvTree>,
 
     // Uncategorized trees
-    pub(super) appservice_in_room_cache:
-        RwLock<HashMap<OwnedRoomId, HashMap<String, bool>>>,
     pub(super) lasttimelinecount_cache: Mutex<HashMap<OwnedRoomId, PduCount>>,
 }
 
@@ -454,7 +452,6 @@ impl KeyValueDatabase {
             global: builder.open_tree("global")?,
             server_signingkeys: builder.open_tree("server_signingkeys")?,
 
-            appservice_in_room_cache: RwLock::new(HashMap::new()),
             lasttimelinecount_cache: Mutex::new(HashMap::new()),
         };
 
diff --git a/src/database/key_value/rooms/state_cache.rs b/src/database/key_value/rooms/state_cache.rs
index 78ee1fb7..2c2e215f 100644
--- a/src/database/key_value/rooms/state_cache.rs
+++ b/src/database/key_value/rooms/state_cache.rs
@@ -8,7 +8,6 @@ use ruma::{
 
 use crate::{
     database::KeyValueDatabase,
-    observability::{FoundIn, Lookup, METRICS},
     service::{self, appservice::RegistrationInfo},
     services, utils, Error, Result,
 };
@@ -164,8 +163,6 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
             self.serverroomids.insert(&serverroom_id, &[])?;
         }
 
-        self.appservice_in_room_cache.write().unwrap().remove(room_id);
-
         Ok(real_users)
     }
 
@@ -178,44 +175,21 @@ impl service::rooms::state_cache::Data for KeyValueDatabase {
         room_id: &RoomId,
         appservice: &RegistrationInfo,
     ) -> Result<bool> {
-        let lookup = Lookup::AppserviceInRoom;
+        let bridge_user_id = UserId::parse_with_server_name(
+            appservice.registration.sender_localpart.as_str(),
+            services().globals.server_name(),
+        )
+        .ok();
 
-        let maybe = self
-            .appservice_in_room_cache
-            .read()
-            .unwrap()
-            .get(room_id)
-            .and_then(|map| map.get(&appservice.registration.id))
-            .copied();
-
-        if let Some(b) = maybe {
-            METRICS.record_lookup(lookup, FoundIn::Cache);
-            Ok(b)
-        } else {
-            let bridge_user_id = UserId::parse_with_server_name(
-                appservice.registration.sender_localpart.as_str(),
-                services().globals.server_name(),
-            )
-            .ok();
-
-            let in_room = bridge_user_id.map_or(false, |id| {
-                self.is_joined(&id, room_id).unwrap_or(false)
-            }) || self.room_members(room_id).any(|userid| {
+        let in_room = bridge_user_id
+            .map_or(false, |id| self.is_joined(&id, room_id).unwrap_or(false))
+            || self.room_members(room_id).any(|userid| {
                 userid.map_or(false, |userid| {
                     appservice.users.is_match(userid.as_str())
                 })
             });
 
-            METRICS.record_lookup(lookup, FoundIn::Database);
-            self.appservice_in_room_cache
-                .write()
-                .unwrap()
-                .entry(room_id.to_owned())
-                .or_default()
-                .insert(appservice.registration.id.clone(), in_room);
-
-            Ok(in_room)
-        }
+        Ok(in_room)
     }
 
     /// Makes a user forget a room.
diff --git a/src/service/rooms/state_cache.rs b/src/service/rooms/state_cache.rs
index f3e6144c..9cf264a9 100644
--- a/src/service/rooms/state_cache.rs
+++ b/src/service/rooms/state_cache.rs
@@ -27,6 +27,8 @@ pub(crate) use data::Data;
 
 pub(crate) struct Service {
     db: &'static dyn Data,
+    appservice_in_room_cache:
+        RwLock<HashMap<OwnedRoomId, HashMap<String, bool>>>,
     our_real_users_cache:
         RwLock<HashMap<OwnedRoomId, Arc<HashSet<OwnedUserId>>>>,
 }
@@ -38,6 +40,7 @@ impl Service {
         Self {
             db,
             our_real_users_cache: RwLock::new(HashMap::new()),
+            appservice_in_room_cache: RwLock::new(HashMap::new()),
         }
     }
 
@@ -309,6 +312,7 @@ impl Service {
             .write()
             .unwrap()
             .insert(room_id.to_owned(), our_real_users.clone());
+        self.appservice_in_room_cache.write().unwrap().remove(room_id);
 
         Ok(our_real_users)
     }
@@ -340,7 +344,31 @@ impl Service {
         room_id: &RoomId,
         appservice: &RegistrationInfo,
     ) -> Result<bool> {
-        self.db.appservice_in_room(room_id, appservice)
+        let lookup = Lookup::AppserviceInRoom;
+
+        if let Some(in_room) = self
+            .appservice_in_room_cache
+            .read()
+            .unwrap()
+            .get(room_id)
+            .and_then(|map| map.get(&appservice.registration.id))
+            .copied()
+        {
+            METRICS.record_lookup(lookup, FoundIn::Cache);
+            return Ok(in_room);
+        }
+
+        let in_room = self.db.appservice_in_room(room_id, appservice)?;
+
+        METRICS.record_lookup(lookup, FoundIn::Database);
+        self.appservice_in_room_cache
+            .write()
+            .unwrap()
+            .entry(room_id.to_owned())
+            .or_default()
+            .insert(appservice.registration.id.clone(), in_room);
+
+        Ok(in_room)
     }
 
     /// Makes a user forget a room.