Rust Async 异步编程（三）：使用 Waker 来唤醒任务

Rust Async 异步编程（三）：使用 Waker 来唤醒任务
- 构建一个定时器
- 参考

Rust Async 异步编程（三）：使用 Waker 来唤醒任务

异步编程背后到底藏有什么秘密？究竟是哪只幕后之手在操纵这一切？如果你对这些感兴趣，就继续看下去，否则可以直接跳过，因为本章节的内容对于一个 API 工程师并没有太多帮助。

但是如果你希望能深入理解 Rust 的 async/.await 代码是如何工作、理解运行时和性能，甚至未来想要构建自己的 async 运行时或相关工具，那么本章节终究不会辜负于你。

构建一个定时器

对于 Future 来说，第一次被 poll 时无法完成任务是很正常的。但它需要确保在未来一旦准备好时，可以通知执行器再次对其进行 poll 进而继续往下执行，该通知就是通过 Waker 类型完成的。

Waker 提供了一个 wake() 方法可以用于告诉执行器：相关的任务可以被唤醒了，此时执行器就可以对相应的 Future 再次进行 poll 操作。

下面一起来实现一个简单的定时器 Future。为了让例子尽量简单，当计时器创建时，我们会启动一个线程接着让该线程进入睡眠，等睡眠结束后再通知给 Future。

注意本例子还会在后面继续使用，因此我们重新创建一个工程来演示：使用 cargo new --lib timer_future 来创建一个新工程，在 lib 包的根路径 src/lib.rs 中添加以下内容：

use std::{
    future::Future,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
    thread,
    time::Duration,
};

继续来实现 Future 定时器，之前提到：新建线程在睡眠结束后会需要将状态同步给定时器 Future，由于是多线程环境，我们需要使用 Arc<Mutex<T>> 来作为一个共享状态，用于在新线程和 Future 定时器间共享。

pub struct TimerFuture {
    shared_state: Arc<Mutex<SharedState>>,
}

/// 在 Future 和等待的线程间共享状态
struct SharedState {
    /// 定时（睡眠）是否结束
    completed: bool,

    /// 当睡眠结束后，线程可以用 waker 通知 TimerFuture 来唤醒任务
    waker: Option<Waker>,
}

下面给出 Future 的具体实现：

impl Future for TimerFuture {
    type Output = ();
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // 通过检查共享状态，来确定定时器是否已经完成
        let mut shared_state = self.shared_state.lock().unwrap();
        if shared_state.completed {
            Poll::Ready(())
        } else {
            // 设置 waker，这样新线程在睡眠（计时）结束后可以唤醒当前的任务，接着再次对 Future 进行 poll 操作,
            // 
            // 下面的 clone 每次被 pol 时都会发生一次，实际上，应该是只 clone 一次更加合理。
            // 选择每次都 clone 的原因是：TimerFuture 可以在执行器的不同任务间移动，如果只克隆一次，那么获取到的 waker 可能已经被篡改并指向了其它任务，最终导致执行器运行了错误的任务。
            shared_state.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}

代码很简单，只要新线程设置了 shared_state.completed = true，那任务就能顺利结束。如果没有设置，会为当前的任务克隆一份 Waker，这样新线程就可以使用它来唤醒当前的任务。

最后，再来创建一个 API 用于构建定时器和启动计时线程：

impl TimerFuture {
    /// 创建一个新的 TimerFuture，在指定的时间结束后，该 Future 可以完成
    pub fn new(duration: Duration) -> Self {
        let shared_state = Arc::new(Mutex::new(SharedState {
            completed: false,
            waker: None,
        }));

        // 创建新线程
        let thread_shared_state = shared_state.clone();
        thread::spawn(move || {
            // 睡眠指定时间实现计时功能
            thread::sleep(duration);
            let mut shared_state = thread_shared_state.lock().unwrap();
            // 通知执行器定时器已经完成，可以继续 poll 对应的 Future 了
            shared_state.completed = true;
            if let Some(waker) = shared_state.waker.take() {
                waker.wake()
            }
        });

        TimerFuture { shared_state }
    }
}