重设文件结构+解决报错和断线问题

src/client.rs

@@ -43,7 +43,10 @@ async fn close_existing_connection(from_id: &str) {
 
     let mut connections = CONNECTIONS.lock().await;
     let already_done = connections.get(&from_id.to_string());
-    println!("关闭之前绑定的 WebSocket 连接: {} {:?}", from_id, already_done);
+    println!(
+        "关闭之前绑定的 WebSocket 连接: {} {:?}",
+        from_id, already_done
+    );
     if let Some(mut old_connection) = connections.remove(from_id) {
         // 尝试优雅地关闭旧连接
         if let Err(e) = old_connection.sender.close().await {
@@ -55,19 +58,25 @@ async fn close_existing_connection(from_id: &str) {
 
 // 更新 Redis 中 connected 集合
 async fn update_connected_redis() {
-    let connections = CONNECTIONS.lock().await;
-    let from_ids: Vec<String> = connections.keys().cloned().collect();
-    let mut con = REDIS_POOL.get_connection().expect("Failed to get Redis connection");
+    // 改进后
+    let from_ids: Vec<String>;
+    {
+        let connections = CONNECTIONS.lock().await;
+        from_ids = connections.keys().cloned().collect();
+    }
+    let mut con = REDIS_POOL
+        .get_connection()
+        .expect("Failed to get Redis connection");
 
     // 先清空集合
     if let Err(e) = con.del::<_, ()>("connected") {
         println!("Failed to delete connected key in Redis: {}", e);
     }
 
-    // 将每个 fromId 依次添加到集合中
-    for from_id in from_ids {
-        if let Err(e) = con.sadd::<_, _, ()>("connected", from_id) {
-            println!("Failed to add fromId to connected set in Redis: {}", e);
+    // 将 from_ids 一次性批量添加到集合中
+    if !from_ids.is_empty() {
+        if let Err(e) = con.sadd::<_, _, ()>("connected", from_ids.as_slice()) {
+            println!("Failed to add fromIds to connected set in Redis: {}", e);
         }
     }
 }
@@ -105,6 +114,25 @@ fn handle_handshake(
     Ok(connection_params)
 }
 
+async fn handle_connection_error(from_id: &str) {
+    // 从全局连接映射中移除该连接
+    {
+        let mut connections = CONNECTIONS.lock().await;
+        connections.remove(from_id);
+    }
+
+    // 更新 Redis 中的 connected 集合
+    update_connected_redis().await;
+
+    // 取消对应的任务
+    let mut tasks = TASKS.lock().unwrap();
+    if let Some(task) = tasks.remove(from_id) {
+        task.abort();
+    }
+
+    println!("断开与用户id: {} 的连接并完成清理操作", from_id);
+}
+
 pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), Error> {
     let must_existed_params = ["deviceId", "fromId", "wsPwd"];
     let mut connection_params = None;
@@ -126,7 +154,7 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
             }
         }
     })
-        .await
+    .await
     {
         Ok(ws) => ws,
         Err(e) => {
@@ -159,7 +187,8 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
 
                 loop {
                     let mut connections = CONNECTIONS.lock().await;
-                    if let Some(current_connection) = connections.get_mut(&from_id_clone) { // 使用克隆后的 from_id
+                    if let Some(current_connection) = connections.get_mut(&from_id_clone) {
+                        // 使用克隆后的 from_id
                         let receiver_ref = &mut current_connection.receiver;
                         let sender_ref = &mut current_connection.sender;
 
@@ -177,11 +206,19 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
                                                             println!("收到客户端心跳消息 {:?}", &data);
                                                             handle_heartbeat(&mut last_heartbeat_time);
                                                             if let Ok(json_str) = make_common_resp(Default::default(), "Heart") {
-                                                                sender_ref.send(Message::text(json_str)).await.unwrap();
+                                                                if let Err(e) = sender_ref.send(Message::text(json_str)).await {
+                                                                    println!("发送心跳信息失败: {}", e);
+                                                                    handle_connection_error(&from_id_clone).await;
+                                                                    break;
+                                                                }
                                                             }
                                                         },
                                                         _ => {
-                                                            handle_other_message(sender_ref, &data).await.unwrap();
+                                                            if let Err(e) = handle_other_message(sender_ref, &data).await {
+                                                                println!("Failed to handle other message: {}", e);
+                                                                handle_connection_error(&from_id_clone).await;
+                                                                break;
+                                                            }
                                                         }
                                                     }
                                                 }
@@ -193,10 +230,12 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
                                     }
                                     Some(Err(e)) => {
                                         println!("接受客户端消息出错: {}", e);
+                                        handle_connection_error(&from_id_clone).await;
                                         break;
                                     }
                                     None => {
                                         println!("客户端断开连接");
+                                        handle_connection_error(&from_id_clone).await;
                                         break;
                                     }
                                 }
@@ -213,7 +252,7 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
                 }
 
                 println!("断开与用户id: {}，连接", from_id_clone); // 使用克隆后的 from_id
-                // 从全局连接映射中移除该连接
+                                                                   // 从全局连接映射中移除该连接
                 {
                     let mut connections = CONNECTIONS.lock().await;
                     connections.remove(&from_id_clone); // 使用克隆后的 from_id
@@ -231,4 +270,4 @@ pub(crate) async fn handle_client(stream: tokio::net::TcpStream) -> Result<(), E
     }
 
     Ok(())
-}
\ No newline at end of file
+}

src/config/config.rs

-pub const STATIC_WS_PWD: &str = "Q8kFm5LzJ2Ab";
\ No newline at end of file
+pub const STATIC_WS_PWD: &str = "Q8kFm5LzJ2Ab";
+pub const STATIC_ADDR: &str = "0.0.0.0:12345";

src/config/mod.rs

-pub mod config;
\ No newline at end of file
+pub mod config;

src/handles/handle_messages.rs

-
+use crate::utils::json_utils::{make_common_resp, MessageData};
 use futures::SinkExt;
 use tungstenite::{Error, Message};
-use crate::utils::json_utils::{make_common_resp, MessageData};
 
 // 处理其他类型消息
 pub(crate) async fn handle_other_message(
@@ -14,4 +13,4 @@ pub(crate) async fn handle_other_message(
         sender.send(Message::text(json_str)).await?;
     }
     Ok(())
-}
\ No newline at end of file
+}

src/handles/heartbeat.rs

@@ -2,4 +2,4 @@ use std::time::Instant;
 
 pub fn handle_heartbeat(last_time: &mut Instant) {
     *last_time = Instant::now();
-}
\ No newline at end of file
+}

src/handles/mod.rs

 pub mod handle_messages;
-pub mod heartbeat;
\ No newline at end of file
+pub mod heartbeat;

src/main.rs

 extern crate core;
 
 mod client;
-mod utils;
 mod config;
 mod handles;
 mod typing;
+mod utils;
 
 use client::handle_client;
+use config::config::STATIC_ADDR as addr;
 use tokio::net::TcpListener;
 
 #[tokio::main]
 async fn main() {
-    let addr = "0.0.0.0:12345";
     let listener = TcpListener::bind(addr).await.unwrap();
     while let Ok((stream, _)) = listener.accept().await {
         tokio::spawn(handle_client(stream));
     }
 }
-

src/typing/mod.rs

-pub mod used_typed;
\ No newline at end of file
+pub mod used_typed;

src/typing/used_typed.rs

 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
+use tokio::sync::Mutex as AsyncMutex;
 use tokio::task::JoinHandle;
 use tokio_tungstenite::WebSocketStream;
 use tungstenite::Message;
-use tokio::sync::Mutex as AsyncMutex;
 
 // 自定义结构体来存储发送器和接收器
 #[derive(Debug)]
@@ -15,4 +15,4 @@ pub struct Connection {
 // 全局连接映射，存储 fromId 到 Connection 的映射
 pub type ConnectionMap = Arc<AsyncMutex<HashMap<String, Connection>>>;
 // 全局任务映射，存储 fromId 到 JoinHandle 的映射
-pub type TaskMap = Arc<Mutex<HashMap<String, JoinHandle<()>>>>;
\ No newline at end of file
+pub type TaskMap = Arc<Mutex<HashMap<String, JoinHandle<()>>>>;

src/utils/json_utils.rs

@@ -31,4 +31,4 @@ pub fn parse_message(json_str: &str) -> Result<MessageData> {
 // 序列化消息为 JSON 字符串
 pub fn serialize_message(msg: &MessageData) -> Result<String> {
     serde_json::to_string(msg)
-}
\ No newline at end of file
+}

src/utils/mod.rs

 pub mod json_utils;
-pub mod utils;
\ No newline at end of file
+pub mod utils;

src/utils/utils.rs

 use std::collections::HashMap;
 
-pub(crate) fn get_connection_params(connection_url: String) -> Result<HashMap<String, String>, &'static str> {
+pub(crate) fn get_connection_params(
+    connection_url: String,
+) -> Result<HashMap<String, String>, &'static str> {
     if let Some(query_part) = connection_url.split('?').nth(1) {
         let mut params_mapping = HashMap::new();
         for param in query_part.split('&') {
@@ -14,4 +16,4 @@ pub(crate) fn get_connection_params(connection_url: String) -> Result<HashMap<St
     } else {
         Ok(HashMap::new())
     }
-}
\ No newline at end of file
+}