📚 课程目录第4章非阻塞TCP服务器

4. 非阻塞TCP服务器

任务说明

实现一个基于非阻塞socket与poll系统调用的并发TCP回显服务器。服务器监听指定端口，使用单个线程通过poll同时管理多个客户端连接，对每个客户端收到的数据原样返回。该案例用于演示Rust中非阻塞I/O与事件驱动编程模型，不引入async/await或第三方异步运行时。

核心实现

use std::io::{self, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::os::unix::io::{AsRawFd, RawFd};
use std::collections::HashMap;
use libc::pollfd;

const POLL_TIMEOUT_MS: i32 = 1000;

fn set_nonblocking(stream: &TcpStream) -> io::Result<()> {
    let fd = stream.as_raw_fd();
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    let result = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
    if result == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

fn main() -> io::Result<()> {
    let addr = "127.0.0.1:8080";
    let listener = TcpListener::bind(addr)?;
    listener.set_nonblocking(true)?;
    println!("非阻塞TCP服务器监听 {}", addr);

    // fd -> (TcpStream, 读缓冲区)
    let mut clients: HashMap<RawFd, (TcpStream, Vec<u8>)> = HashMap::new();
    let listener_fd = listener.as_raw_fd();

    loop {
        // 构建pollfd数组
        let mut pollfds: Vec<pollfd> = Vec::new();
        pollfds.push(pollfd {
            fd: listener_fd,
            events: libc::POLLIN,
            revents: 0,
        });
        for fd in clients.keys() {
            pollfds.push(pollfd {
                fd: *fd,
                events: libc::POLLIN,
                revents: 0,
            });
        }

        // 等待事件
        let ret = unsafe { libc::poll(pollfds.as_mut_ptr(), pollfds.len() as libc::nfds_t, POLL_TIMEOUT_MS) };
        if ret == -1 {
            return Err(io::Error::last_os_error());
        }
        if ret == 0 {
            continue; // 超时，继续循环
        }

        // 处理监听socket上的新连接
        if pollfds[0].revents & libc::POLLIN != 0 {
            loop {
                match listener.accept() {
                    Ok((stream, peer_addr)) => {
                        set_nonblocking(&stream)?;
                        let fd = stream.as_raw_fd();
                        println!("新连接: {} (fd={})", peer_addr, fd);
                        clients.insert(fd, (stream, Vec::new()));
                    }
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => break,
                    Err(e) => return Err(e),
                }
            }
        }

        // 处理客户端socket上的数据
        for pollfd in pollfds.iter().skip(1) {
            if pollfd.revents & libc::POLLIN != 0 {
                let fd = pollfd.fd;
                if let Some((stream, buf)) = clients.get_mut(&fd) {
                    let mut tmp = [0u8; 1024];
                    loop {
                        match stream.read(&mut tmp) {
                            Ok(0) => {
                                // 连接关闭
                                println!("连接关闭: fd={}", fd);
                                clients.remove(&fd);
                                break;
                            }
                            Ok(n) => {
                                buf.extend_from_slice(&tmp[..n]);
                                // 回显所有已读数据
                                if let Err(e) = stream.write_all(buf) {
                                    eprintln!("写入错误 fd={}: {}", fd, e);
                                    clients.remove(&fd);
                                    break;
                                }
                                buf.clear();
                            }
                            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => break,
                            Err(e) => {
                                eprintln!("读取错误 fd={}: {}", fd, e);
                                clients.remove(&fd);
                                break;
                            }
                        }
                    }
                }
            }
        }
    }
}

依赖配置

在Cargo.toml中添加libc依赖：

[dependencies]
libc = "0.2"

运行与验证

启动服务器：

cargo run

预期输出：非阻塞TCP服务器监听 127.0.0.1:8080

验证并发能力：打开两个终端，分别执行：

# 终端1
nc 127.0.0.1 8080
# 输入 hello，应收到 hello

# 终端2
nc 127.0.0.1 8080
# 输入 world，应收到 world

判定标准：

两个客户端可同时保持连接，互不阻塞；
任一客户端发送数据，服务器立即原样返回；
关闭任一客户端，服务器不退出，继续服务其他连接；
服务器终端打印新连接的地址与fd编号，以及连接关闭时的fd编号。

代码说明

服务器启动后将监听socket设为非阻塞，主循环每次构建包含监听fd与所有客户端fd的pollfd数组，调用poll等待事件。当监听fd可读时，循环accept直到返回WouldBlock，确保接收完所有待处理连接。当客户端fd可读时，循环读取直到WouldBlock，将读到的数据原样写回。连接关闭（读返回0）或发生错误时，从clients表中移除对应fd。整个服务器仅使用单线程，通过poll实现并发，无第三方异步依赖。