From 79d626330c5e20fb83b1b33df1bfb29f2404275c Mon Sep 17 00:00:00 2001
From: Nikita Vilunov <nikita@vilunov.me>
Date: Mon, 15 Apr 2024 01:46:34 +0200
Subject: [PATCH] load rooms at player actor start

---
 crates/lavina-core/src/player.rs    | 12 ++++++++++++
 crates/lavina-core/src/repo/mod.rs  |  2 +-
 crates/lavina-core/src/repo/user.rs | 15 +++++++++++++++
 crates/lavina-core/src/room.rs      | 25 ++++++++++++++++++++++---
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/crates/lavina-core/src/player.rs b/crates/lavina-core/src/player.rs
index 4beef24..4f1a699 100644
--- a/crates/lavina-core/src/player.rs
+++ b/crates/lavina-core/src/player.rs
@@ -284,8 +284,11 @@ impl Player {
         let player = Player {
             player_id,
             storage_id,
+            // connections are empty when the actor is just started
             connections: AnonTable::new(),
+            // room handlers will be loaded later in the started task
             my_rooms: HashMap::new(),
+            // TODO implement and load bans
             banned_from: HashSet::new(),
             rx,
             handle,
@@ -297,6 +300,15 @@ impl Player {
     }
 
     async fn main_loop(mut self) -> Self {
+        let rooms = self.storage.get_rooms_of_a_user(self.storage_id).await.unwrap();
+        for room_id in rooms {
+            let room = self.rooms.get_room(&room_id).await;
+            if let Some(room) = room {
+                self.my_rooms.insert(room_id, room);
+            } else {
+                tracing::error!("Room #{room_id:?} not found");
+            }
+        }
         while let Some(cmd) = self.rx.recv().await {
             match cmd {
                 ActorCommand::AddConnection { sender, promise } => {
diff --git a/crates/lavina-core/src/repo/mod.rs b/crates/lavina-core/src/repo/mod.rs
index cb94d43..5faa7c4 100644
--- a/crates/lavina-core/src/repo/mod.rs
+++ b/crates/lavina-core/src/repo/mod.rs
@@ -51,7 +51,7 @@ impl Storage {
         Ok(res)
     }
 
-    pub async fn retrieve_room_by_name(&mut self, name: &str) -> Result<Option<StoredRoom>> {
+    pub async fn retrieve_room_by_name(&self, name: &str) -> Result<Option<StoredRoom>> {
         let mut executor = self.conn.lock().await;
         let res = sqlx::query_as(
             "select id, name, topic, message_count
diff --git a/crates/lavina-core/src/repo/user.rs b/crates/lavina-core/src/repo/user.rs
index bfa34e7..d836b8f 100644
--- a/crates/lavina-core/src/repo/user.rs
+++ b/crates/lavina-core/src/repo/user.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
 
 use crate::repo::Storage;
+use crate::room::RoomId;
 
 impl Storage {
     pub async fn retrieve_user_id_by_name(&self, name: &str) -> Result<Option<u32>> {
@@ -12,4 +13,18 @@ impl Storage {
 
         Ok(res.map(|(id,)| id))
     }
+
+    pub async fn get_rooms_of_a_user(&self, user_id: u32) -> Result<Vec<RoomId>> {
+        let mut executor = self.conn.lock().await;
+        let res: Vec<(String,)> = sqlx::query_as(
+            "select r.name
+            from memberships m inner join rooms r on m.room_id = r.id
+            where m.user_id = ?;",
+        )
+        .bind(user_id)
+        .fetch_all(&mut *executor)
+        .await?;
+
+        res.into_iter().map(|(room_id,)| RoomId::from(room_id)).collect()
+    }
 }
diff --git a/crates/lavina-core/src/room.rs b/crates/lavina-core/src/room.rs
index 58f5140..66507c0 100644
--- a/crates/lavina-core/src/room.rs
+++ b/crates/lavina-core/src/room.rs
@@ -91,9 +91,28 @@ impl RoomRegistry {
     }
 
     pub async fn get_room(&self, room_id: &RoomId) -> Option<RoomHandle> {
-        let inner = self.0.read().await;
-        let res = inner.rooms.get(room_id);
-        res.map(|r| r.clone())
+        let mut inner = self.0.write().await;
+        if let Some(room_handle) = inner.rooms.get(room_id) {
+            Some(room_handle.clone())
+        } else if let Some(stored_room) = inner.storage.retrieve_room_by_name(&*room_id.0).await.unwrap() {
+            // room exists, but was not loaded
+            log::debug!("Loading room {}...", &room_id.0);
+            let room = Room {
+                storage_id: stored_room.id,
+                room_id: room_id.clone(),
+                subscriptions: HashMap::new(), // TODO figure out how to populate subscriptions
+                members: HashSet::new(),       // TODO load members from storage
+                topic: stored_room.topic.into(),
+                message_count: stored_room.message_count,
+                storage: inner.storage.clone(),
+            };
+            let room_handle = RoomHandle(Arc::new(AsyncRwLock::new(room)));
+            inner.rooms.insert(room_id.clone(), room_handle.clone());
+            inner.metric_active_rooms.inc();
+            Some(room_handle)
+        } else {
+            None
+        }
     }
 
     pub async fn get_all_rooms(&self) -> Vec<RoomInfo> {