oxspeak_server/src/services/realtime_registry.rs

use crate::domain::events::channel_permission::{
    ChannelUserPermissionCreatedEvent, ChannelUserPermissionDeletedEvent,
    ChannelUserPermissionUpdatedEvent,
};
use crate::domain::events::server_permission::ServerUserPermissionUpdatedEvent;
use crate::models::{channel, channel_user, computed_permission::PermissionScopeType};
use crate::permissions::ChannelPermission;
use crate::repositories::Repositories;
use event_bus::EventBus;
use parking_lot::RwLock;
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use uuid::Uuid;

/// In-memory index of the users that can receive events for each channel.
#[derive(Debug, Default)]
pub struct RealtimeRegistry {
    channel_users: RwLock<HashMap<Uuid, HashSet<Uuid>>>,
    user_channels: RwLock<HashMap<Uuid, HashSet<Uuid>>>,
}

impl RealtimeRegistry {
    /// Rebuilds one channel audience after a committed structural change.
    pub async fn refresh_channel(
        &self,
        repositories: &Repositories,
        channel_id: Uuid,
    ) -> anyhow::Result<()> {
        let channel = channel::Entity::find_by_id(channel_id)
            .one(&repositories.channel.context.db)
            .await?
            .ok_or_else(|| anyhow::anyhow!("Channel not found"))?;

        if channel.channel_type == channel::ChannelType::DM {
            let users = channel_user::Entity::find()
                .filter(channel_user::Column::ChannelId.eq(channel_id))
                .all(&repositories.channel.context.db)
                .await?
                .into_iter()
                .map(|member| member.user_id);
            self.set_channel_users(channel_id, users);
        } else {
            let permissions = repositories.computed_permission.get_all().await?;
            self.set_channel_users(
                channel_id,
                permissions.into_iter().filter_map(|permission| {
                    (permission.scope_type == PermissionScopeType::Channel
                        && permission.resource_id == channel_id
                        && ChannelPermission::from_bits_retain(permission.permissions as u64)
                            .contains(ChannelPermission::READ_CHANNEL))
                    .then_some(permission.user_id)
                }),
            );
        }
        Ok(())
    }

    async fn refresh_user(&self, repositories: &Repositories, user_id: Uuid) -> anyhow::Result<()> {
        let channels = repositories
            .computed_permission
            .get_all()
            .await?
            .into_iter()
            .filter(|permission| {
                permission.user_id == user_id
                    && permission.scope_type == PermissionScopeType::Channel
                    && ChannelPermission::from_bits_retain(permission.permissions as u64)
                        .contains(ChannelPermission::READ_CHANNEL)
            })
            .map(|permission| permission.resource_id);
        self.set_user_channels(user_id, channels);
        Ok(())
    }

    pub async fn initialize(&self, repositories: &Repositories) -> anyhow::Result<()> {
        let permissions = repositories.computed_permission.get_all().await?;
        let mut channel_users = HashMap::<Uuid, HashSet<Uuid>>::new();
        let mut user_channels = HashMap::<Uuid, HashSet<Uuid>>::new();

        for permission in permissions {
            if permission.scope_type != PermissionScopeType::Channel
                || !ChannelPermission::from_bits_retain(permission.permissions as u64)
                    .contains(ChannelPermission::READ_CHANNEL)
            {
                continue;
            }
            channel_users
                .entry(permission.resource_id)
                .or_default()
                .insert(permission.user_id);
            user_channels
                .entry(permission.user_id)
                .or_default()
                .insert(permission.resource_id);
        }

        // Les DM n'utilisent pas les permissions calculées : leur audience est
        // définie directement par la table d'appartenance channel_user.
        let dm_channels = channel::Entity::find()
            .filter(channel::Column::ChannelType.eq(channel::ChannelType::DM))
            .all(&repositories.channel.context.db)
            .await?;
        let dm_ids: Vec<_> = dm_channels.into_iter().map(|item| item.id).collect();
        if !dm_ids.is_empty() {
            let members = channel_user::Entity::find()
                .filter(channel_user::Column::ChannelId.is_in(dm_ids))
                .all(&repositories.channel.context.db)
                .await?;
            for member in members {
                channel_users
                    .entry(member.channel_id)
                    .or_default()
                    .insert(member.user_id);
                user_channels
                    .entry(member.user_id)
                    .or_default()
                    .insert(member.channel_id);
            }
        }

        *self.channel_users.write() = channel_users;
        *self.user_channels.write() = user_channels;
        Ok(())
    }

    pub fn users_for_channel(&self, channel_id: Uuid) -> HashSet<Uuid> {
        self.channel_users
            .read()
            .get(&channel_id)
            .cloned()
            .unwrap_or_default()
    }

    pub fn set_channel_users(&self, channel_id: Uuid, users: impl IntoIterator<Item = Uuid>) {
        let users: HashSet<_> = users.into_iter().collect();
        let old = {
            let mut by_channel = self.channel_users.write();
            by_channel
                .insert(channel_id, users.clone())
                .unwrap_or_default()
        };
        let mut by_user = self.user_channels.write();
        for user_id in old.difference(&users) {
            if let Some(channels) = by_user.get_mut(user_id) {
                channels.remove(&channel_id);
                if channels.is_empty() {
                    by_user.remove(user_id);
                }
            }
        }
        for user_id in users {
            by_user.entry(user_id).or_default().insert(channel_id);
        }
    }

    pub fn set_user_channels(&self, user_id: Uuid, channels: impl IntoIterator<Item = Uuid>) {
        let channels: HashSet<_> = channels.into_iter().collect();
        let old = self
            .user_channels
            .write()
            .insert(user_id, channels.clone())
            .unwrap_or_default();
        let mut by_channel = self.channel_users.write();
        for channel_id in old.difference(&channels) {
            if let Some(users) = by_channel.get_mut(channel_id) {
                users.remove(&user_id);
                if users.is_empty() {
                    by_channel.remove(channel_id);
                }
            }
        }
        for channel_id in channels {
            by_channel.entry(channel_id).or_default().insert(user_id);
        }
    }

    pub fn remove_user(&self, user_id: Uuid) {
        if let Some(channels) = self.user_channels.write().remove(&user_id) {
            let mut by_channel = self.channel_users.write();
            for channel_id in channels {
                if let Some(users) = by_channel.get_mut(&channel_id) {
                    users.remove(&user_id);
                    if users.is_empty() {
                        by_channel.remove(&channel_id);
                    }
                }
            }
        }
    }

    pub fn remove_channel(&self, channel_id: Uuid) {
        if let Some(users) = self.channel_users.write().remove(&channel_id) {
            let mut by_user = self.user_channels.write();
            for user_id in users {
                if let Some(channels) = by_user.get_mut(&user_id) {
                    channels.remove(&channel_id);
                    if channels.is_empty() {
                        by_user.remove(&user_id);
                    }
                }
            }
        }
    }

    pub fn start_listening(
        self: &Arc<Self>,
        repositories: Arc<Repositories>,
        event_bus: Arc<EventBus>,
    ) {
        let registry = Arc::clone(self);
        event_bus.on_async_with::<ChannelUserPermissionUpdatedEvent, _>(
            repositories.clone(),
            move |repositories, event| {
                let user_id = event.user_id;
                let registry = Arc::clone(&registry);
                async move {
                    if let Err(error) = registry.refresh_user(&repositories, user_id).await {
                        tracing::error!(%user_id, ?error, "Unable to refresh realtime registry")
                    }
                }
            },
        );

        let registry = Arc::clone(self);
        event_bus.on_async_with::<ChannelUserPermissionCreatedEvent, _>(
            repositories.clone(),
            move |repositories, event| {
                let user_id = event.permission.user_id;
                let registry = Arc::clone(&registry);
                async move {
                    if let Err(error) = registry.refresh_user(&repositories, user_id).await {
                        tracing::error!(%user_id, ?error, "Unable to refresh realtime registry")
                    }
                }
            },
        );

        let registry = Arc::clone(self);
        event_bus.on_async_with::<ChannelUserPermissionDeletedEvent, _>(
            repositories.clone(),
            move |repositories, event| {
                let user_id = event.permission.user_id;
                let registry = Arc::clone(&registry);
                async move {
                    if let Err(error) = registry.refresh_user(&repositories, user_id).await {
                        tracing::error!(%user_id, ?error, "Unable to refresh realtime registry")
                    }
                }
            },
        );

        let registry = Arc::clone(self);
        let repositories = repositories.clone();
        event_bus.on_async_with::<ServerUserPermissionUpdatedEvent, _>(
            repositories,
            move |repositories, event| {
                let user_id = event.user_id;
                let registry = Arc::clone(&registry);
                async move {
                    if let Err(error) = registry.refresh_user(&repositories, user_id).await {
                        tracing::error!(%user_id, ?error, "Unable to refresh realtime registry")
                    }
                }
            },
        );
    }
}

#[cfg(test)]
mod tests {
    use super::RealtimeRegistry;
    use uuid::Uuid;

    #[test]
    fn indexes_dm_members_by_channel_and_user() {
        let registry = RealtimeRegistry::default();
        let channel_id = Uuid::new_v4();
        let first = Uuid::new_v4();
        let second = Uuid::new_v4();

        registry.set_channel_users(channel_id, [first, second]);

        assert_eq!(registry.users_for_channel(channel_id).len(), 2);
        assert!(
            registry
                .user_channels
                .read()
                .get(&first)
                .unwrap()
                .contains(&channel_id)
        );
        assert!(
            registry
                .user_channels
                .read()
                .get(&second)
                .unwrap()
                .contains(&channel_id)
        );

        registry.set_channel_users(channel_id, [second]);
        assert!(!registry.user_channels.read().contains_key(&first));
        assert!(registry.users_for_channel(channel_id).contains(&second));
    }
}