Rust中AsyncWriteExt写入TcpStream字节丢失致阻塞问题

我用tokio-yamux实现了一个TCP消息服务器，这个库只允许用AsyncWriteExt和AsyncReadExt读写字节，符合我的需求。但发送完整消息时出现字节丢失，导致整个应用阻塞：客户端会先发送包含长度的头部，再发X字节内容，让服务器明确要读多少字节。当尝试写入262144字节时，write_all或write方法只写了262136字节，头部标明的总长度是262144，服务器一直在等剩余字节，直接阻塞了，我找不到原因。

客户端代码 #

use futures::prelude::*;
use std::{error::Error, vec};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};
use tokio_yamux::{config::Config, session::Session};

#[tokio::main]
pub async fn main() -> Result<(), Box<dyn Error>> {
    let socket = TcpStream::connect("127.0.0.1:8080").await?;
    println!("[client] connected to server: {:?}", socket.peer_addr()?);
    println!("created stream");

    let mut session = Session::new_client(socket, Config::default());
    let ctrl = session.control();

    let mut handles = Vec::new();

    tokio::spawn(async move {
        loop {
            match session.next().await {
                Some(Ok(_)) => (),
                Some(Err(e)) => {
                    println!("{}", e);
                    break;
                }
                None => {
                    println!("closed");
                    break;
                }
            }
        }
    });

    for _i in 0..1 {
        let mut ctrl_clone = ctrl.clone();
        handles.push(tokio::spawn(async move {
            match ctrl_clone.open_stream().await {
                Ok(mut stream) => {
                    // 标识操作与队列
                    let operation_id = 0;
                    let queue_id = Some(2);

                    let data_to_send = match queue_id {
                        Some(id) => vec![operation_id, id],
                        None => vec![operation_id],
                    };

                    let data_to_send: Vec<u8> = data_to_send
                        .into_iter()
                        .flat_map(|x| i32::to_be_bytes(x))
                        .collect();

                    stream.write_all(&data_to_send).await.unwrap();
                    stream.flush().await.unwrap();

                    // 接收服务器确认
                    let mut buf = [1; 4];
                    let sv_code: i32;
                    loop {
                        match stream.read_exact(&mut buf).await {
                            Ok(0) => {}
                            Ok(n) => {
                                println!("Client: Reading Buffer: n_bytes {:?}", n);
                                sv_code = i32::from_be_bytes(buf);
                                break;
                            }
                            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                                continue;
                            }
                            Err(_e) => {}
                        }
                    }

                    // 发送大消息
                    let data: [u8; 262140] = [1; 262140];
                    let mut vec_data = data.to_vec();
                    
                    let len = data.len() as u32;
                    let len_slices = len.to_be_bytes();

                    for slice in len_slices {
                        vec_data.insert(0, slice);
                    }

                    println!("Total_len: {:?}", vec_data.len());
                    stream.write_all(&vec_data).await.unwrap();
                    println!("N: write_all completed");
                    stream.flush().await.unwrap();

                    println!("Fin write");
                }
                Err(e) => {
                    println!("{:?}", e);
                }
            }
        }));
    }

    for handle in handles {
        let _ = handle.await;
    }

    Ok(())
}

服务器代码 #

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

use std::error::Error;

use futures::prelude::*;
use tokio_yamux::{config::Config, session::Session};

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let addr = "127.0.0.1:8080".to_string();

    let listener = TcpListener::bind(&addr).await.unwrap();

    while let Ok((socket, _)) = listener.accept().await {
        println!("accepted a socket: {:?}", socket.peer_addr());
        let mut session = Session::new_server(socket, Config::default());
        tokio::spawn(async move {
            while let Some(Ok(mut stream)) = session.next().await {
                println!("Server accept a stream from client: id={}", stream.id());
                tokio::spawn(async move {
                    // 读取操作码与队列ID
                    let op_code = stream.read_i32().await.unwrap();
                    println!("op_id: {:?}", op_code);

                    let queue_id = Some(stream.read_i32().await.unwrap());
                    println!("queue_id: {:?}", queue_id);

                    // 发送确认信号
                    let sv_code: i32 = 0;
                    stream.write_all(&sv_code.to_be_bytes()).await.unwrap();
                    stream.flush().await.unwrap();

                    // 接收消息逻辑
                    let mut total_bytes = 0;
                    let mut n_bytes_read = 0;
                    let mut len_slices: [u8; 4] = [0; 4];
                    
                    loop {
                        println!("n_bytes read: {:?}", n_bytes_read);
                        let mut capacity = 65535;
                
                        if n_bytes_read == 0 {
                            capacity = 65539;
                        }
                
                        let mut buffer = Vec::with_capacity(capacity);
                                
                        println!("Blocked?");
                        match stream.read_buf(&mut buffer).await {
                            Ok(0) => continue,
                            Ok(n) => {
                                println!("N: {:?}", n);
                                if n_bytes_read == 0 {
                                    // 提取长度头部
                                    for i in 0..4 {
                                        len_slices[i] = buffer.remove(0);
                                    }
                                    // 统一字节序为大端
                                    total_bytes = u32::from_be_bytes(len_slices);
                                    total_bytes += 4;
                                    // 移除头部后，实际接收的内容字节数是n-4
                                    n_bytes_read += (n - 4);
                                } else {
                                    n_bytes_read += n;
                                }
                
                                if n_bytes_read == total_bytes.try_into().unwrap() {
                                    break;
                                }
                            }
                            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                                println!("Err: TCP -> SV (Write))");
                                continue;
                            }
                            Err(e) => {
                                break;
                            }
                        }
                    }

                    println!("Finished");
                });
            }
        });
    }

    Ok(())
}

原输出 #

客户端 #

[client] connected to server: 127.0.0.1:8080
created stream
Client: Reading Buffer: n_bytes 4
Total_len: 262144
Blocked here

服务器 #

accepted a socket: Ok(127.0.0.1:52430)
Server accept a stream from client: id=1
n_bytes read: 0
Blocked?
N: 65539
n_bytes read: 65539
Blocked?
N: 65535
n_bytes read: 131074
Blocked?
N: 65535
n_bytes read: 196609
Blocked?
N: 65527
n_bytes read: 262136
Blocked?

问题根源与修复 #

1. 字节序不匹配 #

客户端用to_be_bytes()生成长度头部，但服务器用from_le_bytes()解析，导致解析出的总长度完全错误，后续判断接收完成的条件失效。

修复：
客户端和服务器统一使用大端字节序：

• 客户端保持len.to_be_bytes()
• 服务器将解析代码改为u32::from_be_bytes(len_slices)

2. 接收计数错误 #

服务器从buffer中移除4字节头部后，没有将这4字节从n_bytes_read的累加值中扣除，导致n_bytes_read的数值比实际接收的内容多4，永远达不到错误的总长度，引发阻塞。

修复：
在移除头部后，将n_bytes_read的累加值改为n - 4，而非直接加n。

3. 冗余代码清理 #

服务器中的buffer.truncate(n)是冗余操作，read_buf返回的n就是实际读取的字节数，buffer的长度已经等于n，无需截断。

4. 客户端返回值误解 #

write_all方法返回的是Result<(), Error>，不是写入的字节数，原代码中let n = stream.write_all(&vec_data).await.unwrap();获取的n是()，无实际意义，可直接删除该变量。

内容的提问来源于stack exchange，提问作者Tutu ツ