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 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>, pub stream_rx: broadcast::Receiver>, pub typing_rx: broadcast::Receiver>, } // ─── TransportSession ───────────────────────────────────────────────────────── pub struct TransportSession { pub user: WsUserCtx, pub subscriptions: Arc>, pub seq: Arc, pub service: Arc, last_heartbeat: Instant, last_activity: Instant, message_count: AtomicU32, rate_window_start: std::sync::Mutex, } impl TransportSession { pub fn new( user: WsUserCtx, service: Arc, ) -> 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 { self.seq.seq(room).await } pub async fn subscribe_room( &self, room_id: RoomId, ) -> Result { 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; } }