From f285c89006e48b0644f421ae399f1a8eb47e37f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20K=C3=B6sters?= <timo@koesters.xyz>
Date: Thu, 6 Jan 2022 00:15:34 +0100
Subject: [PATCH] fix: make incremental sync efficient again

---
 src/client_server/message.rs |  2 +-
 src/client_server/sync.rs    | 77 ++++++++++++++++++++----------------
 src/database/rooms.rs        |  4 +-
 3 files changed, 47 insertions(+), 36 deletions(-)

diff --git a/src/client_server/message.rs b/src/client_server/message.rs
index 899c45a2..9705e4c0 100644
--- a/src/client_server/message.rs
+++ b/src/client_server/message.rs
@@ -242,7 +242,7 @@ pub async fn get_message_events_route(
             &sender_user,
             &sender_device,
             &body.room_id,
-            lazy_loaded.into_iter().collect(),
+            lazy_loaded,
             next_token,
         );
     }
diff --git a/src/client_server/sync.rs b/src/client_server/sync.rs
index a41e728e..c2014403 100644
--- a/src/client_server/sync.rs
+++ b/src/client_server/sync.rs
@@ -374,7 +374,7 @@ async fn sync_helper(
             let current_state_ids = db.rooms.state_full_ids(current_shortstatehash)?;
 
             let mut state_events = Vec::new();
-            let mut lazy_loaded = Vec::new();
+            let mut lazy_loaded = HashSet::new();
 
             for (shortstatekey, id) in current_state_ids {
                 let (event_type, state_key) = db.rooms.get_statekey_from_short(shortstatekey)?;
@@ -399,7 +399,7 @@ async fn sync_helper(
                             continue;
                         }
                     };
-                    lazy_loaded.push(
+                    lazy_loaded.insert(
                         UserId::parse(state_key.as_ref())
                             .expect("they are in timeline_users, so they should be correct"),
                     );
@@ -456,46 +456,57 @@ async fn sync_helper(
             let since_state_ids = db.rooms.state_full_ids(since_shortstatehash)?;
 
             let mut state_events = Vec::new();
-            let mut lazy_loaded = Vec::new();
+            let mut lazy_loaded = HashSet::new();
 
             for (key, id) in current_state_ids {
-                let pdu = match db.rooms.get_pdu(&id)? {
-                    Some(pdu) => pdu,
-                    None => {
-                        error!("Pdu in state not found: {}", id);
-                        continue;
-                    }
-                };
-
-                let state_key = pdu
-                    .state_key
-                    .as_ref()
-                    .expect("state events have state keys");
+                if body.full_state || since_state_ids.get(&key) != Some(&id) {
+                    let pdu = match db.rooms.get_pdu(&id)? {
+                        Some(pdu) => pdu,
+                        None => {
+                            error!("Pdu in state not found: {}", id);
+                            continue;
+                        }
+                    };
 
-                if pdu.kind != EventType::RoomMember {
-                    if body.full_state || since_state_ids.get(&key) != Some(&id) {
-                        state_events.push(pdu);
+                    if pdu.kind == EventType::RoomMember {
+                        match UserId::parse(
+                            pdu.state_key
+                                .as_ref()
+                                .expect("State event has state key")
+                                .clone(),
+                        ) {
+                            Ok(state_key_userid) => {
+                                lazy_loaded.insert(state_key_userid);
+                            }
+                            Err(e) => error!("Invalid state key for member event: {}", e),
+                        }
                     }
-                    continue;
-                }
-
-                // Pdu has to be a member event
-                let state_key_userid = UserId::parse(state_key.as_ref())
-                    .expect("they are in timeline_users, so they should be correct");
 
-                if body.full_state || since_state_ids.get(&key) != Some(&id) {
-                    lazy_loaded.push(state_key_userid);
                     state_events.push(pdu);
-                } else if timeline_users.contains(state_key)
-                    && (!db.rooms.lazy_load_was_sent_before(
+                }
+                for (_, event) in &timeline_pdus {
+                    if lazy_loaded.contains(&event.sender) {
+                        continue;
+                    }
+
+                    if !db.rooms.lazy_load_was_sent_before(
                         &sender_user,
                         &sender_device,
                         &room_id,
-                        &state_key_userid,
-                    )? || lazy_load_send_redundant)
-                {
-                    lazy_loaded.push(state_key_userid);
-                    state_events.push(pdu);
+                        &event.sender,
+                    )? || lazy_load_send_redundant
+                    {
+                        let pdu = match db.rooms.get_pdu(&id)? {
+                            Some(pdu) => pdu,
+                            None => {
+                                error!("Pdu in state not found: {}", id);
+                                continue;
+                            }
+                        };
+
+                        lazy_loaded.insert(event.sender.clone());
+                        state_events.push(pdu);
+                    }
                 }
             }
 
diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index b957b55d..600f46df 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -120,7 +120,7 @@ pub struct Rooms {
     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>>>>,
+        Mutex<HashMap<(Box<UserId>, Box<DeviceId>, Box<RoomId>, u64), HashSet<Box<UserId>>>>,
     pub(super) stateinfo_cache: Mutex<
         LruCache<
             u64,
@@ -3495,7 +3495,7 @@ impl Rooms {
         user_id: &UserId,
         device_id: &DeviceId,
         room_id: &RoomId,
-        lazy_load: Vec<Box<UserId>>,
+        lazy_load: HashSet<Box<UserId>>,
         count: u64,
     ) {
         self.lazy_load_waiting.lock().unwrap().insert(