move our_real_users_cache to service

src/service/rooms/state_cache.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashSet, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::{Arc, RwLock},
+};
 
 use ruma::{
     events::{
@@ -12,19 +15,32 @@ use ruma::{
 };
 use tracing::warn;
 
-use crate::{service::appservice::RegistrationInfo, services, Error, Result};
+use crate::{
+    observability::{FoundIn, Lookup, METRICS},
+    service::appservice::RegistrationInfo,
+    services, Error, Result,
+};
 
 mod data;
 
 pub(crate) use data::Data;
 
 pub(crate) struct Service {
-    pub(crate) db: &'static dyn Data,
+    db: &'static dyn Data,
+    our_real_users_cache:
+        RwLock<HashMap<OwnedRoomId, Arc<HashSet<OwnedUserId>>>>,
 }
 
 type RoomsLeft = (OwnedRoomId, Vec<Raw<AnySyncStateEvent>>);
 
 impl Service {
+    pub(crate) fn new(db: &'static dyn Data) -> Self {
+        Self {
+            db,
+            our_real_users_cache: RwLock::new(HashMap::new()),
+        }
+    }
+
     /// Update current membership data.
     #[tracing::instrument(skip(self, last_state))]
     pub(crate) fn update_membership(
@@ -283,8 +299,18 @@ impl Service {
     }
 
     #[tracing::instrument(skip(self, room_id))]
-    pub(crate) fn update_joined_count(&self, room_id: &RoomId) -> Result<()> {
-        self.db.update_joined_count(room_id)
+    pub(crate) fn update_joined_count(
+        &self,
+        room_id: &RoomId,
+    ) -> Result<Arc<HashSet<OwnedUserId>>> {
+        let our_real_users = Arc::new(self.db.update_joined_count(room_id)?);
+
+        self.our_real_users_cache
+            .write()
+            .unwrap()
+            .insert(room_id.to_owned(), our_real_users.clone());
+
+        Ok(our_real_users)
     }
 
     #[tracing::instrument(skip(self, room_id))]
@@ -292,7 +318,20 @@ impl Service {
         &self,
         room_id: &RoomId,
     ) -> Result<Arc<HashSet<OwnedUserId>>> {
-        self.db.get_our_real_users(room_id)
+        let lookup = Lookup::OurRealUsers;
+
+        if let Some(our_real_users) =
+            self.our_real_users_cache.read().unwrap().get(room_id).cloned()
+        {
+            METRICS.record_lookup(lookup, FoundIn::Cache);
+            return Ok(our_real_users);
+        }
+
+        let our_real_users = self.update_joined_count(room_id)?;
+
+        METRICS.record_lookup(lookup, FoundIn::Database);
+
+        Ok(our_real_users)
     }
 
     #[tracing::instrument(skip(self, room_id, appservice))]

src/service/rooms/state_cache/data.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
 
 use ruma::{
     events::{AnyStrippedStateEvent, AnySyncStateEvent},
@@ -23,12 +23,10 @@ pub(crate) trait Data: Send + Sync {
     ) -> Result<()>;
     fn mark_as_left(&self, user_id: &UserId, room_id: &RoomId) -> Result<()>;
 
-    fn update_joined_count(&self, room_id: &RoomId) -> Result<()>;
-
-    fn get_our_real_users(
+    fn update_joined_count(
         &self,
         room_id: &RoomId,
-    ) -> Result<Arc<HashSet<OwnedUserId>>>;
+    ) -> Result<HashSet<OwnedUserId>>;
 
     fn appservice_in_room(
         &self,