From 3eaac27789c336a4e2b5fc579416b00a166ddd18 Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <public@benj.me>
Date: Tue, 21 May 2024 11:08:59 +0200
Subject: [PATCH] sdk: introduce a `WeakClient` data structure

---
 crates/matrix-sdk/src/client/mod.rs      | 32 ++++++++++++++++++++++--
 crates/matrix-sdk/src/event_cache/mod.rs | 12 ++++-----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/crates/matrix-sdk/src/client/mod.rs b/crates/matrix-sdk/src/client/mod.rs
index 692c7b8b3..e253ff76f 100644
--- a/crates/matrix-sdk/src/client/mod.rs
+++ b/crates/matrix-sdk/src/client/mod.rs
@@ -19,7 +19,7 @@ use std::{
     fmt::{self, Debug},
     future::Future,
     pin::Pin,
-    sync::{Arc, Mutex as StdMutex, RwLock as StdRwLock},
+    sync::{Arc, Mutex as StdMutex, RwLock as StdRwLock, Weak},
 };
 
 use eyeball::{SharedObservable, Subscriber};
@@ -335,7 +335,10 @@ impl ClientInner {
         #[cfg(feature = "e2e-encryption")]
         client.e2ee.initialize_room_key_tasks(&client);
 
-        let _ = client.event_cache.get_or_init(|| async { EventCache::new(&client) }).await;
+        let _ = client
+            .event_cache
+            .get_or_init(|| async { EventCache::new(WeakClient::from_inner(&client)) })
+            .await;
 
         client
     }
@@ -2115,6 +2118,31 @@ impl Client {
     }
 }
 
+/// A weak reference to the inner client, useful when trying to get a handle
+/// on the owning client.
+#[derive(Clone)]
+pub(crate) struct WeakClient {
+    client: Weak<ClientInner>,
+}
+
+impl WeakClient {
+    /// Construct a [`WeakClient`] from a `Arc<ClientInner>`.
+    pub fn from_inner(client: &Arc<ClientInner>) -> Self {
+        Self { client: Arc::downgrade(client) }
+    }
+
+    /// Construct a [`WeakClient`] from a [`Client`].
+    #[cfg(test)]
+    pub fn from_client(client: &Client) -> Self {
+        Self::from_inner(&client.inner)
+    }
+
+    /// Attempts to get a [`Client`] from this [`WeakClient`].
+    pub fn get(&self) -> Option<Client> {
+        self.client.upgrade().map(|inner| Client { inner })
+    }
+}
+
 // The http mocking library is not supported for wasm32
 #[cfg(all(test, not(target_arch = "wasm32")))]
 pub(crate) mod tests {
diff --git a/crates/matrix-sdk/src/event_cache/mod.rs b/crates/matrix-sdk/src/event_cache/mod.rs
index 9ae2e5727..0e649d392 100644
--- a/crates/matrix-sdk/src/event_cache/mod.rs
+++ b/crates/matrix-sdk/src/event_cache/mod.rs
@@ -44,7 +44,7 @@
 use std::{
     collections::BTreeMap,
     fmt::Debug,
-    sync::{Arc, OnceLock, Weak},
+    sync::{Arc, OnceLock},
 };
 
 use eyeball::Subscriber;
@@ -69,7 +69,7 @@ use self::{
     paginator::{Paginator, PaginatorError},
     store::{Gap, RoomEvents},
 };
-use crate::{client::ClientInner, Client, Room};
+use crate::{client::WeakClient, Client, Room};
 
 mod linked_chunk;
 mod pagination;
@@ -159,10 +159,10 @@ impl Debug for EventCache {
 
 impl EventCache {
     /// Create a new [`EventCache`] for the given client.
-    pub(crate) fn new(client: &Arc<ClientInner>) -> Self {
+    pub(crate) fn new(client: WeakClient) -> Self {
         Self {
             inner: Arc::new(EventCacheInner {
-                client: Arc::downgrade(client),
+                client,
                 multiple_room_updates_lock: Default::default(),
                 by_room: Default::default(),
                 drop_handles: Default::default(),
@@ -297,7 +297,7 @@ impl EventCache {
 struct EventCacheInner {
     /// A weak reference to the inner client, useful when trying to get a handle
     /// on the owning client.
-    client: Weak<ClientInner>,
+    client: WeakClient,
 
     /// A lock used when many rooms must be updated at once.
     ///
@@ -316,7 +316,7 @@ struct EventCacheInner {
 
 impl EventCacheInner {
     fn client(&self) -> Result<Client> {
-        Ok(Client { inner: self.client.upgrade().ok_or(EventCacheError::ClientDropped)? })
+        self.client.get().ok_or(EventCacheError::ClientDropped)
     }
 
     /// Clears all the room's data.