Use an atomic int instead of abusing the strong count of an arc.

Author: bnjbvr

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

@@ -331,13 +331,14 @@ impl EventCache {
     /// The auto-shrink mechanism works this way:
     ///
     /// - Each time there's a new subscriber to a [`RoomEventCache`], it will
-    ///   get-or-create a new shared `Arc<()>`. When that subscriber is dropped,
-    ///   and the number of shared references is about to drop to 0, we notify
-    ///   this task (below) that this is the case, with the room id.
-    /// - This task here, owned by the [`EventCacheInner`], will listen to such
-    ///   notifications that a room may be shrunk; it will look at the number of
-    ///   shared references again, and will auto-shrink if the number of shared
-    ///   references is 0.
+    ///   increment the active number of listeners to that room, aka
+    ///   [`RoomEventCacheState::listener_count`].
+    /// - When that subscriber is dropped, it will decrement that count; and
+    ///   notify the task below if it reached 0.
+    /// - The task spawned here, owned by the [`EventCacheInner`], will listen
+    ///   to such notifications that a room may be shrunk. It will attempt an
+    ///   auto-shrink, by letting the inner state decide whether this is a good
+    ///   time to do so (new listeners might have spawned in the meanwhile).
     #[instrument(skip_all)]
     async fn auto_shrink_linked_chunk_task(
         inner: Arc<EventCacheInner>,

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

@@ -18,7 +18,10 @@ use std::{
     collections::BTreeMap,
     fmt,
     ops::{Deref, DerefMut},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
 use events::Gap;
@@ -75,18 +78,15 @@ pub struct RoomEventCacheListener {
     auto_shrink_sender: mpsc::Sender<AutoShrinkChannelPayload>,
 
     /// Shared instance of the auto-shrinker.
-    auto_shrinker: Arc<()>,
+    listener_count: Arc<AtomicUsize>,
 }
 
 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, such that there'll remain only the one
-            // in the `RoomEventCacheState`; notify the auto-shrink task that it
-            // may have to trigger for this room.
+        let previous_listener_count = self.listener_count.fetch_sub(1, Ordering::SeqCst);
+        if previous_listener_count == 1 {
+            // We were the last instance of the listener; let the auto-shrinker know by
+            // notifying it of our room id.
 
             let mut room_id = self.room_id.clone();
 
@@ -145,16 +145,17 @@ impl RoomEventCache {
     /// Subscribe to this room updates, after getting the initial list of
     /// events.
     pub async fn subscribe(&self) -> (Vec<TimelineEvent>, RoomEventCacheListener) {
-        let mut state = self.inner.state.write().await;
+        let state = self.inner.state.read().await;
         let events = state.events().events().map(|(_position, item)| item.clone()).collect();
 
-        let listener_count = state.auto_shrinker.get_or_insert_with(|| Arc::new(())).clone();
+        state.listener_count.fetch_add(1, Ordering::SeqCst);
+
         let recv = self.inner.sender.subscribe();
         let listener = RoomEventCacheListener {
             recv,
-            auto_shrinker: listener_count,
             room_id: self.inner.room_id.clone(),
             auto_shrink_sender: self.inner.auto_shrink_sender.clone(),
+            listener_count: state.listener_count.clone(),
         };
 
         (events, listener)
@@ -627,7 +628,7 @@ pub(super) enum LoadMoreEventsBackwardsOutcome {
 
 // Use a private module to hide `events` to this parent module.
 mod private {
-    use std::sync::Arc;
+    use std::sync::{atomic::AtomicUsize, Arc};
 
     use eyeball_im::VectorDiff;
     use matrix_sdk_base::{
@@ -669,10 +670,8 @@ mod private {
         /// 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<()>>,
+        /// An atomic count of the current number of listeners.
+        pub(super) listener_count: Arc<AtomicUsize>,
     }
 
     impl RoomEventCacheState {
@@ -729,7 +728,7 @@ mod private {
                 events,
                 deduplicator,
                 waited_for_initial_prev_token: false,
-                auto_shrinker: None,
+                listener_count: Default::default(),
             })
         }
 
@@ -911,21 +910,19 @@ mod private {
             Ok(Some(self.events.updates_as_vector_diffs()))
         }
 
-        /// Automatically shrink the room if there are no listeners.
+        /// Automatically shrink the room if there are no listeners, as
+        /// indicated by the atomic number of active listeners.
         #[must_use = "Updates as `VectorDiff` must probably be propagated via `RoomEventCacheUpdate`"]
         pub(crate) async fn auto_shrink_if_no_listeners(
             &mut self,
         ) -> Result<Option<Vec<VectorDiff<TimelineEvent>>>, EventCacheError> {
-            let mut diffs = None;
-            if let Some(auto_shrinker) = self.auto_shrinker.as_mut() {
+            if self.listener_count.load(std::sync::atomic::Ordering::SeqCst) == 0 {
                 // If we are the last strong reference to the auto-shrinker, we can shrink the
                 // events data structure to its last chunk.
-                if Arc::strong_count(auto_shrinker) == 1 {
-                    self.auto_shrinker = None;
-                    diffs = self.shrink_to_last_chunk().await?;
-                }
+                self.shrink_to_last_chunk().await
+            } else {
+                Ok(None)
             }
-            Ok(diffs)
         }
 
         /// Removes the bundled relations from an event, if they were present.
@@ -2203,7 +2200,7 @@ mod tests {
         {
             // 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());
+            assert_eq!(state.listener_count.load(std::sync::atomic::Ordering::SeqCst), 0);
         }
 
         // Getting the events will only give us the latest chunk.