move pdu_cache to service

src/service/rooms/timeline.rs

@@ -1,9 +1,10 @@
 use std::{
     cmp::Ordering,
     collections::{BTreeMap, HashSet},
-    sync::Arc,
+    sync::{Arc, Mutex},
 };
 
+use lru_cache::LruCache;
 use ruma::{
     api::{client::error::ErrorKind, federation},
     canonical_json::to_canonical_value,
@@ -30,6 +31,7 @@ use tracing::{error, info, warn};
 use super::{short::ShortRoomId, state_compressor::CompressedStateEvent};
 use crate::{
     api::server_server,
+    observability::{FoundIn, Lookup, METRICS},
     service::{
         appservice::NamespaceRegex,
         globals::{marker, SigningKeys},
@@ -111,10 +113,21 @@ impl Ord for PduCount {
 }
 
 pub(crate) struct Service {
-    pub(crate) db: &'static dyn Data,
+    db: &'static dyn Data,
+    pdu_cache: Mutex<LruCache<OwnedEventId, Arc<PduEvent>>>,
 }
 
 impl Service {
+    pub(crate) fn new(
+        db: &'static dyn Data,
+        pdu_cache_capacity: usize,
+    ) -> Self {
+        Self {
+            db,
+            pdu_cache: Mutex::new(LruCache::new(pdu_cache_capacity)),
+        }
+    }
+
     #[tracing::instrument(skip(self))]
     pub(crate) fn first_pdu_in_room(
         &self,
@@ -174,7 +187,24 @@ impl Service {
         &self,
         event_id: &EventId,
     ) -> Result<Option<Arc<PduEvent>>> {
-        self.db.get_pdu(event_id)
+        let lookup = Lookup::Pdu;
+
+        if let Some(p) = self.pdu_cache.lock().unwrap().get_mut(event_id) {
+            METRICS.record_lookup(lookup, FoundIn::Cache);
+            return Ok(Some(Arc::clone(p)));
+        }
+
+        if let Some(pdu) = self.db.get_pdu(event_id)? {
+            METRICS.record_lookup(lookup, FoundIn::Database);
+            self.pdu_cache
+                .lock()
+                .unwrap()
+                .insert(event_id.to_owned(), Arc::clone(&pdu));
+            Ok(Some(pdu))
+        } else {
+            METRICS.record_lookup(lookup, FoundIn::Nothing);
+            Ok(None)
+        }
     }
 
     /// Returns the pdu.
@@ -203,7 +233,9 @@ impl Service {
         pdu_json: &CanonicalJsonObject,
         pdu: &PduEvent,
     ) -> Result<()> {
-        self.db.replace_pdu(pdu_id, pdu_json, pdu)
+        self.db.replace_pdu(pdu_id, pdu_json)?;
+        self.pdu_cache.lock().unwrap().remove(&(*pdu.event_id).to_owned());
+        Ok(())
     }
 
     /// Creates a new persisted data unit and adds it to a room.

src/service/rooms/timeline/data.rs

@@ -76,7 +76,6 @@ pub(crate) trait Data: Send + Sync {
         &self,
         pdu_id: &PduId,
         pdu_json: &CanonicalJsonObject,
-        pdu: &PduEvent,
     ) -> Result<()>;
 
     /// Returns an iterator over all events and their tokens in a room that