move lasttimelinecount_cache to service

src/service/rooms/timeline.rs

@@ -1,6 +1,6 @@
 use std::{
     cmp::Ordering,
-    collections::{BTreeMap, HashSet},
+    collections::{hash_map, BTreeMap, HashMap, HashSet},
     sync::{Arc, Mutex},
 };
 
@@ -115,6 +115,7 @@ impl Ord for PduCount {
 pub(crate) struct Service {
     db: &'static dyn Data,
     pdu_cache: Mutex<LruCache<OwnedEventId, Arc<PduEvent>>>,
+    lasttimelinecount_cache: Mutex<HashMap<OwnedRoomId, PduCount>>,
 }
 
 impl Service {
@@ -125,6 +126,7 @@ impl Service {
         Self {
             db,
             pdu_cache: Mutex::new(LruCache::new(pdu_cache_capacity)),
+            lasttimelinecount_cache: Mutex::new(HashMap::new()),
         }
     }
 
@@ -145,7 +147,37 @@ impl Service {
         sender_user: &UserId,
         room_id: &RoomId,
     ) -> Result<PduCount> {
-        self.db.last_timeline_count(sender_user, room_id)
+        let lookup = Lookup::LastTimelineCount;
+
+        match self
+            .lasttimelinecount_cache
+            .lock()
+            .unwrap()
+            .entry(room_id.to_owned())
+        {
+            hash_map::Entry::Vacant(v) => {
+                if let Some(last_count) = self
+                    .pdus_until(sender_user, room_id, PduCount::MAX)?
+                    .find_map(|x| match x {
+                        Ok(x) => Some(x),
+                        Err(error) => {
+                            error!(%error, "Bad pdu in pdus_since");
+                            None
+                        }
+                    })
+                {
+                    METRICS.record_lookup(lookup, FoundIn::Database);
+                    Ok(*v.insert(last_count.0))
+                } else {
+                    METRICS.record_lookup(lookup, FoundIn::Nothing);
+                    Ok(PduCount::Normal(0))
+                }
+            }
+            hash_map::Entry::Occupied(o) => {
+                METRICS.record_lookup(lookup, FoundIn::Cache);
+                Ok(*o.get())
+            }
+        }
     }
 
     /// Returns the `count` of this pdu's id.
@@ -334,7 +366,11 @@ impl Service {
         let pdu_id = PduId::new(pdu_id);
 
         // Insert pdu
-        self.db.append_pdu(&pdu_id, pdu, &pdu_json, count2)?;
+        self.db.append_pdu(&pdu_id, pdu, &pdu_json)?;
+        self.lasttimelinecount_cache
+            .lock()
+            .unwrap()
+            .insert(pdu.room_id.clone(), PduCount::Normal(count2));
 
         drop(insert_token);
 

src/service/rooms/timeline/data.rs

@@ -6,12 +6,6 @@ use super::PduCount;
 use crate::{service::rooms::timeline::PduId, PduEvent, Result};
 
 pub(crate) trait Data: Send + Sync {
-    fn last_timeline_count(
-        &self,
-        sender_user: &UserId,
-        room_id: &RoomId,
-    ) -> Result<PduCount>;
-
     /// Returns the `count` of this pdu's id.
     fn get_pdu_count(&self, event_id: &EventId) -> Result<Option<PduCount>>;
 
@@ -60,7 +54,6 @@ pub(crate) trait Data: Send + Sync {
         pdu_id: &PduId,
         pdu: &PduEvent,
         json: &CanonicalJsonObject,
-        count: u64,
     ) -> Result<()>;
 
     // Adds a new pdu to the backfilled timeline