121 lines
4.3 KiB
Rust
121 lines
4.3 KiB
Rust
use std::time::Instant;
|
|
use uuid::Uuid;
|
|
|
|
use crate::error::RoomError;
|
|
use super::{RoomConnectionManager, CONNECTION_COOLDOWN, MAX_CONNECTIONS_PER_ROOM, ROOM_IDLE_TIMEOUT};
|
|
|
|
impl RoomConnectionManager {
|
|
pub async fn check_room_connection_rate(&self, room_id: Uuid, user_id: Uuid) -> Result<(), RoomError> {
|
|
let mut map = self.connection_rate.write().await;
|
|
let key = (room_id, user_id);
|
|
if let Some(last) = map.get(&key) {
|
|
if last.elapsed() < CONNECTION_COOLDOWN {
|
|
self.metrics.ws_rate_limit_hits.increment(1);
|
|
return Err(RoomError::RateLimited(format!(
|
|
"Connection cooldown active, retry in {}s",
|
|
CONNECTION_COOLDOWN.saturating_sub(last.elapsed()).as_secs()
|
|
)));
|
|
}
|
|
}
|
|
map.insert(key, Instant::now());
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn check_project_connection_rate(&self, project_id: Uuid, user_id: Uuid) -> Result<(), RoomError> {
|
|
let mut map = self.connection_rate.write().await;
|
|
let key = (project_id, user_id);
|
|
if let Some(last) = map.get(&key) {
|
|
if last.elapsed() < CONNECTION_COOLDOWN {
|
|
self.metrics.ws_rate_limit_hits.increment(1);
|
|
return Err(RoomError::RateLimited(format!(
|
|
"Connection cooldown active, retry in {}s",
|
|
CONNECTION_COOLDOWN.saturating_sub(last.elapsed()).as_secs()
|
|
)));
|
|
}
|
|
}
|
|
map.insert(key, Instant::now());
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn check_user_connection_rate(&self, user_id: Uuid) -> Result<(), RoomError> {
|
|
let mut map = self.connection_rate.write().await;
|
|
let key = (Uuid::nil(), user_id);
|
|
if let Some(last) = map.get(&key) {
|
|
if last.elapsed() < CONNECTION_COOLDOWN {
|
|
self.metrics.ws_rate_limit_hits.increment(1);
|
|
return Err(RoomError::RateLimited(format!(
|
|
"Connection cooldown active, retry in {}s",
|
|
CONNECTION_COOLDOWN.saturating_sub(last.elapsed()).as_secs()
|
|
)));
|
|
}
|
|
}
|
|
map.insert(key, Instant::now());
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn cleanup_rate_limit(&self) {
|
|
let mut map = self.connection_rate.write().await;
|
|
map.retain(|_, instant| instant.elapsed() < CONNECTION_COOLDOWN * 2);
|
|
const MAX_RATE_ENTRIES: usize = MAX_CONNECTIONS_PER_ROOM * 10;
|
|
if map.len() > MAX_RATE_ENTRIES {
|
|
let mut entries: Vec<_> = map.iter().collect();
|
|
entries.sort_by(|a, b| a.1.cmp(b.1));
|
|
let keep_count = entries.len() / 2;
|
|
let to_remove: Vec<_> = entries.into_iter().take(keep_count).map(|(k, _)| *k).collect();
|
|
for key in to_remove {
|
|
map.remove(&key);
|
|
}
|
|
}
|
|
drop(map);
|
|
self.cleanup_idle_rooms().await;
|
|
}
|
|
|
|
pub async fn cleanup_idle_rooms(&self) {
|
|
let now = Instant::now();
|
|
let activity = self.room_last_activity.read().await;
|
|
let idle_room_ids: Vec<Uuid> = activity
|
|
.iter()
|
|
.filter(|(_, last_time)| now.duration_since(**last_time) > ROOM_IDLE_TIMEOUT)
|
|
.map(|(room_id, _)| *room_id)
|
|
.collect();
|
|
drop(activity);
|
|
|
|
if idle_room_ids.is_empty() {
|
|
return;
|
|
}
|
|
|
|
{
|
|
let mut counts = self.room_subscriber_count.write().await;
|
|
let mut rooms = self.room_inner.write().await;
|
|
for room_id in &idle_room_ids {
|
|
if let Some(sender) = rooms.remove(room_id) {
|
|
let count = counts.remove(&room_id).unwrap_or(1);
|
|
self.metrics.users_online.decrement(count as f64);
|
|
drop(sender);
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
let mut stream_map = self.room_stream_inner.write().await;
|
|
for room_id in &idle_room_ids {
|
|
stream_map.remove(room_id);
|
|
}
|
|
}
|
|
|
|
{
|
|
let mut txs = self.room_shutdown_txs.write().await;
|
|
for room_id in &idle_room_ids {
|
|
txs.remove(room_id);
|
|
}
|
|
}
|
|
|
|
{
|
|
let mut activity = self.room_last_activity.write().await;
|
|
for room_id in &idle_room_ids {
|
|
activity.remove(room_id);
|
|
}
|
|
}
|
|
}
|
|
}
|