该示例收录于 Axum 官方 git 库，地址为 https://github.com/tokio-rs/axum/tree/main/examples/chat，可以直接在该地址下载最新版本的代码。

Cargo.toml

[package]
name = "example-chat"
version = "0.1.0"
edition = "2021"
publish = false

[dependencies]
axum = { path = "../../axum", features = ["ws"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

chat.html

<!DOCTYPE html>
<html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>WebSocket Chat</title>
    </head>
    <body>
        <h1>WebSocket Chat Example</h1>
<!-- 用户输入用户 -->
        <input id="username" style="display:block; width:100px; box-sizing: border-box" type="text" placeholder="username">
<!-- 该按钮点击之后，即加入聊天室 -->
        <button id="join-chat" type="button">Join Chat</button>
<!-- 该 textarea 只用于展示聊天室内容 -->
        <textarea id="chat" style="display:block; width:600px; height:400px; box-sizing: border-box" cols="30" rows="10"></textarea>
<!-- 聊天内容 -->
        <input id="input" style="display:block; width:600px; box-sizing: border-box" type="text" placeholder="chat">

        <script>
            const username = document.querySelector("#username");
            const join_btn = document.querySelector("#join-chat");
            const textarea = document.querySelector("#chat");
            const input = document.querySelector("#input");

            join_btn.addEventListener("click", function(e) {

// 用户输入用户名，申请加入聊天室之后，立马将该按钮置灰，不允许重复点击
                this.disabled = true;

// 连接服务端，并创建一个新的 WebSocket 实例
                const websocket = new WebSocket("ws://localhost:3001/websocket");

// 当 WebSocket 连接成功之后，立马发送用户名信息
                websocket.onopen = function() {
                    console.log("connection opened");
                    websocket.send(username.value);
                }

                const btn = this;
// WebSocket 关闭时，解除当前按钮的禁用状态
                websocket.onclose = function() {
                    console.log("connection closed");
                    btn.disabled = false;
                }

// 当收到新消息时，直接将消息添加至 textarea 的最后面
                websocket.onmessage = function(e) {
                    console.log("received message: "+e.data);
                    textarea.value += e.data+"\r\n";
                }

// 当用户按下回车键时，发送用户输入的消息，并清空输入框
                input.onkeydown = function(e) {
                    if (e.key == "Enter") {
                        websocket.send(input.value);
                        input.value = "";
                    }
                }
            });
        </script>
    </body>
</html>
//! Example chat application.
//!
//! Run with
//!
//! ```not_rust
//! cargo run -p example-chat
//! ```

use axum::{
    extract::{
        ws::{Message, Utf8Bytes, WebSocket, WebSocketUpgrade},
        State,
    },
    response::{Html, IntoResponse},
    routing::get,
    Router,
};
use futures_util::{sink::SinkExt, stream::StreamExt};
use std::{
    collections::HashSet,
    sync::{Arc, Mutex},
};
use tokio::sync::broadcast;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

// Our shared state
// 共享的状态
// 
// 我们需要用户名是唯一的，这个用于跟踪哪些用户名是已经被占用了。所以在这里面，使用了：
//
// - HashSet 集合数据结构，存储不重复的值，查找/插入
// - Mutext 互斥锁，保证多线程安全访问
// 
// `tokio::sync::broadcast` 是一个**多生产者多消费者**的广播通道，一个消息可以被所有订阅者接收，非常适合"向所有已连接客户端发送消息"的场景。有具有以下特性
// 
// - **一对多**：一条消息 → 所有 Receiver
// - **有界容量**：创建时指定缓冲区大小
// - **滞后处理**：慢消费者会被踢掉（Lagged 错误）
// - **Clone**：Sender 和 Receiver 都可以克隆
struct AppState {
    // We require unique usernames. This tracks which usernames have been taken.
    user_set: Mutex<HashSet<String>>,
    // Channel used to send messages to all connected clients.
    tx: broadcast::Sender<String>,
}

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| format!("{}=trace", env!("CARGO_CRATE_NAME")).into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    // Set up application state for use with with_state().
    let user_set = Mutex::new(HashSet::new());
    let (tx, _rx) = broadcast::channel(100);

    let app_state = Arc::new(AppState { user_set, tx });

    let app = Router::new()
        .route("/", get(index))
        .route("/websocket", get(websocket_handler))
        .with_state(app_state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3001")
        .await
        .unwrap();
    tracing::debug!("listening on {}", listener.local_addr().unwrap());
    axum::serve(listener, app).await;
}

async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    ws.on_upgrade(|socket| websocket(socket, state))
}

// This function deals with a single websocket connection, i.e., a single
// connected client / user, for which we will spawn two independent tasks (for
// receiving / sending chat messages).
async fn websocket(stream: WebSocket, state: Arc<AppState>) {
    // By splitting, we can send and receive at the same time.
    let (mut sender, mut receiver) = stream.split();

    // Username gets set in the receive loop, if it's valid.
    let mut username = String::new();
    // Loop until a text message is found.
    while let Some(Ok(message)) = receiver.next().await {
        if let Message::Text(name) = message {
            // If username that is sent by client is not taken, fill username string.
            check_username(&state, &mut username, name.as_str());

            // If not empty we want to quit the loop else we want to quit function.
            if !username.is_empty() {
                break;
            } else {
                // Only send our client that username is taken.
                let _ = sender
                    .send(Message::Text(Utf8Bytes::from_static(
                        "Username already taken.",
                    )))
                    .await;

                return;
            }
        }
    }

    // We subscribe *before* sending the "joined" message, so that we will also
    // display it to our client.
    let mut rx = state.tx.subscribe();

    // Now send the "joined" message to all subscribers.
    let msg = format!("{username} joined.");
    tracing::debug!("{msg}");
    let _ = state.tx.send(msg);

    // Spawn the first task that will receive broadcast messages and send text
    // messages over the websocket to our client.
    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            // In any websocket error, break loop.
            if sender.send(Message::text(msg)).await.is_err() {
                break;
            }
        }
    });

    // Clone things we want to pass (move) to the receiving task.
    let tx = state.tx.clone();
    let name = username.clone();

    // Spawn a task that takes messages from the websocket, prepends the user
    // name, and sends them to all broadcast subscribers.
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(Message::Text(text))) = receiver.next().await {
            // Add username before message.
            let _ = tx.send(format!("{name}: {text}"));
        }
    });

    // If any one of the tasks run to completion, we abort the other.
    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    };

    // Send "user left" message (similar to "joined" above).
    let msg = format!("{username} left.");
    tracing::debug!("{msg}");
    let _ = state.tx.send(msg);

    // Remove username from map so new clients can take it again.
    state.user_set.lock().unwrap().remove(&username);
}

fn check_username(state: &AppState, string: &mut String, name: &str) {
    let mut user_set = state.user_set.lock().unwrap();

    if !user_set.contains(name) {
        user_set.insert(name.to_owned());

        string.push_str(name);
    }
}

// Include utf-8 file at **compile** time.
async fn index() -> Html<&'static str> {
    Html(std::include_str!("../chat.html"))
}

在上面的示例中，#[tokio::main]的作用是什么样的？

核心作用

#[tokio::main] 是一个属性宏（attribute macro），它会自动帮你搭建 Tokio 异步运行时，让你能直接写 async fn main()。

展开前后对比

我们写下下面这样的代码：

#[tokio::main]
async fn main() {
    println!("Hello from async!");
    
    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    
    println!("Done!");
}

宏实际展开后（大致等价于）：

fn main() {
    tokio::runtime::Runtime::new()
        .unwrap()
        .block_on(async {
            println!("Hello from async!");
            
            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
            
            println!("Done!");
        });
}

它具体做了三件事：

┌─────────────────────────────────────────────┐
│            #[tokio::main]                   │
├─────────────────────────────────────────────┤
│                                             │
│  1. 创建 Tokio Runtime（异步运行时）           │
│     ├─ 线程池（worker threads）               │
│     ├─ I/O 事件循环（epoll/kqueue/iocp）      │
│     └─ 定时器                                │
│                                             │
│  2. 把你的 async fn main() 包进 block_on()    │
│     → 在 Runtime 上执行这个 Future            │
│                                             │
│  3. 等待主 Future 完成，然后清理退出            │
│                                             │
└─────────────────────────────────────────────┘

没有 #[tokio::main] 会怎样？

// ❌ 编译错误！
async fn main() {
    // .await 只能在 async 上下文中使用，
    // 但标准 Rust 的 main 必须是 fn main()，不能是 async fn main()
    tokio::time::sleep(...).await;
}
error[E0752]: `main` function is not allowed to be `async`
 --> src/main.rs:1:1

  |
1 | async fn main() {

  | ^^^^^ `main` function is not allowed to be `async`

带配置参数的版本

// 默认：多线程运行时
#[tokio::main]
async fn main() { }

// 等价于：
#[tokio::main(flavor = "multi_thread")]
async fn main() { }

// 单线程运行时（不需要 Sync，开销更小）
#[tokio::main(flavor = "current_thread")]
async fn main() { }

// 自定义工作线程数
#[tokio::main(flavor = "multi_thread", worker_threads = 4)]
async fn main() { }

// 完整写法
#[tokio::main]
async fn main() {
    // 等价于手动：
    // let rt = tokio::runtime::Builder::new_multi_thread()
    //     .worker_threads(num_cpus::get())  // 默认用所有 CPU 核心
    //     .enable_all()                     // 启用 I/O 和定时器
    //     .build()
    //     .unwrap();
    // rt.block_on(async { ... });
}

可视化理解

程序启动
  │
  ▼
#[tokio::main] 创建 Runtime
  │
  ├── Worker Thread 0 ────┐
  ├── Worker Thread 1 ────┤
  ├── Worker Thread 2 ────┤  ← 线程池
  └── ...                 ┘
  │
  ▼
block_on(async { main() })  ← 主 task 在这个 Runtime 上运行
  │
  │   主 task 中 spawn 子任务：
  │
  ├── tokio::spawn(处理客户端1)  ← 分配到线程池
  ├── tokio::spawn(处理客户端2)  
  ├── tokio::spawn(广播消息)    
  └── ...
  │
  ▼
main() 返回后，Runtime 等待所有任务完成，然后退出

在聊天室示例中的意义

回顾之前的聊天室代码：

#[tokio::main]  // ← 创建异步运行时
async fn main() {
    let chat = Arc::new(ChatRoom::new());
    let listener = TcpListener::bind("127.0.0.1:8080").await.unwrap();
    //           ▲                           ▲
    //           │                           └── .await 需要 Runtime
    //           └── TcpListener 需要 Runtime 的 I/O 事件循环

    loop {
        let (socket, addr) = listener.accept().await.unwrap();
        //                            ▲
        //                            └── 非阻塞等待连接，需要 Runtime
        
        tokio::spawn(async move {
            // ▲ 在 Runtime 的线程池中并发运行
            // 多个客户端可以同时被处理
        });
    }
}

没有 #[tokio::main]，上述所有 .await 和 tokio::spawn 都无法工作。

问题答案
它是什么？属性宏，让你写 async fn main()
它做了什么？创建 Tokio Runtime + 执行你的 async 代码
为什么需要它？Rust 原生不支持 async main，需要 Runtime
能手动替代吗？可以，手动创建 Runtime 然后 block_on

一句话概括：#[tokio::main] 就是帮你把异步代码"接"到同步世界的那个桥梁。

标签: Rust, WebSocket, Chat, 聊天室

添加新评论