moduforge_collaboration/
ws_server.rs

1use std::net::SocketAddr;
2use std::sync::Arc;
3use tokio::net::{TcpListener, TcpStream};
4use tokio_tungstenite::{accept_async, tungstenite::Message};
5use futures_util::{StreamExt, SinkExt};
6use dashmap::DashMap;
7use uuid::Uuid;
8use serde::{Deserialize, Serialize};
9use tokio::sync::mpsc;
10use tokio::sync::Mutex;
11use crate::{Result, TransmissionError, YrsManager, ClientInfo, SyncService};
12use anyhow;
13use yrs::updates::decoder::Decode;
14use moduforge_core::runtime::ForgeRuntime;
15use yrs::Transact;
16use yrs::ReadTxn;
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub enum WsMessage {
20    /// 客户端加入房间
21    JoinRoom { room_id: String },
22    /// 客户端离开房间
23    LeaveRoom { room_id: String },
24    /// Yrs update 数据 (二进制)
25    YrsUpdate { room_id: String, update: Vec<u8> },
26    /// 客户端请求同步
27    YrsSyncRequest { room_id: String, state_vector: Vec<u8> },
28    /// JSON格式的状态同步消息
29    StateSync { 
30        room_id: String, 
31        operation: String,
32        data: serde_json::Value,
33        timestamp: u64 
34    },
35    /// 心跳
36    Ping,
37    /// 心跳响应
38    Pong,
39    /// 错误消息
40    Error { message: String },
41    /// 服务器通知消息
42    Notification { message: String },
43}
44
45/// 内部广播消息
46#[derive(Debug, Clone)]
47pub enum BroadcastMessage {
48    /// 向房间广播文本消息
49    ToRoom { room_id: String, message: Arc<String>, exclude_client: Option<String> },
50    /// 向房间广播二进制消息
51    ToRoomBinary { room_id: String, data: Arc<Vec<u8>>, exclude_client: Option<String> },
52    /// 向特定客户端发送消息
53    ToClient { client_id: String, message: Arc<String> },
54    /// 向特定客户端发送二进制消息
55    ToClientBinary { client_id: String, data: Arc<Vec<u8>> },
56}
57
58/// A message to be sent to a single client's websocket.
59/// Using Arc to avoid cloning the underlying data for every client in a broadcast.
60#[derive(Debug, Clone)]
61pub enum OutgoingMessage {
62    Text(Arc<String>),
63    Binary(Arc<Vec<u8>>),
64    // We can add other types like Close, etc. if needed
65}
66
67impl From<OutgoingMessage> for Message {
68    fn from(msg: OutgoingMessage) -> Self {
69        match msg {
70            OutgoingMessage::Text(text) => Message::Text((*text).clone()),
71            OutgoingMessage::Binary(data) => Message::Binary((*data).clone()),
72        }
73    }
74}
75
76/// 客户端连接信息（扩展版）
77#[derive(Debug, Clone)]
78pub struct ClientConnection {
79    pub info: ClientInfo,
80    pub sender: mpsc::UnboundedSender<OutgoingMessage>,
81}
82
83pub struct WebSocketServer {
84    yrs_manager: Arc<YrsManager>,
85    /// 客户端连接信息和发送通道
86    clients: Arc<DashMap<String, ClientConnection>>,
87    /// 房间到客户端的映射
88    room_clients: Arc<DashMap<String, Vec<String>>>,
89    /// 广播消息通道
90    broadcast_tx: mpsc::UnboundedSender<BroadcastMessage>,
91    broadcast_rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<BroadcastMessage>>>,
92}
93
94impl WebSocketServer {
95    pub fn new(yrs_manager: Arc<YrsManager>) -> Self {
96        let (broadcast_tx, broadcast_rx) = mpsc::unbounded_channel();
97        
98        Self {
99            yrs_manager,
100            clients: Arc::new(DashMap::new()),
101            room_clients: Arc::new(DashMap::new()),
102            broadcast_tx,
103            broadcast_rx: Arc::new(tokio::sync::Mutex::new(broadcast_rx)),
104        }
105    }
106
107    /// 启动 WebSocket 服务器
108    pub async fn start(
109        self: Arc<Self>, 
110        addr: SocketAddr, 
111        sync_service: Arc<SyncService>,
112        runtime: Arc<Mutex<ForgeRuntime>>
113    ) -> Result<()> {
114        let listener = TcpListener::bind(addr).await?;
115        tracing::info!("WebSocket 服务器启动在: {}", addr);
116
117        // 启动广播处理任务
118        self.clone().start_broadcast_handler().await;
119
120        while let Ok((stream, addr)) = listener.accept().await {
121            let server = self.clone();
122            let sync_service = sync_service.clone();
123            let runtime = runtime.clone();
124            
125            tokio::spawn(async move {
126                if let Err(e) = server.handle_connection(
127                    stream, 
128                    addr,
129                    sync_service,
130                    runtime,
131                ).await {
132                    tracing::error!("处理连接失败: {}", e);
133                }
134            });
135        }
136
137        Ok(())
138    }
139
140    /// 启动广播消息处理器
141    async fn start_broadcast_handler(self: Arc<Self>) {
142        let clients = self.clients.clone();
143        let room_clients = self.room_clients.clone();
144        let broadcast_rx = self.broadcast_rx.clone();
145        
146        tokio::spawn(async move {
147            let mut rx = broadcast_rx.lock().await;
148            
149            while let Some(broadcast_msg) = rx.recv().await {
150                Self::handle_broadcast_message(broadcast_msg, &clients, &room_clients).await;
151            }
152        });
153    }
154
155    /// 处理广播消息
156    async fn handle_broadcast_message(
157        message: BroadcastMessage,
158        clients: &DashMap<String, ClientConnection>,
159        room_clients: &DashMap<String, Vec<String>>,
160    ) {
161        match message {
162            BroadcastMessage::ToRoom { room_id, message, exclude_client } => {
163                let outgoing_msg = OutgoingMessage::Text(message);
164                if let Some(client_list) = room_clients.get(&room_id) {
165                    for client_id in client_list.iter() {
166                        if let Some(ref exclude) = exclude_client {
167                            if client_id == exclude {
168                                continue;
169                            }
170                        }
171                        
172                        if let Some(client) = clients.get(client_id) {
173                            let _ = client.sender.send(outgoing_msg.clone());
174                        }
175                    }
176                }
177            }
178            BroadcastMessage::ToRoomBinary { room_id, data, exclude_client } => {
179                let outgoing_msg = OutgoingMessage::Binary(data);
180                if let Some(client_list) = room_clients.get(&room_id) {
181                    for client_id in client_list.iter() {
182                        if let Some(ref exclude) = exclude_client {
183                            if client_id == exclude {
184                                continue;
185                            }
186                        }
187                        
188                        if let Some(client) = clients.get(client_id) {
189                            let _ = client.sender.send(outgoing_msg.clone());
190                        }
191                    }
192                }
193            }
194            BroadcastMessage::ToClient { client_id, message } => {
195                if let Some(client) = clients.get(&client_id) {
196                    let _ = client.sender.send(OutgoingMessage::Text(message));
197                }
198            }
199            BroadcastMessage::ToClientBinary { client_id, data } => {
200                if let Some(client) = clients.get(&client_id) {
201                    let _ = client.sender.send(OutgoingMessage::Binary(data));
202                }
203            }
204        }
205    }
206
207    async fn handle_connection(
208        self: Arc<Self>,
209        stream: TcpStream,
210        addr: SocketAddr,
211        sync_service: Arc<SyncService>,
212        runtime: Arc<Mutex<ForgeRuntime>>,
213    ) -> Result<()> {
214        let ws_stream = accept_async(stream).await
215            .map_err(|e| TransmissionError::WebSocketError(e.to_string()))?;
216        
217        let client_id = Uuid::new_v4().to_string();
218        tracing::info!("新客户端连接: {} ({})", client_id, addr);
219
220        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
221        let (msg_tx, mut msg_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
222
223        // 存储客户端连接信息
224        let client_info = ClientInfo {
225            id: client_id.clone(),
226            room_id: String::new(),
227            connected_at: std::time::SystemTime::now(),
228        };
229        
230        let client_connection = ClientConnection {
231            info: client_info,
232            sender: msg_tx,
233        };
234        
235        self.clients.insert(client_id.clone(), client_connection);
236
237        let client_id_clone = client_id.clone();
238        
239        // 启动发送任务
240        let send_task = tokio::spawn(async move {
241            while let Some(message) = msg_rx.recv().await {
242                if ws_sender.send(message.into()).await.is_err() {
243                    break;
244                }
245            }
246        });
247
248        // 处理接收消息循环
249        while let Some(msg) = ws_receiver.next().await {
250            match msg {
251                Ok(Message::Text(text)) => {
252                    if let Ok(ws_msg) = serde_json::from_str::<WsMessage>(&text) {
253                        if let Err(e) = self.handle_message(
254                            &client_id,
255                            ws_msg,
256                            &sync_service,
257                            &runtime,
258                        ).await {
259                            tracing::error!("处理消息失败: {}" , e);
260                            let error_msg = WsMessage::Error {
261                                message: e.to_string()
262                            };
263                            if let Ok(json) = serde_json::to_string(&error_msg) {
264                                let _ = self.broadcast_tx.send(BroadcastMessage::ToClient {
265                                    client_id: client_id.clone(),
266                                    message: Arc::new(json),
267                                });
268                            }
269                        }
270                    }
271                }
272                Ok(Message::Binary(_data)) => {
273                    // 二进制消息现在应该通过YrsUpdate JSON消息包装
274                    tracing::warn!("收到未包装的二进制消息，已忽略");
275                }
276                Ok(Message::Close(_)) => {
277                    break;
278                }
279                Err(_) => {
280                    break;
281                }
282                _ => {}
283            }
284        }
285
286        // 清理连接
287        send_task.abort();
288        self.cleanup_client(&client_id_clone).await;
289
290        Ok(())
291    }
292
293    async fn handle_message(
294        &self,
295        client_id: &str,
296        message: WsMessage,
297        sync_service: &Arc<SyncService>,
298        runtime: &Arc<Mutex<ForgeRuntime>>,
299    ) -> Result<()> {
300        match message {
301            WsMessage::JoinRoom { room_id } => {
302                if let Some(mut client) = self.clients.get_mut(client_id) {
303                    client.info.room_id = room_id.clone();
304                }
305                self.room_clients.entry(room_id.clone()).or_default().push(client_id.to_string());
306                let _doc = self.yrs_manager.get_or_create_doc(&room_id);
307                tracing::info!("客户端 {} 加入房间 {}", client_id, room_id);
308                Ok(())
309            }
310            WsMessage::LeaveRoom { room_id } => {
311                if let Some(mut client_list) = self.room_clients.get_mut(&room_id) {
312                    client_list.retain(|id| id != client_id);
313                }
314                if let Some(mut client) = self.clients.get_mut(client_id) {
315                    client.info.room_id = String::new();
316                }
317                tracing::info!("客户端 {} 离开房间 {}", client_id, room_id);
318                Ok(())
319            }
320            WsMessage::YrsUpdate { room_id, update } => {
321                let doc = self.yrs_manager.get_doc(&room_id).ok_or_else(|| TransmissionError::RoomNotFound(room_id.clone()))?;
322                if let Err(e) = apply_update_to_doc(&doc, &update) {
323                    tracing::error!("应用更新失败: {}", e);
324                } else {
325                    self.broadcast_tx.send(BroadcastMessage::ToRoomBinary {
326                        room_id,
327                        data: Arc::new(update),
328                        exclude_client: Some(client_id.to_string()),
329                    }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))?;
330                }
331                Ok(())
332            }
333            WsMessage::YrsSyncRequest { room_id, state_vector } => {
334                tracing::debug!("收到来自 {} 的同步请求，房间: {}", client_id, room_id);
335                
336                let tree = runtime.lock().await.doc().get_inner().clone();
337                let diff_update = sync_service.handle_sync_request(&room_id, &tree, &state_vector).await?;
338
339                if !diff_update.is_empty() {
340                    self.broadcast_tx.send(BroadcastMessage::ToClientBinary {
341                        client_id: client_id.to_string(),
342                        data: Arc::new(diff_update),
343                    }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))?;
344                }
345
346                Ok(())
347            }
348            WsMessage::Ping => {
349                let pong = WsMessage::Pong;
350                let json = serde_json::to_string(&pong)?;
351                self.broadcast_tx.send(BroadcastMessage::ToClient {
352                    client_id: client_id.to_string(),
353                    message: Arc::new(json),
354                }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))?;
355                Ok(())
356            }
357            _ => Ok(()),
358        }
359    }
360
361    async fn cleanup_client(
362        &self,
363        client_id: &str,
364    ) {
365        if let Some((_, client)) = self.clients.remove(client_id) {
366            tracing::info!("客户端断开连接: {}", client_id);
367            let room_id = client.info.room_id;
368
369            if !room_id.is_empty() {
370                let should_remove_doc = if let Some(mut room) = self.room_clients.get_mut(&room_id) {
371                    room.retain(|id| id != client_id);
372                    room.is_empty()
373                } else {
374                    false
375                };
376
377                if should_remove_doc {
378                    if self.room_clients.remove(&room_id).is_some() {
379                        self.yrs_manager.remove_doc(&room_id);
380                        tracing::info!("房间 {} 已空，移除Yrs Doc", room_id);
381                    }
382                }
383            }
384        }
385    }
386
387    /// 获取在线客户端数量
388    pub fn client_count(&self) -> usize {
389        self.clients.len()
390    }
391
392    /// 获取房间数量
393    pub fn room_count(&self) -> usize {
394        self.room_clients.len()
395    }
396
397    /// 🚀 主动推送消息到房间
398    pub fn broadcast_to_room(&self, room_id: &str, message: String, exclude_client: Option<String>) -> Result<()> {
399        self.broadcast_tx.send(BroadcastMessage::ToRoom { 
400            room_id: room_id.to_string(), 
401            message: Arc::new(message), 
402            exclude_client 
403        }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))
404    }
405
406    /// 🚀 主动推送二进制数据到房间
407    pub fn broadcast_binary_to_room(&self, room_id: &str, data: Vec<u8>, exclude_client: Option<String>) -> Result<()> {
408        self.broadcast_tx.send(BroadcastMessage::ToRoomBinary { 
409            room_id: room_id.to_string(), 
410            data: Arc::new(data), 
411            exclude_client 
412        }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))
413    }
414
415    /// 🚀 主动推送消息到特定客户端
416    pub fn send_to_client(&self, client_id: &str, message: String) -> Result<()> {
417        self.broadcast_tx.send(BroadcastMessage::ToClient { 
418            client_id: client_id.to_string(), 
419            message: Arc::new(message)
420        }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))
421    }
422
423    /// 🚀 主动推送二进制数据到特定客户端
424    pub fn send_binary_to_client(&self, client_id: &str, data: Vec<u8>) -> Result<()> {
425        self.broadcast_tx.send(BroadcastMessage::ToClientBinary { 
426            client_id: client_id.to_string(), 
427            data: Arc::new(data)
428        }).map_err(|e| TransmissionError::Other(anyhow::anyhow!(e.to_string())))
429    }
430}
431
432// Helper functions extracted for clarity
433fn apply_update_to_doc(doc: &Arc<yrs::Doc>, update: &[u8]) -> crate::Result<()> {
434    let u = yrs::Update::decode_v1(update)?;
435    let mut txn = doc.transact_mut();
436    txn.apply_update(u)?;
437    Ok(())
438}
439
440#[allow(dead_code)]
441fn get_diff_update_from_doc(doc: &Arc<yrs::Doc>, sv: &[u8]) -> std::result::Result<Vec<u8>, yrs::encoding::read::Error> {
442    let state_vector = yrs::StateVector::decode_v1(sv)?;
443    let txn = doc.transact();
444    Ok(txn.encode_diff_v1(&state_vector))
445}
moduforge_collaboration/ws_server.rs

moduforge_collaboration/
ws_server.rs
