Rust 异步编程实战：构建高性能 WebSocket 服务

2.2 连接管理与消息广播

在实际应用中，我们需要管理多个 WebSocket 连接，并支持消息广播。我们可以使用 tokio::sync::broadcast 通道来实现消息广播。

use axum::{routing::get, extract::ws::{WebSocket, Message, WebSocketUpgrade}, response::IntoResponse, Router};
use tracing_subscriber::prelude::*;
use tracing::info;
use tokio::sync::broadcast;

async fn ws_handler(ws: WebSocketUpgrade, tx: broadcast::Sender<String>) -> impl IntoResponse {
    info!("New WebSocket connection");
    let rx = tx.subscribe();
    ws.on_upgrade(|socket| handle_socket(socket, tx, rx))
}

async fn handle_socket(mut socket: WebSocket, tx: broadcast::Sender<String>, mut rx: broadcast::Receiver<String>) {
    // 发送欢迎消息
    if let Err(e) = socket.send(Message::Text("Welcome to WebSocket server".to_string())).await {
        info!("Error sending welcome message: {}", e);
        return;
    }
    tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if let Err(e) = socket.send(Message::Text(msg)).await {
                info!("Error sending broadcast message: {}", e);
                break;
            }
        }
    });
    while let Some(msg) = socket.recv().await {
        match msg {
            Ok(msg) => {
                match msg {
                    Message::Text(text) => {
                        info!("Received message: {}", text);
                        // 广播消息
                        if let Err(e) = tx.send(text) {
                            info!("Error broadcasting message: {}", e);
                            break;
                        }
                    }
                    Message::Binary(data) => {
                        info!("Received binary message ({} bytes)", data.len());
                    }
                    Message::Ping(data) => {
                        info!("Received ping");
                        if let Err(e) = socket.send(Message::Pong(data)).await {
                            info!("Error sending pong: {}", e);
                            break;
                        }
                    }
                    Message::Pong(_) => {
                        info!("Received pong");
                    }
                    Message::Close(_) => {
                        info!("Connection closed");
                        break;
                    }
                }
            }
            Err(e) => {
                info!("Error receiving message: {}", e);
                break;
            }
        }
    }
    info!("WebSocket connection closed");
}

#[tokio::main]
async fn main() {
    // 初始化日志
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    // 创建广播通道
    let (tx, _) = broadcast::channel(100);
    // 创建路由
    let app = Router::new().route("/ws", get(ws_handler)).with_state(tx);
    // 启动服务器
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    info!("WebSocket server running on http://0.0.0.0:3000");
    axum::serve(listener, app).await.unwrap();
}

2.3 心跳检测与连接超时

为了确保 WebSocket 连接的有效性，我们需要实现心跳检测与连接超时机制。

use axum::{routing::get, extract::ws::{WebSocket, Message, WebSocketUpgrade}, response::IntoResponse, Router};
use tracing_subscriber::prelude::*;
use tracing::info;
use tokio::sync::broadcast;
use tokio::time::{interval, Duration};

async fn ws_handler(ws: WebSocketUpgrade, tx: broadcast::Sender<String>) -> impl IntoResponse {
    info!("New WebSocket connection");
    let rx = tx.subscribe();
    ws.on_upgrade(|socket| handle_socket(socket, tx, rx))
}

async fn handle_socket(mut socket: WebSocket, tx: broadcast::Sender<String>, mut rx: broadcast::Receiver<String>) {
    // 发送欢迎消息
    if let Err(e) = socket.send(Message::Text("Welcome to WebSocket server".to_string())).await {
        info!("Error sending welcome message: {}", e);
        return;
    }
    // 心跳检测
    let mut heartbeat_interval = interval(Duration::from_secs(10));
    // 广播接收任务
    let broadcast_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if let Err(e) = socket.send(Message::Text(msg)).await {
                info!("Error sending broadcast message: {}", e);
                break;
            }
        }
    });
    // 消息处理任务
    let message_task = tokio::spawn(async move {
        while let Some(msg) = socket.recv().await {
            match msg {
                Ok(msg) => {
                    match msg {
                        Message::Text(text) => {
                            info!("Received message: {}", text);
                            if let Err(e) = tx.send(text) {
                                info!("Error broadcasting message: {}", e);
                                break;
                            }
                        }
                        Message::Binary(data) => {
                            info!("Received binary message ({} bytes)", data.len());
                        }
                        Message::Ping(data) => {
                            info!("Received ping");
                            if let Err(e) = socket.send(Message::Pong(data)).await {
                                info!("Error sending pong: {}", e);
                                break;
                            }
                        }
                        Message::Pong(_) => {
                            info!("Received pong");
                        }
                        Message::Close(_) => {
                            info!("Connection closed");
                            break;
                        }
                    }
                }
                Err(e) => {
                    info!("Error receiving message: {}", e);
                    break;
                }
            }
        }
    });
    // 心跳检测任务
    let heartbeat_task = tokio::spawn(async move {
        loop {
            tokio::select! {
                _ = heartbeat_interval.tick() => {
                    if let Err(e) = socket.send(Message::Ping(vec![])).await {
                        info!("Error sending ping: {}", e);
                        break;
                    }
                }
                _ = message_task => {
                    break;
                }
                _ = broadcast_task => {
                    break;
                }
            }
        }
    });
    // 等待所有任务完成
    let _ = tokio::try_join!(heartbeat_task, message_task, broadcast_task);
    info!("WebSocket connection closed");
}

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    let (tx, _) = broadcast::channel(100);
    let app = Router::new().route("/ws", get(ws_handler)).with_state(tx);
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    info!("WebSocket server running on http://0.0.0.0:3000");
    axum::serve(listener, app).await.unwrap();
}

三、异步 WebSocket 客户端开发

3.1 使用 Tungstenite 实现 WebSocket 客户端

Tungstenite 是 Rust 社区中常用的 WebSocket 客户端库，它提供了异步和同步两种 API。

在 Cargo.toml 中添加依赖：

[dependencies]
tungstenite = "0.18"
tokio = { version = "1.0", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }

实现 WebSocket 客户端：

use tungstenite::connect;
use tungstenite::Message;
use url::Url;
use tracing_subscriber::prelude::*;
use tracing::info;

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    info!("Connecting to WebSocket server");
    let (mut socket, response) = connect(Url::parse("ws://127.0.0.1:3000/ws").unwrap()).unwrap();
    info!("Connected to WebSocket server, response: {:?}", response);
    // 发送消息
    socket.write_message(Message::Text("Hello, WebSocket!".to_string())).unwrap();
    // 接收消息
    loop {
        match socket.read_message() {
            Ok(msg) => {
                match msg {
                    Message::Text(text) => {
                        info!("Received message: {}", text);
                    }
                    Message::Binary(data) => {
                        info!("Received binary message ({} bytes)", data.len());
                    }
                    Message::Ping(data) => {
                        info!("Received ping");
                        socket.write_message(Message::Pong(data)).unwrap();
                    }
                    Message::Pong(_) => {
                        info!("Received pong");
                    }
                    Message::Close(_) => {
                        info!("Connection closed");
                        break;
                    }
                }
            }
            Err(e) => {
                info!("Error receiving message: {}", e);
                break;
            }
        }
    }
    info!("Client disconnected");
}

3.2 异步 WebSocket 客户端

Tungstenite 也提供了异步 API，我们可以使用 Tokio 的异步运行时来实现异步 WebSocket 客户端。

use tokio_tungstenite::connect_async;
use tungstenite::protocol::Message;
use url::Url;
use tracing_subscriber::prelude::*;
use tracing::info;

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    info!("Connecting to WebSocket server");
    let (ws_stream, response) = connect_async(Url::parse("ws://127.0.0.1:3000/ws").unwrap()).await.unwrap();
    info!("Connected to WebSocket server, response: {:?}", response);
    let (mut write, mut read) = ws_stream.split();
    // 发送消息任务
    tokio::spawn(async move {
        write.send(Message::Text("Hello, WebSocket!".to_string())).await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
        write.send(Message::Text("Another message".to_string())).await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
        write.send(Message::Close(None)).await.unwrap();
    });
    // 接收消息任务
    while let Some(msg) = read.next().await {
        match msg {
            Ok(msg) => {
                match msg {
                    Message::Text(text) => {
                        info!("Received message: {}", text);
                    }
                    Message::Binary(data) => {
                        info!("Received binary message ({} bytes)", data.len());
                    }
                    Message::Ping(data) => {
                        info!("Received ping");
                        write.send(Message::Pong(data)).await.unwrap();
                    }
                    Message::Pong(_) => {
                        info!("Received pong");
                    }
                    Message::Close(_) => {
                        info!("Connection closed");
                        break;
                    }
                }
            }
            Err(e) => {
                info!("Error receiving message: {}", e);
                break;
            }
        }
    }
    info!("Client disconnected");
}

3.3 重连机制

在实际应用中，WebSocket 连接可能会由于网络问题而断开，我们需要实现重连机制。

use tokio_tungstenite::connect_async;
use tungstenite::protocol::Message;
use url::Url;
use tracing_subscriber::prelude::*;
use tracing::info;
use tokio::time::{sleep, Duration};

async fn connect_and_handle() -> Result<(), Box<dyn std::error::Error>> {
    info!("Connecting to WebSocket server");
    let (ws_stream, response) = connect_async(Url::parse("ws://127.0.0.1:3000/ws").unwrap()).await?;
    info!("Connected to WebSocket server, response: {:?}", response);
    let (mut write, mut read) = ws_stream.split();
    // 发送消息任务
    tokio::spawn(async move {
        write.send(Message::Text("Hello, WebSocket!".to_string())).await.unwrap();
        loop {
            tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
            write.send(Message::Text("Ping".to_string())).await.unwrap();
        }
    });
    // 接收消息任务
    while let Some(msg) = read.next().await {
        match msg {
            Ok(msg) => {
                match msg {
                    Message::Text(text) => {
                        info!("Received message: {}", text);
                    }
                    Message::Binary(data) => {
                        info!("Received binary message ({} bytes)", data.len());
                    }
                    Message::Ping(data) => {
                        info!("Received ping");
                        write.send(Message::Pong(data)).await.unwrap();
                    }
                    Message::Pong(_) => {
                        info!("Received pong");
                    }
                    Message::Close(_) => {
                        info!("Connection closed");
                        break;
                    }
                }
            }
            Err(e) => {
                info!("Error receiving message: {}", e);
                break;
            }
        }
    }
    info!("Client disconnected");
    Ok(())
}

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    loop {
        if let Err(e) = connect_and_handle().await {
            info!("Connection error: {}, retrying in 5 seconds", e);
            sleep(Duration::from_secs(5)).await;
        }
    }
}

四、实战项目：构建实时聊天应用

4.1 项目需求与架构设计

我们将构建一个简单的实时聊天应用，支持以下功能：

• 多用户同时聊天
• 消息广播
• 用户加入/离开通知
• 心跳检测
• 连接超时

项目架构设计：

• 使用 Axum 作为 WebSocket 服务端
• 使用 tokio::sync::broadcast 实现消息广播
• 使用 tokio::sync::Mutex 管理用户信息
• 使用 HTML 和 JavaScript 实现客户端界面

4.2 服务端实现

创建 src/main.rs：

use axum::{routing::{get, post}, extract::{ws::{WebSocket, Message, WebSocketUpgrade}, State}, response::{IntoResponse, Html}, Router};
use tracing_subscriber::prelude::*;
use tracing::info;
use tokio::sync::{broadcast, Mutex};
use tokio::time::{interval, Duration};
use std::collections::HashMap;
use std::sync::Arc;

#[derive(Debug, Clone)]
struct User {
    id: String,
    name: String,
}
type Users = Arc<Mutex<HashMap<String, User>>>;

async fn index() -> Html<&'static str> {
    Html(r#"
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Real-Time Chat</title>
<style>
body { font-family: Arial, sans-serif; margin: 0; padding: 0; background-color: #f5f5f5; }
.container { max-width: 800px; margin: 0 auto; padding: 20px; }
h1 { color: #333; text-align: center; }
#chat { height: 400px; overflow-y: scroll; border: 1px solid #ddd; padding: 10px; margin-bottom: 10px; background-color: white; }
.message { margin-bottom: 10px; padding: 10px; background-color: #f0f0f0; border-radius: 5px; }
.system { color: #666; font-style: italic; }
.user { font-weight: bold; }
#message-input { width: 80%; padding: 10px; font-size: 14px; }
#send-btn { padding: 10px 20px; font-size: 14px; background-color: #007bff; color: white; border: none; border-radius: 5px; cursor: pointer; }
#send-btn:hover { background-color: #0056b3; }
</style>
</head>
<body>
<div>
<h1>Real-Time Chat</h1>
<div></div>
<input type="text" placeholder="Enter your message">
<button>Send</button>
</div>
<script>
const ws = new WebSocket('ws://' + location.host + '/ws');
const chat = document.getElementById('chat');
const messageInput = document.getElementById('message-input');
const sendBtn = document.getElementById('send-btn');
ws.onmessage = (event) => {
    const data = JSON.parse(event.data);
    const messageDiv = document.createElement('div');
    messageDiv.className = 'message';
    if (data.type === 'system') {
        messageDiv.className += ' system';
        messageDiv.textContent = data.content;
    } else if (data.type === 'user') {
        messageDiv.innerHTML = `<span>${data.user}:</span> ${data.content}`;
    }
    chat.appendChild(messageDiv);
    chat.scrollTop = chat.scrollHeight;
};
sendBtn.addEventListener('click', () => {
    const message = messageInput.value.trim();
    if (message) {
        ws.send(JSON.stringify({ type: 'message', content: message }));
        messageInput.value = '';
    }
});
messageInput.addEventListener('keypress', (e) => {
    if (e.key === 'Enter') {
        sendBtn.click();
    }
});
</script>
</body>
</html>
"#)
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    users: State<Users>,
    tx: State<broadcast::Sender<String>>,
) -> impl IntoResponse {
    info!("New WebSocket connection");
    let user_id = uuid::Uuid::new_v4().to_string();
    let user_name = format!("User{}", user_id.chars().take(8).collect::<String>());
    users.lock().await.insert(user_id.clone(), User { id: user_id.clone(), name: user_name.clone() });
    let rx = tx.get_ref().subscribe();
    ws.on_upgrade(|socket| handle_socket(socket, users, tx, rx, user_id, user_name))
}

async fn handle_socket(
    mut socket: WebSocket,
    users: State<Users>,
    tx: State<broadcast::Sender<String>>,
    mut rx: broadcast::Receiver<String>,
    user_id: String,
    user_name: String,
) {
    // 发送系统消息通知用户加入
    let join_message = serde_json::json!({"type": "system", "content": format!("{} joined the chat", user_name)});
    tx.get_ref().send(serde_json::to_string(&join_message).unwrap()).unwrap();
    // 发送欢迎消息
    let welcome_message = serde_json::json!({"type": "system", "content": format!("Welcome, {}!", user_name)});
    socket.send(Message::Text(serde_json::to_string(&welcome_message).unwrap())).await.unwrap();
    // 心跳检测
    let mut heartbeat_interval = interval(Duration::from_secs(10));
    // 广播接收任务
    let broadcast_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if let Err(e) = socket.send(Message::Text(msg)).await {
                info!("Error sending broadcast message: {}", e);
                break;
            }
        }
    });
    // 消息处理任务
    let message_task = tokio::spawn(async move {
        while let Some(msg) = socket.recv().await {
            match msg {
                Ok(msg) => {
                    match msg {
                        Message::Text(text) => {
                            info!("Received message: {}", text);
                            let data: serde_json::Value = serde_json::from_str(&text).unwrap();
                            if data["type"] == "message" {
                                let user_message = serde_json::json!({"type": "user", "user": user_name, "content": data["content"]});
                                tx.get_ref().send(serde_json::to_string(&user_message).unwrap()).unwrap();
                            }
                        }
                        Message::Binary(data) => {
                            info!("Received binary message ({} bytes)", data.len());
                        }
                        Message::Ping(data) => {
                            info!("Received ping");
                            socket.send(Message::Pong(data)).await.unwrap();
                        }
                        Message::Pong(_) => {
                            info!("Received pong");
                        }
                        Message::Close(_) => {
                            info!("Connection closed");
                            break;
                        }
                    }
                }
                Err(e) => {
                    info!("Error receiving message: {}", e);
                    break;
                }
            }
        }
    });
    // 心跳检测任务
    let heartbeat_task = tokio::spawn(async move {
        loop {
            tokio::select! {
                _ = heartbeat_interval.tick() => {
                    if let Err(e) = socket.send(Message::Ping(vec![])).await {
                        info!("Error sending ping: {}", e);
                        break;
                    }
                }
                _ = message_task => {
                    break;
                }
                _ = broadcast_task => {
                    break;
                }
            }
        }
    });
    // 等待所有任务完成
    let _ = tokio::try_join!(heartbeat_task, message_task, broadcast_task);
    // 发送系统消息通知用户离开
    let leave_message = serde_json::json!({"type": "system", "content": format!("{} left the chat", user_name)});
    tx.get_ref().send(serde_json::to_string(&leave_message).unwrap()).unwrap();
    // 从用户列表中删除
    users.lock().await.remove(&user_id);
    info!("WebSocket connection closed");
}

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new("info"))
        .with(tracing_subscriber::fmt::layer())
        .init();
    let users = Arc::new(Mutex::new(HashMap::new()));
    let (tx, _) = broadcast::channel(100);
    let app = Router::new()
        .route("/", get(index))
        .route("/ws", get(ws_handler))
        .with_state(users)
        .with_state(tx);
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    info!("WebSocket server running on http://0.0.0.0:3000");
    axum::serve(listener, app).await.unwrap();
}

在 Cargo.toml 中添加依赖：

[dependencies]
axum = { version = "0.5", features = ["ws"] }
tokio = { version = "1.0", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
uuid = { version = "1.1", features = ["v4"] }

4.3 客户端实现

客户端界面已经包含在服务端的 index 函数中，使用 HTML 和 JavaScript 实现。当用户访问 http://localhost:3000 时，会显示聊天界面。

4.4 性能测试与优化

我们可以使用 ab（Apache Bench）工具测试服务端的 HTTP 性能：

ab -n 1000 -c 100 http://localhost:3000/

或者使用 wrk 工具测试 WebSocket 连接的性能：

wrk -t 12 -c 400 -d 30s http://localhost:3000/ws

性能优化方法：

• 使用连接池：对于数据库等资源，使用连接池可以避免频繁创建和销毁连接。
• 优化消息处理：减少消息处理的耗时，提高处理效率。
• 使用压缩算法：对消息进行压缩，减少传输数据量。
• 使用负载均衡：对于高并发场景，使用负载均衡可以分散请求压力。

五、常见问题与最佳实践

5.1 WebSocket 连接建立失败

问题：客户端无法与服务器建立 WebSocket 连接。

解决方案：

1. 检查服务器是否正在运行。
2. 检查服务器地址和端口是否正确。
3. 检查防火墙是否阻止了 WebSocket 连接。
4. 检查服务器的 CORS 策略是否允许跨域请求。

5.2 消息发送失败

问题：客户端发送消息失败。

解决方案：

1. 检查 WebSocket 连接是否已经关闭。
2. 检查消息格式是否正确。
3. 检查服务器的消息处理逻辑是否有错误。
4. 检查网络连接是否正常。

5.3 连接超时

问题：WebSocket 连接在一段时间后自动断开。

解决方案：

1. 实现心跳检测与连接超时机制。
2. 检查服务器的超时设置。
3. 检查网络连接是否稳定。

5.4 消息乱序

问题：客户端接收到的消息顺序不正确。

解决方案：

1. 使用可靠的消息队列实现消息广播。
2. 在消息中添加时间戳或序列号，客户端根据时间戳或序列号进行排序。

5.5 性能问题

问题：WebSocket 服务的性能不佳。

解决方案：

1. 优化服务器的消息处理逻辑。
2. 使用连接池和异步编程提高处理效率。
3. 使用负载均衡分散请求压力。
4. 对消息进行压缩，减少传输数据量。

六、总结

WebSocket 协议为实时应用提供了高效的通信方式，Rust 的异步编程能力使得构建高性能 WebSocket 服务变得简单。在本章中，我们介绍了 WebSocket 协议的基本概念、异步服务端和客户端的开发，以及实战项目的实现。

我们使用 Axum 框架实现了 WebSocket 服务端，支持消息广播、心跳检测和连接超时机制。我们还使用 Tungstenite 库实现了异步 WebSocket 客户端，支持重连机制。最后，我们构建了一个实时聊天应用，展示了 WebSocket 在实际项目中的应用。

通过学习本章内容，读者可以掌握 Rust 异步编程中 WebSocket 服务的开发方法，并了解常见问题和最佳实践。希望读者能够将这些知识应用到实际项目中，构建高性能的实时应用。