tokio_tungstenite多所有权wss流异步循环与写入权限问题

需求概述 #

• 获取可读写的WebSocket流，先向服务器发送初始内容启动数据接收
• 运行一个非阻塞的无限后台循环，读取服务器响应并在收到特定消息时自动回写
• 主程序能根据信号，复用同一个WebSocket写流发送额外数据

遇到的问题 #

• 若将读写端都移入tokio::spawn的任务，主程序会失去写流所有权，无法完成第三点需求
• 若不await tokio::spawn返回的JoinHandle，后台任务可能因主程序提前退出而无法执行

解决方案 #

核心解决思路是实现写流的安全共享和保障主程序生命周期：

1. 用Arc<Mutex<WriteHalf>>包装写流，通过引用计数和互斥锁实现多任务安全访问，避免所有权转移导致主程序丢失写流控制权
2. 让主程序保持运行状态（如等待外部信号），确保后台任务有足够执行时间；最后等待后台任务收尾，避免资源泄漏

修正后的完整代码示例 #

use tokio::sync::{Arc, Mutex};
use tokio_tungstenite::{connect_async, tungstenite::Message};

#[tokio::main]
async fn main() {
    // 1. 建立WebSocket连接并拆分读写流
    let (stream, _) = connect_async("wss://your-websocket-url.com").await.unwrap();
    let (write, read) = stream.split();
    
    // 用Arc<Mutex>包装写流，实现多任务安全共享
    let shared_write = Arc::new(Mutex::new(write));
    
    // 发送初始数据启动会话
    let mut locked_write = shared_write.lock().await;
    locked_write.send(Message::Text("initial-message".to_string())).await.unwrap();
    drop(locked_write); // 主动释放锁，避免后台任务阻塞

    // 2. 启动后台读写循环任务
    let read_handle = {
        let shared_write = Arc::clone(&shared_write);
        tokio::spawn(async move {
            while let Some(msg_result) = read.next().await {
                match msg_result {
                    Ok(Message::Text(text)) => {
                        match text.as_str() {
                            "ping" => {
                                let mut write = shared_write.lock().await;
                                write.send(Message::Text("pong".to_string())).await.unwrap();
                            }
                            "error" => break,
                            _ => println!("Received: {}", text),
                        }
                    }
                    Err(e) => {
                        eprintln!("Read error: {}", e);
                        break;
                    }
                    _ => {}
                }
            }
        })
    };

    // 3. 主程序逻辑：等待信号后发送额外数据
    // 示例：等待Ctrl+C信号触发发送操作
    tokio::signal::ctrl_c().await.unwrap();
    let mut write = shared_write.lock().await;
    write.send(Message::Text("shutdown-message".to_string())).await.unwrap();

    // 等待后台任务正常结束
    read_handle.await.unwrap();
}

关键说明 #

• 所有权共享机制：Arc提供线程安全的引用计数，允许后台任务和主程序持有写流的引用；Mutex确保同一时间只有一个任务能操作写流，避免并发写入冲突
• 任务执行保障：主程序通过tokio::signal::ctrl_c().await阻塞等待，避免提前退出导致后台任务被终止；最后await read_handle确保后台任务完成收尾工作
• 锁的主动释放：发送初始数据后主动drop(locked_write)释放锁，避免主程序长时间持有锁导致后台任务无法获取写流，引发死锁

内容的提问来源于stack exchange，提问作者Emad Akhras