Rust 如何實現 async/await

• Future
• Wake & Context
  • 為什麼需要 executor ？
  • 什麼是 waker ？
• async/await
• Executor
  • Waker struct 到 ArcWake trait
  • FuturesUnordered
  • 單線程 executor
  • 線程池 executor
• 總結

非同步編程在 Rust 中的地位非常高，很多 crate 尤其是多IO操作的都使用了 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 目錄下，但是這些東西一般看看文檔就知道大概做了什麼，懂得非同步的實現原理就知道大概是怎麼實現的，如果實在不懂還是可以去看源碼。