Rust?實(shí)現(xiàn)?async/await的詳細(xì)代碼
異步編程在 Rust 中的地位非常高,很多 crate 尤其是多IO操作的都使用了 async/await.
首先弄清楚異步編程的幾個(gè)基本概念:
Future
Future 代表一個(gè)可在未來某個(gè)時(shí)候獲取返回值的 task,為了獲取這個(gè) task 的執(zhí)行狀況,F(xiàn)uture 提供了一個(gè)函數(shù)用于判斷該 task 是否執(zhí)行返回。
trait Future { type Output; fn poll(self: Pin<&mut self>, cx: &mut Context<'_>) -> Poll<Self::Output>; }
poll 函數(shù)就是一個(gè) Future 用于檢查自己的 task 是否已經(jīng)完成,例如我可以創(chuàng)建一個(gè)與某個(gè) IP 建立 TCP 連接的 struct,在構(gòu)建時(shí)完成建立連接的工作,然后實(shí)現(xiàn) Future trait 時(shí)檢查連接是否已經(jīng)建立完成。根據(jù)建立情況返回 enum Poll 中的兩個(gè)元素之一:
- Poll::Pending: task 還在等待
- Poll::Ready(result): task 攜帶 result 返回
實(shí)際上,基于 async 定義的函數(shù)和代碼塊也會(huì)被編譯器編譯為 Future。但是 async 函數(shù)或代碼塊無法顯式地返回 Pending,因此一般只能完成一些簡(jiǎn)單的調(diào)用其他 Future 的工作。復(fù)雜的異步過程通常還是交由實(shí)現(xiàn)了 Future trait 的類型完成。
Wake & Context
你可能會(huì)好奇上面 poll 函數(shù)簽名里的 cx 參數(shù)的作用,在 Rust 官方文檔的定義中,Context 暫時(shí)只用于獲取 Waker,而 Waker 的作用是用于提醒 executor 該 task 已經(jīng)準(zhǔn)備好運(yùn)行了。
為什么需要 executor ?
同樣以上面的建立 TCP 連接的例子來說,在網(wǎng)絡(luò)卡頓時(shí),進(jìn)行一次 poll 可能都沒有建立連接,如果沒有設(shè)置 timeout 之類的東西的話,就需要進(jìn)行多次 poll。這樣的 Future 多了以后,我們可能會(huì)想,不妨將所有的 Future 都存儲(chǔ)在一起,然后另起一個(gè)線程用于循環(huán)遍歷所有的 Future 是否已經(jīng) ready,如果 ready 則返回結(jié)果。這就是一個(gè)非常簡(jiǎn)單的單線程 executor 的雛形。
也就是說,executor 是一個(gè)托管運(yùn)行 task 的工具,類似于多線程,多線程要成功運(yùn)行需要一個(gè)調(diào)度器進(jìn)行調(diào)度。但是多線程至少需要語言層面甚至操作系統(tǒng)層面的支持,而 executor,如果你翻看 Rust 的官方文檔的話,會(huì)發(fā)現(xiàn)沒有任何關(guān)于 executor 的實(shí)現(xiàn)。實(shí)際上,Rust 選擇將 executor 的實(shí)現(xiàn)交給第三方,自己只保留相關(guān)的交互接口(我在隔壁C++看了看,似乎也是一樣的做法,并沒有一個(gè)官方的 executor 實(shí)現(xiàn),我唯一所知的在語言層面提供支持的只有Golang 的 goroutine)。
什么是 waker ?
上面講述的輪詢所有的 Future 是否已經(jīng)完成實(shí)際是最低效的一種做法,當(dāng) Future 多了以后會(huì)帶來相當(dāng)多的 CPU 損耗。考慮到這點(diǎn),Rust 還提供了一種機(jī)制可以用于通知 executor 某個(gè) Future 是否應(yīng)該被輪詢,當(dāng)然這只是其中的一種解決方式,實(shí)際上 Waker 的 wake 函數(shù)可以被實(shí)現(xiàn)為任何邏輯,取決于 executor。
在我看來,Waker 的內(nèi)部定義相當(dāng)不簡(jiǎn)潔,相當(dāng)不 Rust。Waker 內(nèi)部定義有一個(gè) RawWaker,RawWaker 包含一個(gè) RawWakerVTable,RawWakerVTable 定義了四個(gè)函數(shù)指針,executor 要實(shí)現(xiàn) Waker 就需要定義這四種類型的函數(shù)然后賦值給 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 ()) }
之所以沒有設(shè)計(jì)為 trait 形式,主要是 clone 函數(shù),受限于 Rust 的 trait object safety,trait 中的任何函數(shù)的參數(shù)或返回值如果包含 Self 且有 type bound Sized,則不符合 trait object safe 規(guī)范,這樣的 trait 可以被定義,可以被實(shí)現(xiàn),但是無法與 dyn 一起進(jìn)行動(dòng)態(tài)綁定。
而 clones 函數(shù)又是必須的,因?yàn)?future 可能還會(huì)接著調(diào)用 future 的 poll 方法,就需要再 clone 一個(gè) context 傳入。
或許可以用 Box<dyn Waker>
或者 Arc<dyn Waker>
之類的,但是這些都不比 raw pointer 靈活,所以最終 Rust 還是選擇定義一個(gè)包含函數(shù)指針的 struct。
async/await
這兩個(gè)關(guān)鍵字可以說是異步編程領(lǐng)域的標(biāo)志。,但在 Rust 中這兩個(gè)關(guān)鍵字只是起到語法糖的作用,并不是異步的核心。
async 用于快速創(chuàng)建 Future,不管是函數(shù)還是代碼塊或者lambda表達(dá)式,都可以在前面加上 async 關(guān)鍵字快速變成 Future。對(duì)于
async fn bar() { foo().await; }
編譯器會(huì)自動(dòng)生成類似下面的代碼
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 里面點(diǎn)生成 HIR 看到。
Executor
前面講到 wake 的時(shí)候,其實(shí)現(xiàn)與具體的 executor 相關(guān),但是我覺得如果不從 executor 的實(shí)現(xiàn)角度看一下比較難以理解,只能淺顯地知道 wake 是告訴 executor 準(zhǔn)備再 poll 一遍。
Rust 中我知道的 async runtime lib 就是 futures-rs 和 tokio,前者在 GitHub 上是 rust-lang 官方組織推出的 repo,而后者雖然不清楚是否有官方參與,但是功能明顯比前者豐富,據(jù)我所知使用異步的項(xiàng)目大部分都是使用 tokio。
我這里選擇更簡(jiǎn)單的 futures-rs 講一下其 executor 的實(shí)現(xiàn),雖然其更加輕量但起碼也是官方推出的,有質(zhì)量保證。
Waker struct 到 ArcWake trait
futures-rs 還是將標(biāo)準(zhǔn)庫(kù)里面的 Waker 封裝成了 ArcWake trait,并且是 pub 的。和 raw pointer 打交道畢竟是 unsafe 的,與其滿篇的 unsafe 亂飛,不如將 unsafe 限制在一定的范圍內(nèi)。
Waker 本質(zhì)上是一個(gè)變量的指針(data)帶著四個(gè)函數(shù)指針的結(jié)構(gòu)體(RawWakerVTable),因此在定義函數(shù)指針時(shí)只需要將指針強(qiáng)轉(zhuǎn)成實(shí)現(xiàn)某個(gè) trait 的泛型,再調(diào)用該 trait 的對(duì)應(yīng)方法不就可以了。以 wake 函數(shù)為例:
trait Wake { fn wake(self) { Wake::wake_by_ref(&self); } fn wake_by_ref(&self); } unsafe fn wake<T: WakeTrait>(data: *const ()) {//對(duì)應(yīng)RawWakerVTable里的函數(shù)指針 let v = data.cast::<T>(); v.wake(); }
這樣就實(shí)現(xiàn)了 Waker struct 到 Waker trait 的轉(zhuǎn)換。盡管如此,我們還需要一個(gè)結(jié)構(gòu)體用來表示 Waker,滿足下列條件:
- 實(shí)現(xiàn) Deref trait,在引用時(shí)返回 &std::task::Waker
- 為了滿足 Rust 的 safety rules,需要手動(dòng)管理data的內(nèi)存,顯然某個(gè)實(shí)現(xiàn)了 Wake 的類型不會(huì)為了創(chuàng)建 waker 就交出自己的擁有權(quán),因此只能通過傳入的引用轉(zhuǎn)成指針來創(chuàng)建 ManuallyDrop 實(shí)例,并考慮到 Deref trait 和后續(xù)的 Context 創(chuàng)建,需要通過 PhantomData 來管理 lifetime annotation
從而創(chuàng)建 WakeRef 結(jié)構(gòu)體:
use std::mem::ManuallyDrop; use std::task::Waker; use std::marker::PhantomData; struct WakeRef<'a> { waker: ManuallyDrop<Waker>, _marker: PhantomData<&'a ()> }
如何根據(jù)引用創(chuàng)建 WakeRef 實(shí)例:
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, ...))}),//...省略的是創(chuàng)建RawWakerVTable的過程 _marker: PhantomData } }
實(shí)現(xiàn) Deref
use std::task::Waker; impl std::ops::Deref for WakeRef<'_> { type Target = Waker; fn deref(&self) -> &Waker { &self.waker } }
因此對(duì)于某個(gè)實(shí)現(xiàn) Wake 的類型來說,只需要傳入引用就可以用 Context::from_waker(&waker) 來創(chuàng)建 context 了。
在 futures-rs 中,由于涉及到多線程,所以上述的其實(shí)并不安全,需要將普通引用改成 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 的四個(gè)函數(shù)也與這個(gè)有關(guān),以 wake 函數(shù)為例:
unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) { let arc: Arc<T> = Arc::from_raw(data.cast::<T>()); ArcWake::wake(arc); }
FuturesUnordered
FuturesUnordered 是一個(gè) Future 的托管容器,其有一條鏈表維護(hù)所有的 Future,再通過一個(gè)隊(duì)列維護(hù)所有需要運(yùn)行的 Future(當(dāng)然這里都不是 collections 里面那種普通的鏈表和隊(duì)列,由于 FuturesUnordered 其實(shí)要與單線程和線程池 executor 共用,所以這兩個(gè)數(shù)據(jù)結(jié)構(gòu)其實(shí)還涉及很多原子化操作,在保證原子化且無鎖的前提下要設(shè)計(jì)一個(gè)鏈表還挺麻煩的)。
struct FuturesUnordered<Fut> { ready_to_run_queue: Arc<ReadyToRunQueue<Fut>>,//需要運(yùn)行的Future隊(duì)列 head_all: AtomicPtr<Task<Fut>>,//所有Future組成的鏈表 is_terminated: AtomicBool }
這里重點(diǎn)看 FuturesUnordered 如何實(shí)現(xiàn) Waker,F(xiàn)uturesUnordered 將 Future 看作一個(gè)個(gè) Task 。
struct Task<Fut> { future: UnsafeCell<Option<Fut>>, next_all: AtomicPtr<Task<Fut>>,//下一個(gè)Task節(jié)點(diǎn) len_all: UnsafeCell<usize>,//鏈表長(zhǎng)度 next_ready_to_run: AtomicPtr<Task<Fut>>,//下一個(gè)要運(yùn)行的Task ready_to_run_queue: Weak<ReadyToRunQueue<Fut>>, queued: AtomicBool,//是否在Task鏈表內(nèi)(Task運(yùn)行時(shí)需要從鏈表上摘下) woken: AtomicBool//是否已經(jīng)調(diào)用wake函數(shù) }
為 Task 實(shí)現(xiàn) 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(); } } }
當(dāng)一個(gè) Task 運(yùn)行(被poll)時(shí),其被從 FuturesUnordered 的 ready_to_run_queue 上摘下來,而在 wake 中又會(huì)重新放回去。因此,如果 Future 內(nèi)部調(diào)用了 wake,則 Task 會(huì)再被放到 ready_to_run_queue 上運(yùn)行,如果沒有則不會(huì)。
所以每個(gè) Future 使用的 context 其實(shí)是來自于 Task:
let waker = Task::waker_ref(task); let mut cx = Context::from_waker(&waker); future.poll(&mut cx);
FuturesUnordered 本身實(shí)現(xiàn)了 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,根據(jù)返回結(jié)果返回:
- Poll::Pending: ready_to_run_queue 為空或所有 Future 已經(jīng) poll 了一遍
- Poll::Ready(Some(res)): 某個(gè) Future 返回 Ready(res)
- Poll::Ready(None): Task 鏈表為空,所有 Task 都已經(jīng)結(jié)束返回
值得注意的是,在第一種情況下,所有的 Future 都 poll 了一遍,F(xiàn)uturesUnordered 會(huì)調(diào)用一次 wake,告訴 executor FuturesUnordered 已經(jīng)運(yùn)行了一個(gè)輪回,wake 具體的實(shí)現(xiàn)則取決于 executor。
單線程 executor
單線程 executor 允許在單線程上復(fù)用任意數(shù)量的 task,官方建議盡量在多I/O、只需要在 I/O 操作之間完成很少的工作的場(chǎng)景下使用。
struct LocalPool { pool: FuturesUnordered<LocalFutureObj<'static, ()>>, incoming: Rc<Incoming> }
單線程 executor 將 Waker 的 wake 與線程的 wake 綁定,當(dāng)調(diào)用 wake 時(shí),如果 executor 線程處于 park(即阻塞) 狀態(tài),則 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的運(yùn)行方式,只要返回Ready就表明executor結(jié)束運(yùn)行。 return t; } while !thread_notify.unparked.swap(false, Ordering::Acquire) { thread::park(); } } }) }
從 FutureUnordered 的角度來看,在 poll 一遍之后,如果需要繼續(xù)運(yùn)行,則調(diào)用 wake,將 unparked token 置為 true,此時(shí)線程不會(huì)陷入阻塞;否則 executor 線程會(huì)主動(dòng)陷入阻塞。由于 FutureUnordered 和 executor 實(shí)際處于同一線程,因此此時(shí) executor 只能從其他線程 unpark。
這種設(shè)計(jì)節(jié)省了 CPU 資源,使得線程只在有 Future 需要 poll 時(shí)需要運(yùn)行,沒有則掛起,再有了就又可以繼續(xù)運(yùn)行。
線程池 executor
線程池顯然要比單線程 executor 更加復(fù)雜,隨便一想就想到其至少要實(shí)現(xiàn)以下幾點(diǎn):
- 新 spawn 一個(gè) Future,如何分配到某個(gè)線程
- 類似于單線程,在線程沒有被調(diào)用 wake 時(shí)主動(dòng)阻塞
對(duì)于第一點(diǎn),使用多生產(chǎn)者單消費(fèi)者管道 mpsc 進(jìn)行 Future 的分發(fā),實(shí)際的模型其實(shí)應(yīng)該是多消費(fèi)者單生產(chǎn)者,但是 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> }
當(dāng) executor spawn 一個(gè)新的 future 時(shí),只需要將其封裝為一個(gè) 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 主要分為以下狀態(tài):
- POLLING: 正在poll
- REPOLL: 正在 poll 的 Task 如果調(diào)用 wake 會(huì)變成 REPOLL 狀態(tài)
- WAITING: Task 正在等待
- COMPLETE:Task 已經(jīng)完成
為 Task 在不同狀態(tài)間的轉(zhuǎn)換,有些轉(zhuǎn)換是自動(dòng)的,比如 poll 返回 Ready 時(shí)自動(dòng)進(jìn)入 COMPLETE 狀態(tài),在 REPOLL 狀態(tài)會(huì)通過調(diào)用 wait 函數(shù)再次進(jìn)入 POLLING 狀態(tài)重復(fù)運(yùn)行一次 poll 函數(shù);有些轉(zhuǎn)換則需要調(diào)用函數(shù),比如從 WAITING 進(jìn)入 POLLING 需要調(diào)用 Task 的 run 函數(shù)才能運(yùn)行。poll 返回 Pending 時(shí)根據(jù) Future 是否調(diào)用 wake 函數(shù)分別進(jìn)入 REPOLL 和 WAITING 狀態(tài)。
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 對(duì)待 Pending 的方式,相同點(diǎn)在于如果 Future 沒有調(diào)用 wake,則放棄 Future,F(xiàn)uture 要運(yùn)行只能重新 spawn。不同點(diǎn):
- 線程池:如果 Future 調(diào)用 wake,所在的線程阻塞式調(diào)用 poll 直到返回 Ready 或者 Future 放棄調(diào)用 wake
- 單線程:調(diào)用 wake 不會(huì)立刻再屌用 poll,但加入到 ready_to_run_queue 里面在下一次循環(huán)中被 poll
總結(jié)
本文只是一篇介紹 Rust 異步編程的原理,并通過具體的倉(cāng)庫(kù)稍微深挖一下實(shí)現(xiàn)的過程。具體的原因還是官方文檔的介紹非常模糊,以我來說,第一次看到 Waker 完全不知道怎么用,底層到底是干了什么,"Future be ready to run again" 又是什么意思。如果不稍微看一下 runtime lib 的源碼,有些東西很難理解。
本文只是簡(jiǎn)單介紹了一個(gè) futures-rs 的實(shí)現(xiàn),executor 方面都忽略了很多細(xì)節(jié)。而 futures-rs 還有大量的擴(kuò)展代碼藏在 util 目錄下,但是這些東西一般看看文檔就知道大概做了什么,懂得異步的實(shí)現(xiàn)原理就知道大概是怎么實(shí)現(xiàn)的,如果實(shí)在不懂還是可以去看源碼。
到此這篇關(guān)于Rust 如何實(shí)現(xiàn) async/await的文章就介紹到這了,更多相關(guān)Rust 實(shí)現(xiàn) async await內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家!
相關(guān)文章
Rust調(diào)用函數(shù)操作符?.?和?::?的區(qū)別詳解
在Rust中,.和::操作符都可以用來調(diào)用方法,但它們的用法有所不同,所以本文就將詳細(xì)的給大家介紹一下.和::操作符的區(qū)別,感興趣的同學(xué)跟著小編一起來學(xué)習(xí)吧2023-07-07如何在Rust中處理命令行參數(shù)和環(huán)境變量
在本章節(jié)中, 我們探討了Rust處理命令行參數(shù)的常見的兩種方式和處理環(huán)境變量的兩種常見方式,感興趣的朋友一起看看吧2023-12-12rust使用Atomic創(chuàng)建全局變量和使用操作方法
從 Rust1.34 版本后,就正式支持原子類型,原子指的是一系列不可被 CPU 上下文交換的機(jī)器指令,這些指令組合在一起就形成了原子操作,這篇文章主要介紹了rust使用Atomic創(chuàng)建全局變量和使用,需要的朋友可以參考下2024-05-05