feat(event cache): auto-shrink a room event cache's chunk after all listeners are left

Author: bnjbvr

crates/matrix-sdk/src/event_cache/room/mod.rs

@@ -14,7 +14,12 @@
 
 //! All event cache types for a single room.
 
-use std::{collections::BTreeMap, fmt, sync::Arc};
+use std::{
+    collections::BTreeMap,
+    fmt,
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
 
 use events::Gap;
 use eyeball_im::VectorDiff;
@@ -27,9 +32,12 @@ use ruma::{
     serde::Raw,
     EventId, OwnedEventId, OwnedRoomId, RoomVersionId,
 };
-use tokio::sync::{
-    broadcast::{Receiver, Sender},
-    Notify, RwLock,
+use tokio::{
+    spawn,
+    sync::{
+        broadcast::{Receiver, Sender},
+        Notify, RwLock,
+    },
 };
 use tracing::{trace, warn};
 
@@ -55,6 +63,84 @@ impl fmt::Debug for RoomEventCache {
     }
 }
 
+/// A RAII data structure that will automatically shrink the current room's
+/// linked chunk on drop.
+pub(super) struct AutoShrinkLinkedChunk {
+    /// The related room event cache, wrapped in an option so as to take it from
+    /// the drop implementation.
+    room_event_cache: Option<RoomEventCache>,
+}
+
+impl Drop for AutoShrinkLinkedChunk {
+    fn drop(&mut self) {
+        let room_event_cache = self.room_event_cache.take().expect("drop called multiple times");
+
+        // I hear you from the future: "but, spawning a detached task in a drop
+        // implementation is real bad! Maybe there will be multiple shrinks
+        // happening at the same time, and that's bad!". However, this can't
+        // happen, because the whole `state` variable is guarded by a fair lock,
+        // which will run queries in the order they happen. Should be fine™.
+        spawn(async move {
+            let mut state = room_event_cache.inner.state.write().await;
+
+            if let Err(err) = state.shrink_to_last_chunk().await {
+                warn!("error when trying to shrink the linked chunk: {err}");
+            }
+        });
+    }
+}
+
+/// Thin wrapper for a room event cache listener, so as to trigger side-effects
+/// when all listeners are gone.
+#[allow(missing_debug_implementations)]
+pub struct RoomEventCacheListener {
+    /// Underlying receiver of the room event cache's updates.
+    recv: Receiver<RoomEventCacheUpdate>,
+
+    /// Shared instance of the auto-shrinker.
+    auto_shrinker: Arc<AutoShrinkLinkedChunk>,
+}
+
+impl Drop for RoomEventCacheListener {
+    fn drop(&mut self) {
+        if Arc::strong_count(&self.auto_shrinker) == 2 {
+            // There are only two instances of the auto-shrinker: the one in the
+            // `RoomEventCacheState`, and the one here.
+
+            // The one here is about to be dropped. The only remaining one will thus be the
+            // one in the [`RoomEventCacheState`]. Reset it manually, so as to
+            // trigger its drop.
+
+            let room_event_cache = self
+                .auto_shrinker
+                .room_event_cache
+                .as_ref()
+                .expect("drop called multiple times")
+                .clone();
+
+            spawn(async move {
+                let mut state = room_event_cache.inner.state.write().await;
+                // Decrease the strong count of the auto-shrinker to 1.
+                state.auto_shrinker = None;
+            });
+        }
+    }
+}
+
+impl Deref for RoomEventCacheListener {
+    type Target = Receiver<RoomEventCacheUpdate>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.recv
+    }
+}
+
+impl DerefMut for RoomEventCacheListener {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.recv
+    }
+}
+
 impl RoomEventCache {
     /// Create a new [`RoomEventCache`] using the given room and store.
     pub(super) fn new(
@@ -77,11 +163,20 @@ impl RoomEventCache {
 
     /// Subscribe to this room updates, after getting the initial list of
     /// events.
-    pub async fn subscribe(&self) -> (Vec<TimelineEvent>, Receiver<RoomEventCacheUpdate>) {
-        let state = self.inner.state.read().await;
+    pub async fn subscribe(&self) -> (Vec<TimelineEvent>, RoomEventCacheListener) {
+        let mut state = self.inner.state.write().await;
         let events = state.events().events().map(|(_position, item)| item.clone()).collect();
 
-        (events, self.inner.sender.subscribe())
+        let listener_count = state
+            .auto_shrinker
+            .get_or_insert_with(|| {
+                Arc::new(AutoShrinkLinkedChunk { room_event_cache: Some(self.clone()) })
+            })
+            .clone();
+        let recv = self.inner.sender.subscribe();
+        let listener = RoomEventCacheListener { recv, auto_shrinker: listener_count };
+
+        (events, listener)
     }
 
     /// Return a [`RoomPagination`] API object useful for running
@@ -556,7 +651,7 @@ mod private {
     use ruma::{serde::Raw, OwnedEventId, OwnedRoomId};
     use tracing::{debug, error, instrument, trace};
 
-    use super::{events::RoomEvents, LoadMoreEventsBackwardsOutcome};
+    use super::{events::RoomEvents, AutoShrinkLinkedChunk, LoadMoreEventsBackwardsOutcome};
     use crate::event_cache::{deduplicator::Deduplicator, EventCacheError};
 
     /// State for a single room's event cache.
@@ -584,6 +679,11 @@ mod private {
         /// the first time we try to run backward pagination. We reset
         /// that upon clearing the timeline events.
         pub waited_for_initial_prev_token: bool,
+
+        /// A shared auto-shrinker, shared among all the live subscribers.
+        ///
+        /// If set to `None`, means there are no
+        pub(super) auto_shrinker: Option<Arc<AutoShrinkLinkedChunk>>,
     }
 
     impl RoomEventCacheState {
@@ -640,6 +740,7 @@ mod private {
                 events,
                 deduplicator,
                 waited_for_initial_prev_token: false,
+                auto_shrinker: None,
             })
         }
 
@@ -1986,4 +2087,124 @@ mod tests {
         assert_eq!(outcome.events[0].event_id().as_deref(), Some(evid1));
         assert!(outcome.reached_start);
     }
+
+    #[cfg(not(target_arch = "wasm32"))] // This uses the cross-process lock, so needs time support.
+    #[async_test]
+    async fn test_auto_shrink_after_all_subscribers_are_gone() {
+        use std::ops::Not as _;
+
+        use eyeball_im::VectorDiff;
+        use tokio::task::yield_now;
+
+        use crate::{assert_let_timeout, event_cache::RoomEventCacheUpdate};
+
+        let room_id = room_id!("!galette:saucisse.bzh");
+
+        let client = MockClientBuilder::new("http://localhost".to_owned()).build().await;
+
+        let f = EventFactory::new().room(room_id);
+
+        let evid1 = event_id!("$1");
+        let evid2 = event_id!("$2");
+
+        let ev1 = f.text_msg("hello world").sender(*ALICE).event_id(evid1).into_event();
+        let ev2 = f.text_msg("howdy").sender(*BOB).event_id(evid2).into_event();
+
+        // Fill the event cache store with an initial linked chunk with 2 events chunks.
+        {
+            let store = client.event_cache_store();
+            let store = store.lock().await.unwrap();
+            store
+                .handle_linked_chunk_updates(
+                    room_id,
+                    vec![
+                        Update::NewItemsChunk {
+                            previous: None,
+                            new: ChunkIdentifier::new(0),
+                            next: None,
+                        },
+                        Update::PushItems {
+                            at: Position::new(ChunkIdentifier::new(0), 0),
+                            items: vec![ev1],
+                        },
+                        Update::NewItemsChunk {
+                            previous: Some(ChunkIdentifier::new(0)),
+                            new: ChunkIdentifier::new(1),
+                            next: None,
+                        },
+                        Update::PushItems {
+                            at: Position::new(ChunkIdentifier::new(1), 0),
+                            items: vec![ev2],
+                        },
+                    ],
+                )
+                .await
+                .unwrap();
+        }
+
+        let event_cache = client.event_cache();
+        event_cache.subscribe().unwrap();
+        event_cache.enable_storage().unwrap();
+
+        client.base_client().get_or_create_room(room_id, matrix_sdk_base::RoomState::Joined);
+        let room = client.get_room(room_id).unwrap();
+        let (room_event_cache, _drop_handles) = room.event_cache().await.unwrap();
+
+        // Sanity check: lazily loaded, so only includes one item at start.
+        let (events1, mut stream1) = room_event_cache.subscribe().await;
+        assert_eq!(events1.len(), 1);
+        assert_eq!(events1[0].event_id().as_deref(), Some(evid2));
+        assert!(stream1.is_empty());
+
+        // Force loading the full linked chunk by back-paginating.
+        let outcome = room_event_cache.pagination().run_backwards_once(20).await.unwrap();
+        assert_eq!(outcome.events.len(), 1);
+        assert_eq!(outcome.events[0].event_id().as_deref(), Some(evid1));
+        assert!(outcome.reached_start.not());
+
+        // We also get an update about the loading from the store. Ignore it, for this
+        // test's sake.
+        assert_let_timeout!(
+            Ok(RoomEventCacheUpdate::UpdateTimelineEvents { diffs, .. }) = stream1.recv()
+        );
+        assert_eq!(diffs.len(), 1);
+        assert_matches!(&diffs[0], VectorDiff::Insert { index: 0, value } => {
+            assert_eq!(value.event_id().as_deref(), Some(evid1));
+        });
+
+        assert!(stream1.is_empty());
+
+        // Have another listener subscribe to the event cache.
+        // Since it's not the first one, and the previous one loaded some more events,
+        // the second listener seems them all.
+        let (events2, stream2) = room_event_cache.subscribe().await;
+        assert_eq!(events2.len(), 2);
+        assert_eq!(events2[0].event_id().as_deref(), Some(evid1));
+        assert_eq!(events2[1].event_id().as_deref(), Some(evid2));
+        assert!(stream2.is_empty());
+
+        // Drop the first stream, and wait a bit.
+        drop(stream1);
+        yield_now().await;
+
+        // The second stream remains undisturbed.
+        assert!(stream2.is_empty());
+
+        // Now drop the second stream, and wait a bit.
+        drop(stream2);
+        yield_now().await;
+
+        // The linked chunk must have auto-shrunk by now.
+
+        {
+            // Check the inner state: there's no more shared auto-shrinker.
+            let state = room_event_cache.inner.state.read().await;
+            assert!(state.auto_shrinker.is_none());
+        }
+
+        // Getting the events will only give us the latest chunk.
+        let (events3, _stream2) = room_event_cache.subscribe().await;
+        assert_eq!(events3.len(), 1);
+        assert_eq!(events3[0].event_id().as_deref(), Some(evid2));
+    }
 }