feat: lazy loading

src/database/rooms.rs

@@ -28,7 +28,7 @@ use ruma::{
     push::{Action, Ruleset, Tweak},
     serde::{CanonicalJsonObject, CanonicalJsonValue, Raw},
     state_res::{self, RoomVersion, StateMap},
-    uint, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
+    uint, DeviceId, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
 };
 use serde::Deserialize;
 use serde_json::value::to_raw_value;
@@ -79,6 +79,8 @@ pub struct Rooms {
     pub(super) userroomid_leftstate: Arc<dyn Tree>,
     pub(super) roomuserid_leftcount: Arc<dyn Tree>,
 
+    pub(super) lazyloadedids: Arc<dyn Tree>, // LazyLoadedIds = UserId + DeviceId + RoomId + LazyLoadedUserId
+
     pub(super) userroomid_notificationcount: Arc<dyn Tree>, // NotifyCount = u64
     pub(super) userroomid_highlightcount: Arc<dyn Tree>,    // HightlightCount = u64
 
@@ -117,6 +119,8 @@ pub struct Rooms {
     pub(super) shortstatekey_cache: Mutex<LruCache<u64, (EventType, String)>>,
     pub(super) our_real_users_cache: RwLock<HashMap<Box<RoomId>, Arc<HashSet<Box<UserId>>>>>,
     pub(super) appservice_in_room_cache: RwLock<HashMap<Box<RoomId>, HashMap<String, bool>>>,
+    pub(super) lazy_load_waiting:
+        Mutex<HashMap<(Box<UserId>, Box<DeviceId>, Box<RoomId>, u64), Vec<Box<UserId>>>>,
     pub(super) stateinfo_cache: Mutex<
         LruCache<
             u64,
@@ -3466,4 +3470,94 @@ impl Rooms {
 
         Ok(())
     }
+
+    #[tracing::instrument(skip(self))]
+    pub fn lazy_load_was_sent_before(
+        &self,
+        user_id: &UserId,
+        device_id: &DeviceId,
+        room_id: &RoomId,
+        ll_user: &UserId,
+    ) -> Result<bool> {
+        let mut key = user_id.as_bytes().to_vec();
+        key.push(0xff);
+        key.extend_from_slice(&device_id.as_bytes());
+        key.push(0xff);
+        key.extend_from_slice(&room_id.as_bytes());
+        key.push(0xff);
+        key.extend_from_slice(&ll_user.as_bytes());
+        Ok(self.lazyloadedids.get(&key)?.is_some())
+    }
+
+    #[tracing::instrument(skip(self))]
+    pub fn lazy_load_mark_sent(
+        &self,
+        user_id: &UserId,
+        device_id: &DeviceId,
+        room_id: &RoomId,
+        lazy_load: Vec<Box<UserId>>,
+        count: u64,
+    ) {
+        self.lazy_load_waiting.lock().unwrap().insert(
+            (
+                user_id.to_owned(),
+                device_id.to_owned(),
+                room_id.to_owned(),
+                count,
+            ),
+            lazy_load,
+        );
+    }
+
+    #[tracing::instrument(skip(self))]
+    pub fn lazy_load_confirm_delivery(
+        &self,
+        user_id: &UserId,
+        device_id: &DeviceId,
+        room_id: &RoomId,
+        since: u64,
+    ) -> Result<()> {
+        if let Some(user_ids) = self.lazy_load_waiting.lock().unwrap().remove(&(
+            user_id.to_owned(),
+            device_id.to_owned(),
+            room_id.to_owned(),
+            since,
+        )) {
+            let mut prefix = user_id.as_bytes().to_vec();
+            prefix.push(0xff);
+            prefix.extend_from_slice(&device_id.as_bytes());
+            prefix.push(0xff);
+            prefix.extend_from_slice(&room_id.as_bytes());
+            prefix.push(0xff);
+
+            for ll_id in user_ids {
+                let mut key = prefix.clone();
+                key.extend_from_slice(&ll_id.as_bytes());
+                self.lazyloadedids.insert(&key, &[])?;
+            }
+        }
+
+        Ok(())
+    }
+
+    #[tracing::instrument(skip(self))]
+    pub fn lazy_load_reset(
+        &self,
+        user_id: &Box<UserId>,
+        device_id: &Box<DeviceId>,
+        room_id: &Box<RoomId>,
+    ) -> Result<()> {
+        let mut prefix = user_id.as_bytes().to_vec();
+        prefix.push(0xff);
+        prefix.extend_from_slice(&device_id.as_bytes());
+        prefix.push(0xff);
+        prefix.extend_from_slice(&room_id.as_bytes());
+        prefix.push(0xff);
+
+        for (key, _) in self.lazyloadedids.scan_prefix(prefix) {
+            self.lazyloadedids.remove(&key)?;
+        }
+
+        Ok(())
+    }
 }