在Future的poll方法中使用Mio轮询无法获取TCP数据

核心问题原因 #

1. Future的poll方法被阻塞：你在poll里直接调用mio::Poll::poll并设置5秒超时，这会阻塞当前工作线程，完全违背异步Runtime的设计——Runtime需要线程能快速处理多个任务，不能被单个Future卡住。
2. 未关联Mio Waker与Future Waker：Mio检测到事件后，无法主动通知Runtime唤醒你的Future，只能等超时或者线程被其他任务唤醒，导致事件无法及时被处理。

修复方案 #

关键调整点 #

• 永远不要在Future::poll中阻塞调用，改用非阻塞轮询Mio事件。
• 使用Mio的Waker将Future的任务唤醒逻辑与Mio的事件通知绑定，确保Mio检测到事件时能立刻触发Future的重新poll。
• 每次poll时更新Mio Waker关联的Future Waker，因为任务可能在不同线程间转移。

修改后的代码 #

首先更新结构体和初始化逻辑：

use mio::net::{TcpListener, TcpStream};
use mio::{Events, Interest, Poll as MioPoll, Token, Waker as MioWaker};
use std::io::{self, Read, Write};
use std::time::Duration;
use std::error::Error;

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

const SERVER: Token = Token(0);
const WAKER: Token = Token(1); // 用于Mio Waker的Token

struct ServerWorker {
    server: TcpListener,
    poll: MioPoll,
    mio_waker: MioWaker,
}

impl ServerWorker {
    pub fn new() -> Result<Self, Box<dyn Error>> {
        let addr = "127.0.0.1:13265".parse()?;
        let mut server = TcpListener::bind(addr)?;
        let poll: MioPoll = MioPoll::new()?;
        
        // 注册TCP监听器
        poll.registry()
            .register(&mut server, SERVER, Interest::READABLE)?;
        
        // 创建Mio Waker并注册到Poll中
        let mio_waker = MioWaker::new(poll.registry(), WAKER)?;
        
        Ok(ServerWorker {
            server,
            poll,
            mio_waker,
        })
    }
}

然后修改Future实现：

impl Future for ServerWorker {
    // 调整Output为Result，方便处理错误
    type Output = Result<String, Box<dyn Error>>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut events = Events::with_capacity(128);
        
        // 非阻塞轮询Mio事件：超时设为None表示立刻返回
        match self.poll.poll(&mut events, None) {
            Ok(_) => {
                for event in events.iter() {
                    match event.token() {
                        SERVER => {
                            // 处理新连接
                            let (mut stream, _addr) = self.server.accept()?;
                            let mut received_data = [0; 4096];
                            // 读取实际收到的字节数
                            let bytes_read = stream.read(&mut received_data)?;
                            if bytes_read == 0 {
                                // 连接关闭
                                return Poll::Pending;
                            }
                            let received = String::from_utf8_lossy(&received_data[..bytes_read]).to_string();
                            // 如果只需要处理一次连接，返回Ready；否则处理完后继续返回Pending
                            return Poll::Ready(Ok(received));
                        }
                        WAKER => {
                            // Mio Waker触发，无需额外处理，继续轮询即可
                        }
                        _ => unreachable!(),
                    }
                }
            }
            Err(e) => {
                return Poll::Ready(Err(e.into()));
            }
        }

        // 没有事件，更新Mio Waker关联的当前Future Waker
        self.mio_waker.wake_by_ref().unwrap();
        // 注册当前Waker，确保Mio有事件时能唤醒这个任务
        cx.waker().wake_by_ref();
        
        Poll::Pending
    }
}

额外注意事项 #

1. 错误处理：将Future的Output改为Result，避免直接unwrap导致程序panic，更符合异步编程的错误处理规范。
2. 连接处理：如果需要持续监听新连接，处理完当前连接后不要直接返回Ready，而是继续返回Pending，这样Future会保持活跃，继续处理后续连接。
3. 线程安全：确保ServerWorker实现Send和Sync（Mio的类型默认都是Send/Sync的，所以这里没问题），因为你的Runtime会在多线程间调度任务。
4. Mio Interest模式：默认是水平触发，所以只要套接字处于可读状态，Mio会持续触发事件；如果需要边缘触发，可以修改Interest::READABLE | Interest::EDGE_TRIGGERED，但需要确保处理完所有可用数据。

为什么之前的代码不工作？ #

• 阻塞的poll调用占用了工作线程，导致Runtime无法处理其他任务，同时Mio的事件无法及时传递。
• 没有绑定Mio Waker，Mio检测到事件后无法通知Runtime唤醒你的Future，只能等超时或者线程被其他任务唤醒，这时候事件可能已经被处理或者丢失。

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