Tokio 自定义 Future:手写一个异步状态机,理解 poll 的调用时机
Tokio 自定义 Future手写一个异步状态机理解 poll 的调用时机所以今天这篇文章我会从零开始手写一个自定义Future然后带你一步步画出它的状态转移图。相信我理解了 poll 的调用时机你对异步 Rust 的理解会上一个台阶。一、Future 的本质是一个状态机先看一个最简单的异步任务向一个 HTTP 接口发请求拿到响应后解析 JSON。用async/await写/// 一个简单的异步任务 /// 发HTTP请求 → 读body → 解析JSON async fn fetch_and_parse(url: str) - Resultserde_json::Value, Boxdyn std::error::Error { // 第一步发起请求这里会 yield让出CPU let resp reqwest::get(url).await?; // ← 状态切换点1 // 第二步读取完整响应体 let body resp.text().await?; // ← 状态切换点2 // 第三步解析 JSON let value: serde_json::Value serde_json::from_str(body)?; Ok(value) }看着很简单对吧但编译器或者说async关键字实际生成的代码大概长这样use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; /// 编译器为 fetch_and_parse 自动生成的状态机 /// 每一个 .await 点就是一个状态 enum FetchAndParseFuture { /// 状态0初始状态尚未开始 Start { url: String }, /// 状态1正在等 HTTP 响应回来 /// 内部保存着 reqwest::get 返回的 Future WaitingResponse { // 注意这里保存的是 Future不是结果 // Box::pin 是因为 Future 的大小在编译期不确定 future: PinBoxdyn FutureOutput Resultreqwest::Response, reqwest::Error, }, /// 状态2HTTP 响应已拿到正在读 body WaitingBody { resp: reqwest::Response, }, /// 状态3body 读完了等待 JSON 解析 /// 其实这个状态可以省略解析是同步的 /// 编译器为了通用性还是加了 Parsing { body: String }, /// 终态任务完成 Done, } impl Future for FetchAndParseFuture { type Output Resultserde_json::Value, Boxdyn std::error::Error; fn poll(mut self: Pinmut Self, cx: mut Context_) - PollSelf::Output { // poll 是状态机的推进器 // 每次被调用就根据当前状态决定下一步做什么 loop { match mut *self { FetchAndParseFuture::Start { url } { // 当前是初始状态 → 发起 HTTP 请求 let client reqwest::Client::new(); let req_future Box::pin(client.get(url.as_str()).send()); // 切换到等待响应状态 *self FetchAndParseFuture::WaitingResponse { future: req_future, }; // 继续循环进入下一个状态 continue; } FetchAndParseFuture::WaitingResponse { future } { // 当前在等 HTTP 响应 → poll 内部的 Future match future.as_mut().poll(cx) { Poll::Ready(Ok(resp)) { // 响应回来了切换到等body状态 *self FetchAndParseFuture::WaitingBody { resp }; continue; } Poll::Ready(Err(e)) { // 出错了直接返回错误 *self FetchAndParseFuture::Done; return Poll::Ready(Err(Box::new(e))); } Poll::Pending { // 还没好告诉调度器我还在等 // 调度器会记下等 waker 被触发再调 poll return Poll::Pending; } } } FetchAndParseFuture::WaitingBody { resp } { // 注意这里本来应该有一个读 body 的 Future // 简化起见假设 body 可以直接读 // 实际编译器会生成类似的 WaitingBodyFuture 状态 *self FetchAndParseFuture::Done; return Poll::Ready(Ok(serde_json::Value::Null)); } FetchAndParseFuture::Parsing { .. } { // 解析状态 → 同步操作直接完成 *self FetchAndParseFuture::Done; return Poll::Ready(Ok(serde_json::Value::Null)); } FetchAndParseFuture::Done { // 已经到了终态理论上不应该再被 poll panic!(Future 已经完成了不应该再次 poll); } } } } }看完这个代码我想你对async 函数是状态机这句话应该有感觉了。编译器的做法就是把async函数体切分成多个状态.await点就是状态转换点。二、poll 的调用时机手画状态转移图我们来画一张图看看 poll 到底在什么时候、由谁调用stateDiagram-v2 [*] -- 初始化: Task::spawn(async_fn) state Tokio 运行时 as Runtime { 初始化 -- 首次poll: schedule(task) state 第一次 poll as Poll1 { 首次poll -- 执行到await: Future::poll() 执行到await -- 注册Waker: Poll::Pending 注册Waker -- 挂起任务: 返回Pending } 挂起任务 -- 等待IO事件: task被加入等待队列 state IO就绪回调 as Wake { 等待IO事件 -- Waker被调用: epoll通知IO就绪 Waker被调用 -- 任务重新入队: wake_by_ref() } 任务重新入队 -- 再次poll: schedule(task) state 第N次 poll as PollN { 再次poll -- 从断点继续: Future::poll() 从断点继续 -- 最终完成: Poll::Ready(value) } } 最终完成 -- [*]: 返回结果这张图告诉你几个关键的时间节点第一次 poll任务被 spawn 后运行时会在一个合适的时机通常是当前线程空闲时调用Future::poll()。这时你的 async 函数开始执行直到遇到第一个还没就绪的.await。Waker 注册当 poll 返回Poll::Pending时Future 内部会通过cx.waker()获取一个 Waker 并保存起来。这个 Waker 是一个回调通知告诉运行时等 IO 就绪了记得来叫我。再次 poll当 IO 就绪比如网卡收到了数据epoll/kqueue 通知 Tokio 运行时运行时通过 Waker 把任务重新放回调度队列然后再次调用poll。这一次 poll 会从上次中断的地方继续执行。终态当 poll 返回Poll::Ready(value)任务完成不会再被调度。三、手写一个自定义 Future超时重试器光说不练假把式。我们来手写一个自定义 Future带超时重试的 HTTP 请求。如果第一次请求在 3 秒内没完成就自动重试一次。use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::time::Duration; use tokio::time::{sleep, timeout, Instant}; /// 带超时重试的 HTTP 请求 Future /// - 第一次请求3秒超时 /// - 超时后自动重试一次5秒超时 /// - 两次都超时就返回错误 enum RetryRequestFuture { /// 状态1正在做第一次请求3秒超时 FirstAttempt { // tokio::time::Timeout 也是一个 Future // 它内部包装了另一个 Future 一个计时器 inner: PinBoxdyn FutureOutput ResultString, tokio::time::error::Elapsed, }, /// 状态2第一次超时了正在做重试5秒超时 Retrying { inner: PinBoxdyn FutureOutput ResultString, tokio::time::error::Elapsed, }, /// 终态已完成不管成功还是失败 Done, } /// 模拟一个可能很慢的 HTTP 请求 /// 实际项目中会替换成 reqwest 的调用 async fn simulate_http_request(delay_ms: u64) - String { // 模拟网络延迟 sleep(Duration::from_millis(delay_ms)).await; format!(响应内容延迟{}ms, delay_ms) } impl Future for RetryRequestFuture { /// 输出类型Ok(成功内容) 或 Err(错误信息) type Output ResultString, String; fn poll(mut self: Pinmut Self, cx: mut Context_) - PollSelf::Output { loop { match mut *self { RetryRequestFuture::FirstAttempt { inner } { // poll 内部的超时 Future match inner.as_mut().poll(cx) { Poll::Ready(Ok(result)) { // 第一次就成功了直接返回 println!([重试器] 第一次请求成功); *self RetryRequestFuture::Done; return Poll::Ready(Ok(result)); } Poll::Ready(Err(_elapsed)) { // 第一次超时了 → 切换到重试状态 println!([重试器] 第一次超时准备重试...); let retry_future Box::pin(timeout( Duration::from_secs(5), // 重试时给更长的超时 simulate_http_request(2000), // 假设重试时延迟2秒 )); *self RetryRequestFuture::Retrying { inner: retry_future, }; continue; // 立即进入重试状态 } Poll::Pending { // 还没超时也没完成 → 继续等 return Poll::Pending; } } } RetryRequestFuture::Retrying { inner } { // poll 重试的超时 Future match inner.as_mut().poll(cx) { Poll::Ready(Ok(result)) { println!([重试器] 重试成功); *self RetryRequestFuture::Done; return Poll::Ready(Ok(result)); } Poll::Ready(Err(_)) { // 重试也超时了 → 彻底失败 *self RetryRequestFuture::Done; return Poll::Ready( Err(两次请求均超时本次放弃.to_string()) ); } Poll::Pending { return Poll::Pending; } } } RetryRequestFuture::Done { panic!(已完成的任务不应该被 poll); } } } } } /// 创建重试请求的构造函数 fn retry_request(delay_ms: u64) - RetryRequestFuture { let first_attempt Box::pin(timeout( Duration::from_secs(3), // 第一次给3秒 simulate_http_request(delay_ms), )); RetryRequestFuture::FirstAttempt { inner: first_attempt, } } #[tokio::main] async fn main() { // 测试模拟一个4秒的请求超过3秒超时触发重试 println!( 测试1请求耗时4秒会触发重试 ); match retry_request(4000).await { Ok(result) println!(成功: {}, result), Err(e) println!(失败: {}, e), } println!(\n 测试2请求耗时1秒不会触发重试 ); match retry_request(1000).await { Ok(result) println!(成功: {}, result), Err(e) println!(失败: {}, e), } }跑这个代码你会看到测试1的 4 秒请求超过 3 秒超时自动触发重试重试时模拟 2 秒延迟在 5 秒内完成重试成功测试2的 1 秒请求在 3 秒内完成直接返回四、poll 调用时机的几个反直觉细节初学者包括几个月前的我最容易产出的几个误解我这里统一澄清poll 被频繁调用会浪费 CPU不会。poll 返回Pending后如果不通过 Waker 通知调度器不会再次 poll。不存在空转轮询。Waker 只能唤醒一次不。每次 poll 被调用时都必须重新获取 Waker或者重用之前同一个。因为同一个 Future 可能被多次 poll。async 函数中的 for 循环也是状态是。编译器会把循环展开成状态机的一部分。如果你写for i in 0..10 { something.await; }编译后会变成 10 个状态或者用 loop 结构优化。自定义 Future 需要 unsafe不一定。如果你只是组合已有的 Future就像我上面的重试器不需要 unsafe。只有当你需要直接和操作系统 IO 交互、手动管理内存时才需要 unsafe。五、总结这篇文章我从编译器的视角把async/await生成的代码展开成状态机手画了 poll 调用时机的状态转移图最后手写了一个超时重试器来加深理解。理解 poll 的调用时机最核心的就是记住三点poll 不是循环调用的而是通过 Waker 事件驱动每次 poll 都从上次中断的地方继续不会重新执行Pending 后必须注册 Waker否则任务永远不会被唤醒作为自学者以前觉得这些东西编译器帮我做了不就行了吗现在越来越觉得不理解的抽象层一定会在某个深夜变成你的 bug。与其到时候被卡住不如现在花功夫搞明白。如果这篇文章帮你理清了 poll 的调用逻辑欢迎点赞收藏我们下篇见