slightly better sliding sync

src/service/users/mod.rs

@@ -1,20 +1,36 @@
 mod data;
-use std::{collections::BTreeMap, mem};
+use std::{
+    collections::BTreeMap,
+    mem,
+    sync::{Arc, Mutex},
+};
 
 pub use data::Data;
 use ruma::{
-    api::client::{device::Device, error::ErrorKind, filter::FilterDefinition},
+    api::client::{
+        device::Device,
+        error::ErrorKind,
+        filter::FilterDefinition,
+        sync::sync_events::{self, v4::SyncRequestList},
+    },
     encryption::{CrossSigningKey, DeviceKeys, OneTimeKey},
     events::AnyToDeviceEvent,
     serde::Raw,
     DeviceId, DeviceKeyAlgorithm, DeviceKeyId, OwnedDeviceId, OwnedDeviceKeyId, OwnedMxcUri,
-    OwnedUserId, RoomAliasId, UInt, UserId,
+    OwnedRoomId, OwnedUserId, RoomAliasId, UInt, UserId,
 };
 
 use crate::{services, Error, Result};
 
+pub struct SlidingSyncCache {
+    lists: BTreeMap<String, SyncRequestList>,
+    known_rooms: BTreeMap<String, BTreeMap<OwnedRoomId, bool>>,
+}
+
 pub struct Service {
     pub db: &'static dyn Data,
+    pub connections:
+        Mutex<BTreeMap<(OwnedUserId, OwnedDeviceId, String), Arc<Mutex<SlidingSyncCache>>>>,
 }
 
 impl Service {
@@ -23,6 +39,121 @@ impl Service {
         self.db.exists(user_id)
     }
 
+    pub fn forget_sync_request_connection(
+        &self,
+        user_id: OwnedUserId,
+        device_id: OwnedDeviceId,
+        conn_id: String,
+    ) {
+        self.connections
+            .lock()
+            .unwrap()
+            .remove(&(user_id, device_id, conn_id));
+    }
+
+    pub fn update_sync_request_with_cache(
+        &self,
+        user_id: OwnedUserId,
+        device_id: OwnedDeviceId,
+        request: &mut sync_events::v4::Request,
+    ) -> BTreeMap<String, BTreeMap<OwnedRoomId, bool>> {
+        let Some(conn_id) = request.conn_id.clone() else { return BTreeMap::new(); };
+
+        let cache = &mut self.connections.lock().unwrap();
+        let cached = Arc::clone(
+            cache
+                .entry((user_id, device_id, conn_id))
+                .or_insert_with(|| {
+                    Arc::new(Mutex::new(SlidingSyncCache {
+                        lists: BTreeMap::new(),
+                        known_rooms: BTreeMap::new(),
+                    }))
+                }),
+        );
+        let cached = &mut cached.lock().unwrap();
+        drop(cache);
+
+        for (list_id, list) in &mut request.lists {
+            if let Some(cached_list) = cached.lists.remove(list_id) {
+                if list.sort.is_empty() {
+                    list.sort = cached_list.sort;
+                };
+                if list.room_details.required_state.is_empty() {
+                    list.room_details.required_state = cached_list.room_details.required_state;
+                };
+                list.room_details.timeline_limit = list
+                    .room_details
+                    .timeline_limit
+                    .or(cached_list.room_details.timeline_limit);
+                list.include_old_rooms = list
+                    .include_old_rooms
+                    .clone()
+                    .or(cached_list.include_old_rooms);
+                match (&mut list.filters, cached_list.filters) {
+                    (Some(list_filters), Some(cached_filters)) => {
+                        list_filters.is_dm = list_filters.is_dm.or(cached_filters.is_dm);
+                        if list_filters.spaces.is_empty() {
+                            list_filters.spaces = cached_filters.spaces;
+                        }
+                        list_filters.is_encrypted =
+                            list_filters.is_encrypted.or(cached_filters.is_encrypted);
+                        list_filters.is_invite =
+                            list_filters.is_invite.or(cached_filters.is_invite);
+                        if list_filters.room_types.is_empty() {
+                            list_filters.room_types = cached_filters.room_types;
+                        }
+                        if list_filters.not_room_types.is_empty() {
+                            list_filters.not_room_types = cached_filters.not_room_types;
+                        }
+                        list_filters.room_name_like = list_filters
+                            .room_name_like
+                            .clone()
+                            .or(cached_filters.room_name_like);
+                        if list_filters.tags.is_empty() {
+                            list_filters.tags = cached_filters.tags;
+                        }
+                        if list_filters.not_tags.is_empty() {
+                            list_filters.not_tags = cached_filters.not_tags;
+                        }
+                    }
+                    (_, Some(cached_filters)) => list.filters = Some(cached_filters),
+                    (_, _) => {}
+                }
+                if list.bump_event_types.is_empty() {
+                    list.bump_event_types = cached_list.bump_event_types;
+                };
+            }
+            cached.lists.insert(list_id.clone(), list.clone());
+        }
+
+        cached.known_rooms.clone()
+    }
+
+    pub fn update_sync_known_rooms(
+        &self,
+        user_id: OwnedUserId,
+        device_id: OwnedDeviceId,
+        conn_id: String,
+        list_id: String,
+        new_cached_rooms: BTreeMap<OwnedRoomId, bool>,
+    ) {
+        let cache = &mut self.connections.lock().unwrap();
+        let cached = Arc::clone(
+            cache
+                .entry((user_id, device_id, conn_id))
+                .or_insert_with(|| {
+                    Arc::new(Mutex::new(SlidingSyncCache {
+                        lists: BTreeMap::new(),
+                        known_rooms: BTreeMap::new(),
+                    }))
+                }),
+        );
+        let cached = &mut cached.lock().unwrap();
+        drop(cache);
+
+        cached.known_rooms.insert(list_id, new_cached_rooms);
+    }
+
     /// Check if account is deactivated
     pub fn is_deactivated(&self, user_id: &UserId) -> Result<bool> {
         self.db.is_deactivated(user_id)