lavina/crates/lavina-core/src/room.rs

//! Domain of rooms — chats with multiple participants.
use std::{collections::HashMap, hash::Hash, sync::Arc};

use prometheus::{IntGauge, Registry as MetricRegistry};
use serde::Serialize;
use tokio::sync::RwLock as AsyncRwLock;

use crate::player::{PlayerHandle, PlayerId, Updates};
use crate::prelude::*;
use crate::repo::Storage;

/// Opaque room id
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
pub struct RoomId(Str);
impl RoomId {
    pub fn from(str: impl Into<Str>) -> Result<RoomId> {
        let bytes = str.into();
        if bytes.len() > 32 {
            return Err(anyhow::Error::msg("Room name cannot be longer than 32 symbols"));
        }
        if bytes.contains(' ') {
            return Err(anyhow::Error::msg("Room name cannot contain spaces"));
        }
        Ok(RoomId(bytes))
    }
    pub fn as_inner(&self) -> &Str {
        &self.0
    }
    pub fn into_inner(self) -> Str {
        self.0
    }
}

/// Shared datastructure for storing metadata about rooms.
#[derive(Clone)]
pub struct RoomRegistry(Arc<AsyncRwLock<RoomRegistryInner>>);
impl RoomRegistry {
    pub fn new(metrics: &mut MetricRegistry, storage: Storage) -> Result<RoomRegistry> {
        let metric_active_rooms = IntGauge::new("chat_rooms_active", "Number of alive room actors")?;
        metrics.register(Box::new(metric_active_rooms.clone()))?;
        let inner = RoomRegistryInner {
            rooms: HashMap::new(),
            metric_active_rooms,
            storage,
        };
        Ok(RoomRegistry(Arc::new(AsyncRwLock::new(inner))))
    }

    pub async fn get_or_create_room(&mut self, room_id: RoomId) -> Result<RoomHandle> {
        let mut inner = self.0.write().await;
        if let Some(room_handle) = inner.rooms.get(&room_id) {
            // room was already loaded into memory
            log::debug!("Room {} was loaded already", &room_id.0);
            Ok(room_handle.clone())
        } else if let Some(stored_room) = inner.storage.retrieve_room_by_name(&*room_id.0).await? {
            // 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
                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, room_handle.clone());
            inner.metric_active_rooms.inc();
            Ok(room_handle)
        } else {
            // room does not exist, create it and load
            log::debug!("Creating room {}...", &room_id.0);
            let topic = "New room";
            let id = inner.storage.create_new_room(&*room_id.0, &*topic).await?;
            let room = Room {
                storage_id: id,
                room_id: room_id.clone(),
                subscriptions: HashMap::new(),
                topic: topic.into(),
                message_count: 0,
                storage: inner.storage.clone(),
            };
            let room_handle = RoomHandle(Arc::new(AsyncRwLock::new(room)));
            inner.rooms.insert(room_id, room_handle.clone());
            inner.metric_active_rooms.inc();
            Ok(room_handle)
        }
    }

    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())
    }

    pub async fn get_all_rooms(&self) -> Vec<RoomInfo> {
        let handles = {
            let inner = self.0.read().await;
            let handles = inner.rooms.values().cloned().collect::<Vec<_>>();
            handles
        };
        let mut res = vec![];
        for i in handles {
            res.push(i.get_room_info().await)
        }
        res
    }
}

struct RoomRegistryInner {
    rooms: HashMap<RoomId, RoomHandle>,
    metric_active_rooms: IntGauge,
    storage: Storage,
}

#[derive(Clone)]
pub struct RoomHandle(Arc<AsyncRwLock<Room>>);
impl RoomHandle {
    pub async fn subscribe(&self, player_id: PlayerId, player_handle: PlayerHandle) {
        let mut lock = self.0.write().await;
        lock.add_subscriber(player_id, player_handle).await;
    }

    pub async fn unsubscribe(&self, player_id: &PlayerId) {
        let mut lock = self.0.write().await;
        lock.subscriptions.remove(player_id);
        let update = Updates::RoomLeft {
            room_id: lock.room_id.clone(),
            former_member_id: player_id.clone(),
        };
        lock.broadcast_update(update, player_id).await;
    }

    pub async fn send_message(&self, player_id: PlayerId, body: Str) {
        let mut lock = self.0.write().await;
        let res = lock.send_message(player_id, body).await;
        if let Err(err) = res {
            log::warn!("Failed to send message: {err:?}");
        }
    }

    pub async fn get_room_info(&self) -> RoomInfo {
        let lock = self.0.read().await;
        RoomInfo {
            id: lock.room_id.clone(),
            members: lock.subscriptions.keys().map(|x| x.clone()).collect::<Vec<_>>(),
            topic: lock.topic.clone(),
        }
    }

    pub async fn set_topic(&mut self, changer_id: PlayerId, new_topic: Str) {
        let mut lock = self.0.write().await;
        lock.topic = new_topic.clone();
        let update = Updates::RoomTopicChanged {
            room_id: lock.room_id.clone(),
            new_topic: new_topic.clone(),
        };
        lock.broadcast_update(update, &changer_id).await;
    }
}

struct Room {
    storage_id: u32,
    room_id: RoomId,
    subscriptions: HashMap<PlayerId, PlayerHandle>,
    message_count: u32,
    topic: Str,
    storage: Storage,
}
impl Room {
    async fn add_subscriber(&mut self, player_id: PlayerId, player_handle: PlayerHandle) {
        tracing::info!("Adding a subscriber to room");
        self.subscriptions.insert(player_id.clone(), player_handle);
        let update = Updates::RoomJoined {
            room_id: self.room_id.clone(),
            new_member_id: player_id.clone(),
        };
        self.broadcast_update(update, &player_id).await;
    }

    async fn send_message(&mut self, author_id: PlayerId, body: Str) -> Result<()> {
        tracing::info!("Adding a message to room");
        self.storage.insert_message(self.storage_id, self.message_count, &body, &*author_id.as_inner()).await?;
        self.message_count += 1;
        let update = Updates::NewMessage {
            room_id: self.room_id.clone(),
            author_id: author_id.clone(),
            body,
        };
        self.broadcast_update(update, &author_id).await;
        Ok(())
    }

    async fn broadcast_update(&self, update: Updates, except: &PlayerId) {
        tracing::debug!("Broadcasting an update to {} subs", self.subscriptions.len());
        for (player_id, sub) in &self.subscriptions {
            if player_id == except {
                continue;
            }
            log::info!("Sending a message from room to player");
            sub.update(update.clone()).await;
        }
    }
}

#[derive(Serialize)]
pub struct RoomInfo {
    pub id: RoomId,
    pub members: Vec<PlayerId>,
    pub topic: Str,
}