【譯】Async/Await(四)—— Pinning

Praying發表於2021-02-01

原文標題:Async/Await
原文連結:https://os.phil-opp.com/async-await/#multitasking
公眾號: Rust 碎碎念
翻譯 by: Praying

Pinning

在本文中我們已經與pinning偶遇多次。現在終於可以來討論pinning是什麼以及為什麼需要它?

自引用結構體(Self-Referential Structs)

正如上面所解釋的,狀態機的變換把每個暫停點的區域性變數儲存在一個結構體中。對於像example函式這樣的小例子,這會很直觀且不會導致什麼問題。但是,當變數開始互相引用時,事情就變得困難了。例如,考慮下面的函式:

async fn pin_example() -> i32 {
    let array = [123];
    let element = &array[2];
    async_write_file("foo.txt", element.to_string()).await;
    *element
}

這個函式建立了一個array,其中包含有123。它接著建立了一個對 array 最後一個元素的引用然後把它存入element變數。接下來,它把這個已經轉換為字串的數字非同步地寫入到檔案foo.txt中。最後,它返回了被element引用的數字。

因為這個函式使用了一個的await操作,所以得到的狀態機有三種狀態:啟動(start)、結束(end)和等待寫入(waiting on write)。這個函式沒有傳入引數,所以開始狀態的結構體是空的。和之前一樣,結束狀態也是空的因為函式在這個位置已經結束了。等待寫入狀態的結構體就比較有意思了:

struct WaitingOnWriteState {
    array: [123],
    element: 0x1001c// address of the last array element
}

我們需要把arrayelement變數都儲存起來,因為element在返回值的時候需要,而arrayelement所引用。element是一個引用,它儲存了一個指向被引用元素的指標(也就是一個記憶體地址)。這裡我們假設地址是0x1001c,在實際中,它需要是array的最後一個元素的地址,因此,它取決於結構體在記憶體中所處的位置。帶有這樣的內部指標的結構體被稱為自引用(self-referential)結構體,因為它們通過自己的一個欄位引用了它們自身。

自引用結構體的問題

我們的自引用結構體的內部指標導致了一個基本問題,當我們看到它的記憶體佈局後,這個問題就會變得明顯:

array欄位的起始地址為0x10014element欄位在地址0x10020。它指向了地址0x1001c,因為 array 的最後一個元素的位置就在這裡。此時,一切都沒有問題。但是,當我們試圖把這個結構體移動到一個不同的記憶體地址時,問題就出現了:

我們把結構體往後移動了一下,因此現在它的起始地址為0x10024。當我們把結構體作為函式引數傳遞時或者把它賦值給另一個棧上的變數,就會發生這種情況。問題在於,element欄位仍然指向地址0x1001c,而array的最後一個元素的地址已經變成0x1002c。因此,這個指標是懸垂(dangling)的,並會導致下一次呼叫poll時發生未定義行為。

可能的解決方案

解決這個懸垂指標問題有三種基本方式:

  • 在移動時更新指標:思路是無論什麼時候,只要結構體在記憶體中被移動,就更新內部的指標,因此這個指標在移動後仍然是有效的。不幸的是,這種方式將會需要 Rust 作出很大的改變並且有可能導致巨大的效能開銷。原因是,執行時需要追蹤所有結構體欄位的型別並且在每次移動操作時都要檢查是否需要更新指標。

  • 儲存一個偏移量來取代自引用:為了避免更新指標的需要,編譯器可以把自引用儲存為個結構體開始位置的偏移量。例如,上面的WaitingOnWriteState結構體中的element欄位可以儲存為值為 8 的element_offset欄位。因為,引用指向的 array 裡的元素起始於結構體開頭的 8 位元組。因為偏移位置在結構體移動時是不變的,所以不需要進行欄位更新。

    這種方式的問題在於它需要編譯器去探查所有的自引用。這在編譯時是不可能實現的,因為一個引用的值可能取決於使用者輸入,因此,我們可能再次需要一個執行時系統來分析引用並正確地建立狀態結構體。這不會導致執行時開銷,但是也阻礙了特定的編譯器優化,因此,它可能會再度引起巨大的效能開銷。

  • 禁止移動結構體:正如我們上面所見,懸垂指標僅發生於我們在記憶體中移動結構體時,通過完全禁止在自引用結構體上的移動操作,可以避免這個問題。這種方式的一個顯著優勢在於,它可以在型別系統層面上被實現而不需要額外的執行時開銷。缺點在於,它把處理可能是自引用結構的移動操作的負擔交給了程式設計師。

因為要保證提供零成本抽象(zero cost abstraction)的原則,這意味著抽象不應該引入額外的執行時開銷,所以 Rust 選擇了第三種方案。也因此,pinningAPI 在RFC2349中被提出。接下來,我們將會對這個 API 進行簡要介紹,並解釋它是如何與 async/await 以及 future 一同工作的。

堆上的值(Heap Values)

第一個發現是,在大多數情況下,堆分配(heap allocated)的值已經在記憶體中有了一個固定地址。它們通過呼叫allocate來建立,然後被一個指標型別引用,比如Box<T>。儘管指標型別有可能被移動,但是指標指向的堆上的值仍然保持在相同的記憶體地址,除非它被一個deallocate呼叫來釋放。

使用堆分配,我們可以嘗試去建立一個自引用結構體:

fn main() {
    let mut heap_value = Box::new(SelfReferential {
        self_ptr: 0 as *const _,
    });
    let ptr = &*heap_value as *const SelfReferential;
    heap_value.self_ptr = ptr;
    println!("heap value at: {:p}", heap_value);
    println!("internal reference: {:p}", heap_value.self_ptr);
}

struct SelfReferential {
    self_ptr: *const Self,
}

在 playground 上執行程式碼

我們建立了一個名為SelfReferential的簡單結構體,該結構體僅包含一個單獨的指標欄位。首先,我們使用一個空指標來初始化這個結構體,然後使用Box::new在堆上分配它。接著,我們計算出這個分配在堆上的結構體的記憶體地址並將其儲存到一個ptr變數中。最後,我們通過把ptr變數賦值給self_ptr欄位使得結構體成為自引用的。

當我們在 playground 上執行這段程式碼時,我們看到這個堆上的值的地址和它的內部指標的地址是相等的,這意味著,self_ptr欄位是一個有效的自引用。因為heap_value只是一個指標,移動它(比如,把它作為引數傳入函式)不會改變結構體自身的值,所以self_ptr在指標移動後依然是有效的。

但是,仍然有一種方式來破壞這個示例:我們可以擺脫Box<T>或者替換它的內容:

let stack_value = mem::replace(&mut *heap_value, SelfReferential {
    self_ptr: 0 as *const _,
});
println!("value at: {:p}", &stack_value);
println!("internal reference: {:p}", stack_value.self_ptr);

在 playground 上執行

這裡,我們使用mem::replace函式使用一個新的結構體例項來替換堆分配的值。這使得我們把原始的heap_value移動到棧上,而結構體的self_ptr欄位現在是一個仍然指向舊的堆地址的懸垂指標。當你嘗試在 playground 上執行這個示例時,你會看到列印出的"value at:""internal reference:"這一行確實是輸出的不同的指標。因此,在堆上分配一個值並不能保證自引用的安全。

出現上面的破綻的基本問題是,Box<T>允許我們獲得堆分配值的&mut T引用。這個&mut引用讓使用類似mem::replace或者mem::swap的方法使得堆上值失效成為可能。為了解決這個問題,我們必須阻止建立對自引用結構體的&mut引用。

Pin<Box>和 Unpin

pinning API 以Pin包裝型別和Unpin標記 trait 的形式提供了一個針對&mut T問題的解決方案。這些型別背後的思想是對Pin的所有能被用來獲得對 Unpin trait 上包裝的值的&mut引用的方法(如get_mut或者deref_mut)進行管控。Unpin trait 是一個auto trait,它會為所有的型別自動實現,除了顯式選擇退出(opt-out)的型別。通過讓自引用結構體選擇退出Unpin,就沒有(安全的)辦法從一個Pin<Box<T>>型別獲取一個&mut T。因此,它們的內部的自引用就能保證仍是有效的。

舉個例子,讓我們修改上面的SelfReferential型別來選擇退出Unpin

use core::marker::PhantomPinned;

struct SelfReferential {
    self_ptr: *const Self,
    _pin: PhantomPinned,
}

我們通過新增一個型別為PhantomPinned_pin欄位來選擇退出。這個型別是一個零大小標記型別,它唯一目的就是不去實現Unpin trait。因為 auto trait 的工作方式,有一個欄位不滿足Unpin,那麼整個結構體都會選擇退出Unpin

第二步是把例子中的Box<SelfReferential>改為Pin<Box<SelfReferential>>型別。實現這個的最簡單的方式是使用Box::pin函式,而不是使用Box::new建立堆分配的值。

let mut heap_value = Box::pin(SelfReferential {
    self_ptr: 0 as *const _,
    _pin: PhantomPinned,
});

除了把Box::new改為Box::pin之外,我們還需要在結構體初始化新增新的_pin欄位。因為PhantomPinned是一個零大小型別,我們只需要它的型別名來初始化它。

當我們嘗試執行調整後的示例時,我們看到它無法編譯:

error[E0594]: cannot assign to data in a dereference of `std::pin::Pin<std::boxed::Box<SelfReferential>>`
  --> src/main.rs:10:5
   |
10 |     heap_value.self_ptr = ptr;
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^ cannot assign
   |
   = help: trait `DerefMut` is required to modify through a dereference, but it is not implemented for `std::pin::Pin<std::boxed::Box<SelfReferential>>`

error[E0596]: cannot borrow data in a dereference of `std::pin::Pin<std::boxed::Box<SelfReferential>>` as mutable
  --> src/main.rs:16:36
   |
16 |     let stack_value = mem::replace(&mut *heap_value, SelfReferential {
   |                                    ^^^^^^^^^^^^^^^^ cannot borrow as mutable
   |
   = help: trait `DerefMut` is required to modify through a dereference, but it is not implemented for `std::pin::Pin<std::boxed::Box<SelfReferential>>`

兩個錯誤發生都是因為Pin<Box<SelfReferential>>型別沒有實現DerefMut trait。這也正是我們想要的,因為DerefMut trait 將會返回一個&mut引用,這是我們想要避免的。發生這種情況是因為我們選擇退出了Unpin並把Box::new改為了Box::pin

現在的問題在於,編譯器不僅阻止了第 16 行的移動型別,還禁止了第 10 行的self_ptr的初始化。這會發生時因為編譯器無法區分&mut引用的有效使用和無效使用。為了能夠正常初始化,我們不得不使用不安全的get_unchecked_mut方法:

// safe because modifying a field doesn't move the whole struct
unsafe {
    let mut_ref = Pin::as_mut(&mut heap_value);
    Pin::get_unchecked_mut(mut_ref).self_ptr = ptr;
}

嘗試在 playground 上執行

get_unchecked_mut函式作用於Pin<&mut T>而不是Pin<Box<T>>,所以我們不得不使用Pin::as_mut來對之前的值進行轉換。接著,我們可以使用get_unchecked_mut返回的&mut引用來設定self_ptr欄位。

現在,生下來的唯一的錯誤是mem::replace上的期望錯誤。記住,這個操作試圖把一個堆分配的值移動到棧上,這將會破壞儲存在self_ptr欄位上的自引用。通過選擇退出Unpin和使用Pin<Box<T>>,我們可以在編譯期阻止這個操作,從而安全地使用自引用結構體。正如我們所見,編譯器無法證明自引用的建立是安全的,因此我們需要使用一個不安全的塊(block)並且確認其自身的正確性。

棧 Pinning 和 Pin<&mut T>

在先前的部分,我們學習瞭如何使用Pin<Box<T>>來安全地建立一個堆分配的自引用的值。儘管這種方式能夠很好地工作並且相對安全(除了不安全的構造),但是需要的堆分配也會帶來效能損耗。因為 Rust 一直想要儘可能地提供零成本抽象, 所以 pinning API 也允許去建立Pin<&mut T>例項指向棧分配的值。

不像Pin<Box<T>> 例項那樣能夠擁有被包裝的值的所有權,Pin<&mut T>例項只是暫時地借用被包裝的值。這使得事情變得更加複雜,因為它要求程式設計師自己確認額外的保證。最重要的是,一個Pin<&mut T> 必須在被引用的T的整個生命週期被保持 pinned,這對於棧上的變數很難確認。為了幫助處理這類問題,就有了像pin-utils這樣的 crate。但是我仍然不會推薦 pinning 到棧上除非你真的知道自己在做什麼。

想要更加深入地瞭解,請查閱pin 模組Pin::new_unchecked方法的文件。

Pinning 和 Futures

正如我們在本文中已經看到的,Future::poll方法以Pin<&mut Self>引數的形式來使用 pinning:

fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>

這個方法接收self: Pin<&mut Self>而不是普通的&mut self,其原因在於,從 async/await 建立的 future 例項常常是自引用的。通過把Self包裝進Pin並讓編譯器為由 async/await 生的自引用的 futures 選擇退出Unpin,可以保證這些 futures 在poll呼叫之間在記憶體中不被移動。這就保證了所有的內部引用都是仍然有效的。

值得注意的是,在第一次poll呼叫之前移動 future 是沒問題的。因為事實上 future 是懶惰的(lazy)並且直到它們被第一次輪詢之前什麼事情也不會做。生成的狀態機中的start狀態因此只包含函式引數,而沒有內部引用。為了呼叫poll,呼叫者必須首先把 future 包裝進Pin,這就保證了 future 在記憶體中不會再被移動。因為棧上的 pinning 難以正確操作,所以我推薦一直使用Box::pin組合Pin::as_mut

如果你想了解如何安全地使用棧 pinning 實現一個 future 組合字函式,可以去看一下map 組合子方法的原始碼,以及 pin 文件中的 projections and structural pinning部分

相關文章