gitdataai/libs/transport/handler/session.rs

use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::{Duration, Instant};

use dashmap::DashMap;
use tokio::sync::broadcast;

use models::RoomId;
use queue::{RoomMessageEvent, RoomMessageStreamChunkEvent, TypingEvent};
use service::AppService;

use crate::error::AppTransportError;
use crate::seq::SeqAllocator;
use crate::token::AppTransportTokenContext;

// ─── Constants ────────────────────────────────────────────────────────────────

pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
pub const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(60);
pub const MAX_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
pub const RATE_LIMIT_WINDOW: Duration = Duration::from_secs(1);
pub const MAX_MESSAGES_PER_SECOND: u32 = 1000;
pub const MAX_TEXT_MESSAGE_LEN: usize = 64 * 1024;

// ─── User Context ─────────────────────────────────────────────────────────────

#[derive(Clone)]
pub struct WsUserCtx {
    pub user_id: uuid::Uuid,
    pub device_id: String,
    pub client_id: String,
}

impl From<AppTransportTokenContext> for WsUserCtx {
    fn from(ctx: AppTransportTokenContext) -> Self {
        Self {
            user_id: ctx.user_id,
            device_id: ctx.device_id,
            client_id: ctx.client_id,
        }
    }
}

// ─── Per-Room Subscription ────────────────────────────────────────────────────

pub struct RoomSubscription {
    pub room_id: RoomId,
    pub msg_rx: broadcast::Receiver<Arc<RoomMessageEvent>>,
    pub stream_rx: broadcast::Receiver<Arc<RoomMessageStreamChunkEvent>>,
    pub typing_rx: broadcast::Receiver<Arc<TypingEvent>>,
}

// ─── TransportSession ─────────────────────────────────────────────────────────

pub struct TransportSession {
    pub user: WsUserCtx,
    pub subscriptions: Arc<DashMap<RoomId, RoomSubscription>>,
    pub seq: Arc<SeqAllocator>,
    pub service: Arc<AppService>,
    last_heartbeat: Instant,
    last_activity: Instant,
    message_count: AtomicU32,
    rate_window_start: std::sync::Mutex<Instant>,
}

impl TransportSession {
    pub fn new(
        user: WsUserCtx,
        service: Arc<AppService>,
    ) -> Self {
        Self {
            user,
            subscriptions: Arc::new(DashMap::new()),
            seq: Arc::new(SeqAllocator::new(service.cache.clone(), service.db.clone())),
            service,
            last_heartbeat: Instant::now(),
            last_activity: Instant::now(),
            message_count: AtomicU32::new(0),
            rate_window_start: std::sync::Mutex::new(Instant::now()),
        }
    }

    pub fn touch_heartbeat(&mut self) {
        self.last_heartbeat = Instant::now();
    }

    pub fn touch_activity(&mut self) {
        self.last_activity = Instant::now();
        self.last_heartbeat = Instant::now();
    }

    pub fn heartbeat_elapsed(&self) -> Duration {
        self.last_heartbeat.elapsed()
    }

    pub fn activity_elapsed(&self) -> Duration {
        self.last_activity.elapsed()
    }

    pub fn check_rate_limit(&self) -> bool {
        let mut start = match self.rate_window_start.lock() {
            Ok(guard) => guard,
            Err(poisoned) => {
                tracing::warn!("rate_window_start mutex poisoned, recovering");
                poisoned.into_inner()
            }
        };
        if start.elapsed() > RATE_LIMIT_WINDOW {
            self.message_count.store(0, Ordering::Relaxed);
            *start = Instant::now();
        }
        self.message_count.fetch_add(1, Ordering::Relaxed) >= MAX_MESSAGES_PER_SECOND
    }

    pub async fn next_seq(&self, room: RoomId) -> Result<i64, AppTransportError> {
        self.seq.seq(room).await
    }

    pub async fn subscribe_room(
        &self,
        room_id: RoomId,
    ) -> Result<RoomSubscription, AppTransportError> {
        let manager = &self.service.room.room_manager;
        let rx = manager
            .subscribe(room_id, self.user.user_id)
            .await
            .map_err(|_| AppTransportError::Internal)?;
        let stream_rx = manager.subscribe_room_stream(room_id).await;
        let typing_rx = manager.subscribe_typing(room_id).await;
        Ok(RoomSubscription {
            room_id,
            msg_rx: rx,
            stream_rx,
            typing_rx,
        })
    }

    pub async fn unsubscribe_room(&self, room_id: RoomId) {
        self.service.room.room_manager.unsubscribe(room_id, self.user.user_id).await;
        self.subscriptions.remove(&room_id);
    }

    pub async fn broadcast_typing(&self, room_id: RoomId, action: &str) {
        let event = TypingEvent {
            room_id,
            user_id: self.user.user_id,
            username: self.user.user_id.to_string(),
            avatar_url: None,
            action: action.to_string(),
            sender_type: Some("user".to_string()),
        };
        self.service.room.room_manager.broadcast_typing(room_id, event).await;
    }
}