d3d11f4
use std::collections::HashMap;
use std::time::Duration;
use chat::errors::ChatError;
use chat::services::messages::ListMessageQuery;
use chat::services::service::ChatService;
use chat::share::realtime::RoomEvent;
use futures::{SinkExt, StreamExt};
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use tokio::sync::mpsc;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::{self, Message};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::frames::{ClientFrame, ServerFrame};
pub type ServerWs = WebSocketStream<TokioIo<Upgraded>>;
#[derive(Clone)]
enum Inbox {
Room(Uuid, Result<String, ()>),
User(String),
}
struct Subscription {
cancel: CancellationToken,
handle: tokio::task::JoinHandle<()>,
}
/// Push into the socket inbox without ever deadlocking the forwarder:
/// bounded retry, then drop the payload (gap detection resyncs the client).
async fn best_effort_send(
tx: &mpsc::Sender<Inbox>,
item: Inbox,
cancel: &CancellationToken,
) -> bool {
for _ in 0..1000 {
match tx.try_send(item.clone()) {
Ok(()) => return true,
Err(mpsc::error::TrySendError::Full(_)) => {
tokio::select! {
_ = cancel.cancelled() => return false,
_ = tokio::time::sleep(Duration::from_millis(1)) => {}
}
}
Err(mpsc::error::TrySendError::Closed(_)) => return false,
}
}
false
}
pub async fn handle_socket(socket: ServerWs, service: ChatService, uid: Uuid, session_id: Uuid) {
let _connection = export::metrics::ImConnectionGuard::new();
let (mut sink, mut stream) = socket.split();
let (inbox_tx, mut inbox_rx) = mpsc::channel::<Inbox>(1024);
let user_cancel = CancellationToken::new();
let mut user_rx = match service.realtime.subscribe_user(uid).await {
Ok(rx) => rx,
Err(error) => {
tracing::warn!(%error, "ws user subscribe failed");
let _ = close_with_unsplit(sink, "internal error").await;
return;
}
};
let user_handle = {
let tx = inbox_tx.clone();
let cancel = user_cancel.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = cancel.cancelled() => break,
payload = user_rx.recv() => {
let Ok(payload) = payload else { break };
if !best_effort_send(&tx, Inbox::User(payload), &cancel).await {
break;
}
}
}
}
})
};
let ready = ServerFrame::Ready {
user_id: uid,
session_id,
};
if send_frame(&mut sink, &ready).await.is_err() {
user_cancel.cancel();
let _ = user_handle.await;
service.realtime.unsubscribe_user(uid).await;
return;
}
// 在线标记：推送 worker 用它决定"这个用户已经有实时通道，不需要系统
// 通知"。同一个账号可以同时开多个连接，每条连接在集合里各占一个成员，
// 只有全部断开（或成员过期）才算离线。
let online = OnlineMarker::claim(&service, uid).await;
// 可见在线状态跟随同一条 WS 生命周期：claim 即上线（用户已有的 Away/
// Dnd/隐身状态原样保留、只续期），tick 续期，最后一条连接断开时下线。
service.mark_online(uid).await;
let heartbeat_secs = service.config().chat.ws.heartbeat_secs.max(5);
let idle_limit = Duration::from_secs(heartbeat_secs * 3);
let mut last_activity = tokio::time::Instant::now();
let mut ticker = tokio::time::interval(Duration::from_secs(heartbeat_secs));
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut subscriptions: HashMap<Uuid, Subscription> = HashMap::new();
let mut last_seq: HashMap<Uuid, i64> = HashMap::new();
let max_subscriptions = service.config().chat.ws.max_subscriptions;
let mut last_typing: Option<tokio::time::Instant> = None;
let mut tick_count: u32 = 0;
loop {
tokio::select! {
frame = stream.next() => {
let Some(frame) = frame else { break };
let Ok(frame) = frame else { break };
last_activity = tokio::time::Instant::now();
match frame {
Message::Text(text) => {
let frame: Result<ClientFrame, _> = serde_json::from_str(&text);
match frame {
Ok(frame) => {
handle_client_frame(
&mut sink,
&service,
uid,
frame,
&mut subscriptions,
&mut last_seq,
&mut last_typing,
&inbox_tx,
max_subscriptions,
)
.await;
}
Err(_) => {
let _ = send_frame(&mut sink, &ServerFrame::Error {
message: "invalid frame".to_string(),
}).await;
}
}
}
Message::Ping(payload) => {
let _ = sink.send(Message::Pong(payload)).await;
}
Message::Close(_) => break,
_ => {}
}
}
inbox = inbox_rx.recv() => {
let Some(inbox) = inbox else { break };
match inbox {
Inbox::Room(room_id, Ok(payload)) => {
let gap = check_seq_gap(&payload, &mut last_seq, room_id);
if let Some(frame) = gap
&& send_frame(&mut sink, &frame).await.is_err()
{
break;
}
let frame = ServerFrame::Event {
room_id,
event: serde_json::from_str(&payload)
.unwrap_or(serde_json::Value::Null),
};
if send_frame(&mut sink, &frame).await.is_err() {
break;
}
}
Inbox::Room(room_id, Err(())) => {
let frame = ServerFrame::ResyncRequired {
room_id,
reason: "lagged".to_string(),
};
last_seq.remove(&room_id);
if send_frame(&mut sink, &frame).await.is_err() {
break;
}
}
Inbox::User(payload) => {
let frame = ServerFrame::UserEvent {
event: serde_json::from_str(&payload)
.unwrap_or(serde_json::Value::Null),
};
if send_frame(&mut sink, &frame).await.is_err() {
break;
}
}
}
}
_ = ticker.tick() => {
if last_activity.elapsed() > idle_limit {
let _ = send_frame(&mut sink, &ServerFrame::Bye {
message: "idle timeout".to_string(),
}).await;
break;
}
// 续期必须在 tick 里做而不是只在 ready 时写一次：TTL 到期而
// 进程还活着（例如用户只是在读、没有发帧）时，键消失会让推送
// worker 把他当成离线，于是同一个事件既弹 toast 又弹系统通知；
// 可见在线状态同样靠这里的心跳保活。
online.refresh().await;
service.touch_presence(uid).await;
tick_count = tick_count.wrapping_add(1);
if tick_count.is_multiple_of(2) {
let mut revoked: Vec<Uuid> = Vec::new();
for &room_id in subscriptions.keys() {
if service_can_read(&service, uid, room_id).await.is_err() {
revoked.push(room_id);
}
}
for room_id in revoked {
if let Some(sub) = subscriptions.remove(&room_id) {
sub.cancel.cancel();
let _ = sub.handle.await;
service.realtime.unsubscribe_room(room_id).await;
last_seq.remove(&room_id);
if send_frame(&mut sink, &ServerFrame::ResyncRequired {
room_id,
reason: "access revoked".to_string(),
})
.await
.is_err()
{
break;
}
}
}
}
}
}
}
// 计数归零＝最后一条连接关闭；Redis 故障时 release 返回 None，此时不判定
// 离线，交给 TTL 兜底。
if online.release().await == Some(0) {
service.mark_offline(uid).await;
}
let mut teardown: Vec<(Uuid, tokio::task::JoinHandle<()>)> = Vec::new();
for (room_id, sub) in subscriptions {
sub.cancel.cancel();
teardown.push((room_id, sub.handle));
}
user_cancel.cancel();
// Releasing the local sender unblocks any forwarder parked on a full
// inbox, guaranteeing the awaits below terminate.
drop(inbox_tx);
let _ = user_handle.await;
for (room_id, handle) in teardown {
// Await the forwarder so its broadcast receiver is released before
// the unsubscribe refcount check.
let _ = handle.await;
service.realtime.unsubscribe_room(room_id).await;
}
service.realtime.unsubscribe_user(uid).await;
}
/// 一条 WS 连接的"在线"占位。
///
/// 每条连接是集合里的一个成员（分数＝到期时刻）：多标签页/多实例各自登记，
/// 只有全部断开才算离线；进程被强杀留下的成员会随分数过期被惰性清理，不会
/// 留下永久在线的幽灵，也不会污染在线计数。
struct OnlineMarker {
key: String,
member: String,
ttl_secs: u64,
redis: export::AppRedis,
}
impl OnlineMarker {
async fn claim(service: &ChatService, uid: Uuid) -> Self {
let ttl_secs = service.config().chat.online_ttl_secs();
let marker = Self {
key: export::push::online_key(uid),
// 每条连接一个唯一成员：多标签页/多实例各自登记，互不覆盖。
member: Uuid::now_v7().to_string(),
ttl_secs,
redis: service.state.redis.clone(),
};
marker.refresh().await;
marker
}
/// 登记/续期本连接；过期的旧成员在这里被顺手清掉。续期是幂等的，
/// 心跳不会像计数型占位那样把数字越推越高。
async fn refresh(&self) {
if let Err(error) =
export::redis::zset_touch(&self.redis, &self.key, &self.member, self.ttl_secs).await
{
tracing::warn!(%error, key = %self.key, "online marker refresh failed");
}
}
/// 注销本连接，返回剩余**存活**连接数；归零即最后一条连接退出。
/// Redis 故障返回 `None`，调用方不应据此判定离线（TTL 会兜底）。
async fn release(&self) -> Option<i64> {
match export::redis::zset_release(&self.redis, &self.key, &self.member).await {
Ok(remaining) => Some(remaining),
Err(error) => {
tracing::warn!(%error, key = %self.key, "online marker release failed");
None
}
}
}
}
fn check_seq_gap(
payload: &str,
last_seq: &mut HashMap<Uuid, i64>,
room_id: Uuid,
) -> Option<ServerFrame> {
let event: serde_json::Value = serde_json::from_str(payload).ok()?;
let seq = event.get("seq")?.as_i64()?;
let last = last_seq.get(&room_id).copied();
// Only advance the cursor; edit/pin events may carry older seq values.
if last.is_some_and(|l| seq <= l) {
return None;
}
last_seq.insert(room_id, seq);
match last {
Some(last) if seq > last + 1 => Some(ServerFrame::Gap {
room_id,
from: last + 1,
to: seq,
}),
_ => None,
}
}
#[allow(clippy::too_many_arguments)]
async fn handle_client_frame<S>(
sink: &mut S,
service: &ChatService,
uid: Uuid,
frame: ClientFrame,
subscriptions: &mut HashMap<Uuid, Subscription>,
last_seq: &mut HashMap<Uuid, i64>,
last_typing: &mut Option<tokio::time::Instant>,
inbox_tx: &mpsc::Sender<Inbox>,
max_subscriptions: usize,
) where
S: futures::Sink<Message, Error = tungstenite::Error> + Unpin,
{
match frame {
ClientFrame::Subscribe { room_id, resync } => {
if subscriptions.contains_key(&room_id) {
return;
}
if subscriptions.len() >= max_subscriptions {
let _ = send_frame(
sink,
&ServerFrame::Error {
message: "subscription limit reached".to_string(),
},
)
.await;
return;
}
if service_can_read(service, uid, room_id).await.is_err() {
let _ = send_frame(
sink,
&ServerFrame::Error {
message: "forbidden".to_string(),
},
)
.await;
return;
}
let rx = match service.realtime.subscribe_room(room_id).await {
Ok(rx) => rx,
Err(_) => {
let _ = send_frame(
sink,
&ServerFrame::Error {
message: "subscribe failed".to_string(),
},
)
.await;
return;
}
};
let cancel = CancellationToken::new();
let handle = {
let tx = inbox_tx.clone();
let child = cancel.clone();
let mut rx = rx;
tokio::spawn(async move {
loop {
tokio::select! {
_ = child.cancelled() => break,
result = rx.recv() => {
let item = match result {
Ok(payload) => Inbox::Room(room_id, Ok(payload)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
Inbox::Room(room_id, Err(()))
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
};
// Non-blocking fallback: overflow drops the
// frame; seq gap detection triggers resync.
if !best_effort_send(&tx, item, &child).await {
break;
}
}
}
}
})
};
subscriptions.insert(room_id, Subscription { cancel, handle });
let history_limit = service.config().chat.ws.history_limit;
if let Ok(page) = service
.list_messages(
uid,
room_id,
ListMessageQuery {
before_seq: None,
after_seq: None,
limit: Some(history_limit),
thread_id: None,
author_id: None,
},
)
.await
{
if resync.unwrap_or(true) {
let frame = ServerFrame::History {
room_id,
messages: page.items,
};
let _ = send_frame(sink, &frame).await;
}
if let Some(newest) = page.newest_seq {
last_seq.insert(room_id, newest);
}
}
}
ClientFrame::Unsubscribe { room_id } => {
if let Some(sub) = subscriptions.remove(&room_id) {
sub.cancel.cancel();
let _ = sub.handle.await;
last_seq.remove(&room_id);
service.realtime.unsubscribe_room(room_id).await;
}
}
ClientFrame::Resync { room_id, after_seq } => {
if service_can_read(service, uid, room_id).await.is_err() {
return;
}
if let Ok(page) = service
.list_messages(
uid,
room_id,
ListMessageQuery {
before_seq: None,
after_seq: Some(after_seq),
limit: Some(service.config().chat.ws.history_limit),
thread_id: None,
author_id: None,
},
)
.await
{
if let Some(newest) = page.newest_seq {
last_seq.insert(room_id, newest);
}
let frame = ServerFrame::History {
room_id,
messages: page.items,
};
let _ = send_frame(sink, &frame).await;
}
}
ClientFrame::Typing { room_id, thread_id } => {
// Coarse throttle: at most one typing broadcast per 3s per socket.
let now = tokio::time::Instant::now();
if last_typing.is_some_and(|prev| now.duration_since(prev).as_secs() < 3) {
return;
}
if service_can_read(service, uid, room_id).await.is_ok() {
*last_typing = Some(now);
let event = RoomEvent::Typing {
room_id,
user_id: uid,
thread_id,
ttl_secs: service.config().chat.typing_ttl_secs.max(1),
};
service.realtime.publish_room(&event).await;
}
}
ClientFrame::Ping => {
let _ = send_frame(sink, &ServerFrame::Pong).await;
}
}
}
async fn service_can_read(
service: &ChatService,
uid: Uuid,
room_id: Uuid,
) -> Result<(), ChatError> {
chat::share::lookup::require_room_readable(service, room_id, uid)
.await
.map(|_| ())
}
async fn send_frame<S>(sink: &mut S, frame: &ServerFrame) -> Result<(), tungstenite::Error>
where
S: futures::Sink<Message, Error = tungstenite::Error> + Unpin,
{
let payload = match serde_json::to_string(frame) {
Ok(payload) => payload,
Err(_) => return Ok(()),
};
sink.send(Message::Text(payload.into())).await
}
async fn close_with_unsplit<S>(mut sink: S, message: &str) -> Result<(), tungstenite::Error>
where
S: futures::Sink<Message, Error = tungstenite::Error> + Unpin,
{
let frame = ServerFrame::Bye {
message: message.to_string(),
};
let payload = serde_json::to_string(&frame).unwrap_or_default();
sink.send(Message::Text(payload.into())).await?;
sink.send(Message::Close(None)).await
}