gitdataai/libs/room/src/connection/lifecycle.rs
2026-05-14 10:02:21 +08:00

139 lines
4.4 KiB
Rust

use tokio::sync::broadcast;
use uuid::Uuid;
use super::{RoomConnectionManager, SHUTDOWN_CHANNEL_CAPACITY};
impl RoomConnectionManager {
pub fn subscribe_shutdown(&self) -> broadcast::Receiver<()> {
self.shutdown_tx.subscribe()
}
pub fn trigger_shutdown(&self) {
let _ = self.shutdown_tx.send(());
}
pub async fn register_room(&self, room_id: Uuid) -> broadcast::Receiver<()> {
let mut txs = self.room_shutdown_txs.write().await;
if let Some(tx) = txs.get(&room_id) {
return tx.subscribe();
}
let (tx, rx) = broadcast::channel(SHUTDOWN_CHANNEL_CAPACITY);
txs.insert(room_id, tx);
rx
}
pub async fn shutdown_room(&self, room_id: Uuid) {
{
let txs = self.room_shutdown_txs.read().await;
if let Some(tx) = txs.get(&room_id) {
let _ = tx.send(());
}
}
crate::service::unmark_room_spawned(room_id);
{
let mut counts = self.room_subscriber_count.write().await;
let count = counts.remove(&room_id).unwrap_or(0) as f64;
if count > 0.0 {
self.metrics.users_online.decrement(count);
}
}
{
let mut map = self.room_inner.write().await;
map.remove(&room_id);
}
{
let mut stream_map = self.room_stream_inner.write().await;
stream_map.remove(&room_id);
}
// Remove all streams associated with this room from active_streams and room_to_streams.
{
let mut r2s = self.room_to_streams.write().await;
if let Some(stream_ids) = r2s.remove(&room_id) {
let mut active = self.active_streams.write().await;
for id in stream_ids {
active.remove(&id);
}
}
}
{
let mut txs = self.room_shutdown_txs.write().await;
txs.remove(&room_id);
}
}
pub async fn prune_stale_rooms(&self, active_room_ids: &[Uuid]) {
let mut txs = self.room_shutdown_txs.write().await;
let stale: Vec<Uuid> = txs
.keys()
.filter(|id| !active_room_ids.contains(id))
.copied()
.collect();
for id in &stale {
txs.remove(id);
crate::service::unmark_room_spawned(*id);
}
drop(txs);
let mut counts = self.room_subscriber_count.write().await;
counts.retain(|room_id, _| active_room_ids.contains(room_id));
}
pub async fn register_project(&self, project_id: Uuid) -> broadcast::Receiver<()> {
let mut txs = self.project_shutdown_txs.write().await;
if let Some(tx) = txs.get(&project_id) {
return tx.subscribe();
}
let (tx, rx) = broadcast::channel(SHUTDOWN_CHANNEL_CAPACITY);
txs.insert(project_id, tx);
rx
}
pub async fn shutdown_project(&self, project_id: Uuid) {
{
let txs = self.project_shutdown_txs.read().await;
if let Some(tx) = txs.get(&project_id) {
let _ = tx.send(());
}
}
{
let mut map = self.project_inner.write().await;
map.remove(&project_id);
}
{
let mut txs = self.project_shutdown_txs.write().await;
txs.remove(&project_id);
}
}
pub async fn prune_stale_projects(&self, active_project_ids: &[Uuid]) {
let mut txs = self.project_shutdown_txs.write().await;
txs.retain(|project_id, _| active_project_ids.contains(project_id));
}
pub async fn register_user(&self, user_id: Uuid) -> broadcast::Receiver<()> {
let mut txs = self.user_shutdown_txs.write().await;
if let Some(tx) = txs.get(&user_id) {
return tx.subscribe();
}
let (tx, rx) = broadcast::channel(SHUTDOWN_CHANNEL_CAPACITY);
txs.insert(user_id, tx);
rx
}
pub async fn shutdown_user(&self, user_id: Uuid) {
{
let txs = self.user_shutdown_txs.read().await;
if let Some(tx) = txs.get(&user_id) {
let _ = tx.send(());
}
}
{
let mut map = self.user_inner.write().await;
map.remove(&user_id);
}
{
let mut txs = self.user_shutdown_txs.write().await;
txs.remove(&user_id);
}
}
}