Skip to content
Clone
194c44c
284 lines 8.4 KB Rust
use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, broadcast};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::errors::ChatError;
pub fn room_channel(room_id: Uuid) -> String {
format!("chat:room:{room_id}")
}
pub fn user_channel(uid: i64) -> String {
export::notify::user_channel(uid)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum RoomEvent {
MessagePublished {
room_id: Uuid,
seq: i64,
message: serde_json::Value,
},
MessageEdited {
room_id: Uuid,
seq: i64,
message: serde_json::Value,
},
MessageDeleted {
room_id: Uuid,
seq: i64,
message_id: Uuid,
},
MessagesDeleted {
room_id: Uuid,
seq: i64,
message_ids: Vec<Uuid>,
},
MessagePinned {
room_id: Uuid,
seq: i64,
message_id: Uuid,
pinned: bool,
},
ReactionUpdated {
room_id: Uuid,
message_id: Uuid,
reactions: serde_json::Value,
},
Typing {
room_id: Uuid,
user_id: i64,
thread_id: Option<Uuid>,
ttl_secs: u64,
},
RoomUpdated {
room_id: Uuid,
},
MemberJoined {
room_id: Uuid,
user_id: i64,
},
MemberLeft {
room_id: Uuid,
user_id: i64,
},
ThreadUpdated {
room_id: Uuid,
thread_id: Uuid,
},
/// 语音频道在场集合变化(webhook 加入/离开、对账回写)。
VoicePresence {
room_id: Uuid,
},
/// 举手集合变化(举手、审批、收回、离开)。
VoiceHands {
room_id: Uuid,
},
/// 参与者窗口布局变化(拖动/缩放,房间内所有人共享)。
VoiceLayout {
room_id: Uuid,
},
/// 主持人动作(静音、解除静音、移出、授权)。
VoiceModeration {
room_id: Uuid,
user_id: i64,
action: String,
},
}
impl RoomEvent {
pub fn room_id(&self) -> Uuid {
match self {
Self::MessagePublished { room_id, .. }
| Self::MessageEdited { room_id, .. }
| Self::MessageDeleted { room_id, .. }
| Self::MessagesDeleted { room_id, .. }
| Self::MessagePinned { room_id, .. }
| Self::ReactionUpdated { room_id, .. }
| Self::Typing { room_id, .. }
| Self::RoomUpdated { room_id, .. }
| Self::MemberJoined { room_id, .. }
| Self::MemberLeft { room_id, .. }
| Self::ThreadUpdated { room_id, .. }
| Self::VoicePresence { room_id }
| Self::VoiceHands { room_id }
| Self::VoiceLayout { room_id }
| Self::VoiceModeration { room_id, .. } => *room_id,
}
}
pub fn seq(&self) -> Option<i64> {
match self {
Self::MessagePublished { seq, .. }
| Self::MessageEdited { seq, .. }
| Self::MessageDeleted { seq, .. }
| Self::MessagesDeleted { seq, .. }
| Self::MessagePinned { seq, .. } => Some(*seq),
_ => None,
}
}
}
struct ChannelEntry {
tx: broadcast::Sender<String>,
cancel: CancellationToken,
}
struct Inner {
redis: export::AppRedis,
pubsub_client: redis::Client,
rooms: Mutex<HashMap<Uuid, ChannelEntry>>,
users: Mutex<HashMap<i64, ChannelEntry>>,
}
#[derive(Clone)]
pub struct Realtime {
inner: Arc<Inner>,
}
impl Realtime {
pub fn new(redis: export::AppRedis) -> Result<Self, ChatError> {
let url = redis.config().url().to_string();
let first = url
.split(',')
.map(str::trim)
.find(|s| !s.is_empty())
.ok_or_else(|| ChatError::Internal("empty redis url".to_string()))?;
let pubsub_client = redis::Client::open(first)?;
Ok(Self {
inner: Arc::new(Inner {
redis,
pubsub_client,
rooms: Mutex::new(HashMap::new()),
users: Mutex::new(HashMap::new()),
}),
})
}
pub async fn publish_room(&self, event: &RoomEvent) {
self.publish(&room_channel(event.room_id()), event).await;
}
pub async fn publish_user<T: Serialize>(&self, uid: i64, event: &T) {
self.publish(&user_channel(uid), event).await;
}
async fn publish<T: Serialize>(&self, channel: &str, event: &T) {
let payload = match serde_json::to_vec(event) {
Ok(bytes) => bytes,
Err(error) => {
tracing::warn!(%error, %channel, "chat realtime serialize failed");
return;
}
};
let mut conn = self.inner.redis.connection_clone();
let result: Result<i64, redis::RedisError> = match &mut conn {
export::RedisConnection::Standalone(c) => {
redis::cmd("PUBLISH")
.arg(channel)
.arg(payload.as_slice())
.query_async(c)
.await
}
export::RedisConnection::Cluster(c) => {
redis::cmd("PUBLISH")
.arg(channel)
.arg(payload.as_slice())
.query_async(c)
.await
}
};
if let Err(error) = result {
tracing::warn!(%error, %channel, "chat realtime publish failed");
}
}
pub async fn subscribe_room(
&self,
room_id: Uuid,
) -> Result<broadcast::Receiver<String>, ChatError> {
let channel = room_channel(room_id);
let mut map = self.inner.rooms.lock().await;
if let Some(entry) = map.get(&room_id) {
return Ok(entry.tx.subscribe());
}
let (tx, cancel) = spawn_pump(&self.inner.pubsub_client, &channel).await?;
let rx = tx.subscribe();
map.insert(room_id, ChannelEntry { tx, cancel });
Ok(rx)
}
pub async fn unsubscribe_room(&self, room_id: Uuid) {
let mut map = self.inner.rooms.lock().await;
let remove = map
.get(&room_id)
.is_some_and(|entry| entry.tx.receiver_count() == 0);
if remove && let Some(entry) = map.remove(&room_id) {
entry.cancel.cancel();
}
}
pub async fn subscribe_user(
&self,
uid: i64,
) -> Result<broadcast::Receiver<String>, ChatError> {
let channel = user_channel(uid);
let mut map = self.inner.users.lock().await;
if let Some(entry) = map.get(&uid) {
return Ok(entry.tx.subscribe());
}
let (tx, cancel) = spawn_pump(&self.inner.pubsub_client, &channel).await?;
let rx = tx.subscribe();
map.insert(uid, ChannelEntry { tx, cancel });
Ok(rx)
}
pub async fn unsubscribe_user(&self, uid: i64) {
let mut map = self.inner.users.lock().await;
let remove = map
.get(&uid)
.is_some_and(|entry| entry.tx.receiver_count() == 0);
if remove && let Some(entry) = map.remove(&uid) {
entry.cancel.cancel();
}
}
}
async fn spawn_pump(
client: &redis::Client,
channel: &str,
) -> Result<(broadcast::Sender<String>, CancellationToken), ChatError> {
let mut pubsub = client.get_async_pubsub().await?;
pubsub.subscribe(channel).await?;
let (tx, _) = broadcast::channel::<String>(4096);
let cancel = CancellationToken::new();
let pump_tx = tx.clone();
let pump_cancel = cancel.clone();
let channel_name = channel.to_string();
tokio::spawn(async move {
use futures::StreamExt;
let mut stream = pubsub.on_message();
loop {
tokio::select! {
_ = pump_cancel.cancelled() => break,
msg = stream.next() => {
let Some(msg) = msg else { break };
let payload: Result<String, _> = msg.get_payload();
if let Ok(payload) = payload
&& pump_tx.send(payload).is_err()
&& pump_tx.receiver_count() == 0
{
tracing::debug!(channel = %channel_name, "chat realtime channel idle");
}
}
}
}
tracing::debug!(channel = %channel_name, "chat realtime pump exited");
});
Ok((tx, cancel))
}