Rust 实现 async/await的详细代码

Future

Future 代表一个可在未来某个时候获取返回值的 task，为了获取这个 task 的执行状况，Future 提供了一个函数用于判断该 task 是否执行返回。

trait Future {
  type Output;
  fn poll(self: Pin<&mut self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

poll 函数就是一个 Future 用于检查自己的 task 是否已经完成，例如我可以创建一个与某个 IP 建立 TCP 连接的 struct，在构建时完成建立连接的工作，然后实现 Future trait 时检查连接是否已经建立完成。根据建立情况返回 enum Poll 中的两个元素之一：

Poll::Pending: task 还在等待
Poll::Ready(result): task 携带 result 返回

实际上，基于 async 定义的函数和代码块也会被编译器编译为 Future。但是 async 函数或代码块无法显式地返回 Pending，因此一般只能完成一些简单的调用其他 Future 的工作。复杂的异步过程通常还是交由实现了 Future trait 的类型完成。

Wake & Context

你可能会好奇上面 poll 函数签名里的 cx 参数的作用，在 Rust 官方文档的定义中，Context 暂时只用于获取 Waker，而 Waker 的作用是用于提醒 executor 该 task 已经准备好运行了。

为什么需要 executor ？

同样以上面的建立 TCP 连接的例子来说，在网络卡顿时，进行一次 poll 可能都没有建立连接，如果没有设置 timeout 之类的东西的话，就需要进行多次 poll。这样的 Future 多了以后，我们可能会想，不妨将所有的 Future 都存储在一起，然后另起一个线程用于循环遍历所有的 Future 是否已经 ready，如果 ready 则返回结果。这就是一个非常简单的单线程 executor 的雏形。

也就是说，executor 是一个托管运行 task 的工具，类似于多线程，多线程要成功运行需要一个调度器进行调度。但是多线程至少需要语言层面甚至操作系统层面的支持，而 executor，如果你翻看 Rust 的官方文档的话，会发现没有任何关于 executor 的实现。实际上，Rust 选择将 executor 的实现交给第三方，自己只保留相关的交互接口（我在隔壁C++看了看，似乎也是一样的做法，并没有一个官方的 executor 实现，我唯一所知的在语言层面提供支持的只有Golang 的 goroutine）。

什么是 waker ？

上面讲述的轮询所有的 Future 是否已经完成实际是最低效的一种做法，当 Future 多了以后会带来相当多的 CPU 损耗。考虑到这点，Rust 还提供了一种机制可以用于通知 executor 某个 Future 是否应该被轮询，当然这只是其中的一种解决方式，实际上 Waker 的 wake 函数可以被实现为任何逻辑，取决于 executor。

在我看来，Waker 的内部定义相当不简洁，相当不 Rust。Waker 内部定义有一个 RawWaker，RawWaker 包含一个 RawWakerVTable，RawWakerVTable 定义了四个函数指针，executor 要实现 Waker 就需要定义这四种类型的函数然后赋值给 RawWakerVTable。

struct Waker {
  waker: RawWaker
}
struct RawWaker {
  data: *const (),
  vtable: &'static RawWakerVTable
}
struct RawWakerVTable {
  clone: unsafe fn(*const ()) -> RawWaker,
  wake: unsafe fn(*const ()),
  wake_by_ref: unsafe fn(*const ()),
  drop: unsafe fn(*const ())
}

之所以没有设计为 trait 形式，主要是 clone 函数，受限于 Rust 的 trait object safety，trait 中的任何函数的参数或返回值如果包含 Self 且有 type bound Sized，则不符合 trait object safe 规范，这样的 trait 可以被定义，可以被实现，但是无法与 dyn 一起进行动态绑定。

而 clones 函数又是必须的，因为 future 可能还会接着调用 future 的 poll 方法，就需要再 clone 一个 context 传入。

或许可以用 Box<dyn Waker> 或者 Arc<dyn Waker> 之类的，但是这些都不比 raw pointer 灵活，所以最终 Rust 还是选择定义一个包含函数指针的 struct。

async/await

这两个关键字可以说是异步编程领域的标志。，但在 Rust 中这两个关键字只是起到语法糖的作用，并不是异步的核心。

async 用于快速创建 Future，不管是函数还是代码块或者lambda表达式，都可以在前面加上 async 关键字快速变成 Future。对于

async fn bar() {
  foo().await;
}

编译器会自动生成类似下面的代码

fn bar() -> impl Future {
    std::future::from_generator(move |mut _task_context| {
        let _t = {
            match std::future::IntoFuture::into_future(foo()) {
                mut __awaitee => loop {
                    match unsafe {
                        std::future::Future::poll(
                            std::pin::Pin::new_unchecked(&mut __awaitee),
                            std::future::get_context(_task_context),
                        )
                    } {
                        std::task::Poll::Ready { 0: result } => break result,
                        std::task::Poll::Pending {} => {}
                    }
                    _task_context = (yield ());
                },
            };
        };
        _t
    })
}

Tips：上面的代码可以在 Rust Playground 里面点生成 HIR 看到。

Executor

前面讲到 wake 的时候，其实现与具体的 executor 相关，但是我觉得如果不从 executor 的实现角度看一下比较难以理解，只能浅显地知道 wake 是告诉 executor 准备再 poll 一遍。

Rust 中我知道的 async runtime lib 就是 futures-rs 和 tokio，前者在 GitHub 上是 rust-lang 官方组织推出的 repo，而后者虽然不清楚是否有官方参与，但是功能明显比前者丰富，据我所知使用异步的项目大部分都是使用 tokio。

我这里选择更简单的 futures-rs 讲一下其 executor 的实现，虽然其更加轻量但起码也是官方推出的，有质量保证。

Waker struct 到 ArcWake trait

futures-rs 还是将标准库里面的 Waker 封装成了 ArcWake trait，并且是 pub 的。和 raw pointer 打交道毕竟是 unsafe 的，与其满篇的 unsafe 乱飞，不如将 unsafe 限制在一定的范围内。

Waker 本质上是一个变量的指针(data)带着四个函数指针的结构体(RawWakerVTable)，因此在定义函数指针时只需要将指针强转成实现某个 trait 的泛型，再调用该 trait 的对应方法不就可以了。以 wake 函数为例：

trait Wake {
  fn wake(self) {
    Wake::wake_by_ref(&self);
  }
  fn wake_by_ref(&self);
}
unsafe fn wake<T: WakeTrait>(data: *const ()) {//对应RawWakerVTable里的函数指针
  let v = data.cast::<T>();
  v.wake();
}

这样就实现了 Waker struct 到 Waker trait 的转换。尽管如此，我们还需要一个结构体用来表示 Waker，满足下列条件：

实现 Deref trait，在引用时返回 &std::task::Waker
为了满足 Rust 的 safety rules，需要手动管理data的内存，显然某个实现了 Wake 的类型不会为了创建 waker 就交出自己的拥有权，因此只能通过传入的引用转成指针来创建 ManuallyDrop 实例，并考虑到 Deref trait 和后续的 Context 创建，需要通过 PhantomData 来管理 lifetime annotation

从而创建 WakeRef 结构体：

use std::mem::ManuallyDrop;
use std::task::Waker;
use std::marker::PhantomData;
struct WakeRef<'a> {
  waker: ManuallyDrop<Waker>,
  _marker: PhantomData<&'a ()>
}

如何根据引用创建 WakeRef 实例：

use std::task::{Waker, RawWaker};
fn get_waker<W: Wake>(wake: &W) -> WakeRef<'_> {
  let ptr = wake as *const _ as *const ();
  WakeRef {
    waker: ManuallyDrop::new(unsafe {Waker::from_raw(RawWaker::new(ptr, ...))}),//...省略的是创建RawWakerVTable的过程
    _marker: PhantomData
  }
}

实现 Deref

use std::task::Waker;
impl std::ops::Deref for WakeRef<'_> {
  type Target = Waker;
  fn deref(&self) -> &Waker {
    &self.waker
  }
}

因此对于某个实现 Wake 的类型来说，只需要传入引用就可以用 Context::from_waker(&waker) 来创建 context 了。

在 futures-rs 中，由于涉及到多线程，所以上述的其实并不安全，需要将普通引用改成 Arc 用于在多线程之间传递，Wake trait 也变成了 ArcWake，

trait ArcWake: Send + Sync {
  fn wake(self: Arc<Self>) {
    Self::wake_by_ref(&self)
  }
  
  fn wake_by_ref(arc_self: &Arc<Self>);
}

但是道理差不多。RawWakerVTable 的四个函数也与这个有关，以 wake 函数为例：

unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
  let arc: Arc<T> = Arc::from_raw(data.cast::<T>());
  ArcWake::wake(arc);
}

FuturesUnordered

FuturesUnordered 是一个 Future 的托管容器，其有一条链表维护所有的 Future，再通过一个队列维护所有需要运行的 Future（当然这里都不是 collections 里面那种普通的链表和队列，由于 FuturesUnordered 其实要与单线程和线程池 executor 共用，所以这两个数据结构其实还涉及很多原子化操作，在保证原子化且无锁的前提下要设计一个链表还挺麻烦的）。

struct FuturesUnordered<Fut> {
  ready_to_run_queue: Arc<ReadyToRunQueue<Fut>>,//需要运行的Future队列
  head_all: AtomicPtr<Task<Fut>>,//所有Future组成的链表
  is_terminated: AtomicBool
}

这里重点看 FuturesUnordered 如何实现 Waker，FuturesUnordered 将 Future 看作一个个 Task 。

struct Task<Fut> {
  future: UnsafeCell<Option<Fut>>,
  next_all: AtomicPtr<Task<Fut>>,//下一个Task节点
  len_all: UnsafeCell<usize>,//链表长度
  next_ready_to_run: AtomicPtr<Task<Fut>>,//下一个要运行的Task
  ready_to_run_queue: Weak<ReadyToRunQueue<Fut>>,
  queued: AtomicBool,//是否在Task链表内(Task运行时需要从链表上摘下)
  woken: AtomicBool//是否已经调用wake函数
}

为 Task 实现 ArcWake

impl<Fut> ArcWake for Task<Fut> {
  fn wake_by_ref(arc_self: &Arc<Self>) {
    let inner = match arc_self.ready_to_run_queue.upgrade() {
      Some(inner) => inner,
      None => return,
    };
    
    arc_self.woken.store(true, Relaxed);
    let prev = arc_self.queued.swap(true, SeqCst);
    if !prev {
      inner.enqueue(Arc::as_ptr(arc_self));
      inner.waker.wake();
    }
  }
}

当一个 Task 运行(被poll)时，其被从 FuturesUnordered 的 ready_to_run_queue 上摘下来，而在 wake 中又会重新放回去。因此，如果 Future 内部调用了 wake，则 Task 会再被放到 ready_to_run_queue 上运行，如果没有则不会。

所以每个 Future 使用的 context 其实是来自于 Task：

let waker = Task::waker_ref(task);
let mut cx = Context::from_waker(&waker);
future.poll(&mut cx);

FuturesUnordered 本身实现了 Stream trait

trait Stream {
  type Item;
  fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
}

FuturesUnordered 轮流 poll ready_to_run_queue 里面的 Future，根据返回结果返回：

Poll::Pending: ready_to_run_queue 为空或所有 Future 已经 poll 了一遍
Poll::Ready(Some(res)): 某个 Future 返回 Ready(res)
Poll::Ready(None): Task 链表为空，所有 Task 都已经结束返回

值得注意的是，在第一种情况下，所有的 Future 都 poll 了一遍，FuturesUnordered 会调用一次 wake，告诉 executor FuturesUnordered 已经运行了一个轮回，wake 具体的实现则取决于 executor。

单线程 executor

单线程 executor 允许在单线程上复用任意数量的 task，官方建议尽量在多I/O、只需要在 I/O 操作之间完成很少的工作的场景下使用。

struct LocalPool {
  pool: FuturesUnordered<LocalFutureObj<'static, ()>>,
  incoming: Rc<Incoming>
}

单线程 executor 将 Waker 的 wake 与线程的 wake 绑定，当调用 wake 时，如果 executor 线程处于 park(即阻塞) 状态，则 unpark 线程。

struct ThreadNotify {
  thread: std::thread::Thread,
  unparked: AtomicBool
}
impl ArcWake for ThreadNotify {
  fn wake_by_ref(arc_self: &Arc<Self>) {
    let unparked = arc_self.unparked.swap(true, Ordering::Release);
    if !unparked {
      arc_self.thread.unpark();
    }
  }
}

先看 LocalPool 如何定义 run 操作：

fn run_executor<T, F>(mut f: F) -> T
where
	F: FnMut(&mut Context<'_>) -> Poll<T>
{
  CURRENT_THREAD_NOTIFY.with(|thread_notify| {
    let waker = waker_ref(thread_notify);
    let mut cx = Context::from_waker(&waker);
    loop {
      if let Poll::Ready(t) = f(&mut cx) {//f决定了executor的运行方式，只要返回Ready就表明executor结束运行。
        return t;
      }
      while !thread_notify.unparked.swap(false, Ordering::Acquire) {
        thread::park();
      }
    }
  })
}

从 FutureUnordered 的角度来看，在 poll 一遍之后，如果需要继续运行，则调用 wake，将 unparked token 置为 true，此时线程不会陷入阻塞；否则 executor 线程会主动陷入阻塞。由于 FutureUnordered 和 executor 实际处于同一线程，因此此时 executor 只能从其他线程 unpark。

这种设计节省了 CPU 资源，使得线程只在有 Future 需要 poll 时需要运行，没有则挂起，再有了就又可以继续运行。

线程池 executor

线程池显然要比单线程 executor 更加复杂，随便一想就想到其至少要实现以下几点：

新 spawn 一个 Future，如何分配到某个线程
类似于单线程，在线程没有被调用 wake 时主动阻塞

对于第一点，使用多生产者单消费者管道 mpsc 进行 Future 的分发，实际的模型其实应该是多消费者单生产者，但是 Rust 并不提供这种管道，所以这里使用管道配合 mutex 使用。

struct PoolState {
  tx: Mutex<mpsc::Sender<Message>>,
  rx: Mutex<mpsc::Receiver<Message>>,
  cnt: AtomicUsize,//clone size
  size: usize//pool size
}

将 PoolState 包在 Arc 下就变成了 ThreadPool

struct ThreadPool {
  state: Arc<PoolState>
}

当 executor spawn 一个新的 future 时，只需要将其封装为一个 Task，然后传入管道：

fn spwan_obj_ok(&self, future: FutureObj<'static, ()>) {
  let task = Task {
    future,
    wake_handle: Arc::new(WakeHandl {exec: self.clone(), mutex: UnparkMutex::new()}),
    exec: self.clone()
  };
  self.state.send(Message::Run(task));
}

ThreadPool 也有自定义的 Task：

struct Task {
  future: FutureObj<'static ()>,
  exec: ThreadPool,
  wake_handle: Arc<WakeHandle>
}
struct WakeHandle {
  mutex: UnparkMutex<Task>,
  exec: ThreadPool
}

Task 主要分为以下状态：

POLLING: 正在poll
REPOLL: 正在 poll 的 Task 如果调用 wake 会变成 REPOLL 状态
WAITING： Task 正在等待
COMPLETE：Task 已经完成

为 Task 在不同状态间的转换，有些转换是自动的，比如 poll 返回 Ready 时自动进入 COMPLETE 状态，在 REPOLL 状态会通过调用 wait 函数再次进入 POLLING 状态重复运行一次 poll 函数；有些转换则需要调用函数，比如从 WAITING 进入 POLLING 需要调用 Task 的 run 函数才能运行。poll 返回 Pending 时根据 Future 是否调用 wake 函数分别进入 REPOLL 和 WAITING 状态。

impl Task {
  fn run(self) {
    let Self { mut future, wake_handle, mut exec } = self;
    let waker = waker_ref(&wake_handle);
    let mut cx = Context::from_waker(&waker);
    unsafe {
      wake_handle.mutex.start_poll();
      loop {
        let res = future.poll_unpin(&mut cx);
        match res {
          Poll::Pending => {}
          Poll::Ready(()) => return wake_handle.mutex.complete(),
        }
        let task = Self { future, wake_handle: wake_handle.clone(), exec };
        match wake_handle.mutex.wait(task) {
          Ok(()) => return, // we've waited
          Err(task) => {
            // someone's notified us
            future = task.future;
            exec = task.exec;
          }
        }
      }
    }
  }
}

线程池 executor 和单线程 executor 对待 Pending 的方式，相同点在于如果 Future 没有调用 wake，则放弃 Future，Future 要运行只能重新 spawn。不同点：

线程池：如果 Future 调用 wake，所在的线程阻塞式调用 poll 直到返回 Ready 或者 Future 放弃调用 wake
单线程：调用 wake 不会立刻再屌用 poll，但加入到 ready_to_run_queue 里面在下一次循环中被 poll

总结

本文只是一篇介绍 Rust 异步编程的原理，并通过具体的仓库稍微深挖一下实现的过程。具体的原因还是官方文档的介绍非常模糊，以我来说，第一次看到 Waker 完全不知道怎么用，底层到底是干了什么，"Future be ready to run again" 又是什么意思。如果不稍微看一下 runtime lib 的源码，有些东西很难理解。

本文只是简单介绍了一个 futures-rs 的实现，executor 方面都忽略了很多细节。而 futures-rs 还有大量的扩展代码藏在 util 目录下，但是这些东西一般看看文档就知道大概做了什么，懂得异步的实现原理就知道大概是怎么实现的，如果实在不懂还是可以去看源码。