Rust 執行階段檢查提升通道安全性

在 Rust 中使用通道傳遞資料時，確保其安全性至關重要。本文介紹的技術利用執行階段檢查，有效防止了訊息的誤用和未定義行為的發生。透過原子操作和記憶體屏障，我們可以精確控制訊息的讀寫順序，避免資料競爭。同時，新增的旗標和檢查機制則進一步提升了通道的安全性，防止多次接收和傳送訊息所帶來的問題。此外，Drop 特性的引入有效地解決了資源洩漏的風險，而單一原子變數的最佳化則提升了通道的效能。最後，Sender 和 Receiver 的型別安全設計，則從根本上杜絕了誤用的可能性。

透過執行階段檢查提升安全性

為了提供更安全的介面，我們可以加入檢查機制，使誤用導致的結果變成具有明確訊息的panic，這遠比未定義行為更為可取。

檢查接收順序

首先，我們來處理在訊息準備好之前呼叫receive的問題。這個問題很容易解決，只要在receive方法中驗證ready旗標即可：

/// 如果尚未有訊息可用，則會`panic`。
///
/// 提示：先使用 `is_ready` 檢查。
///
/// 安全性：只呼叫一次！
pub unsafe fn receive(&self) -> T {
    if !self.ready.load(Acquire) {
        panic!("沒有可用的訊息！");
    }
    (*self.message.get()).assume_init_read()
}

這個函式仍然是不安全的，因為使用者仍需負責不多次呼叫此函式。但如果未能先檢查is_ready()，現在將導致panic而非未定義行為。

由於現在receive方法內部有一個acquire-load操作來檢查ready旗標，提供了必要的同步機制，我們可以將is_ready中的load操作的記憶體順序降為Relaxed，因為它現在僅用於指示目的：

pub fn is_ready(&self) -> bool {
    self.ready.load(Relaxed)
}

需要注意的是，ready旗標的總修改順序保證了在is_ready載入true後，receive也會看到true。無論is_ready中使用什麼樣的記憶體順序，都不可能出現is_ready傳回true而receive()仍然panic的情況。

內容解密：

• is_ready方法檢查ready旗標：用於判斷是否有訊息可供接收。
• receive方法中的Acquire順序：確保在讀取訊息之前，所有相關的寫入操作已經完成。
• Relaxed順序在is_ready中的使用：由於receive方法提供了必要的同步，因此在is_ready中使用Relaxed是安全的。

處理多次接收的情況

下一個需要解決的問題是多次呼叫receive的情況。我們可以透過在receive方法中將ready旗標重設為false來使其導致panic：

/// 如果尚未有訊息可用，或者訊息已被消費，則會`panic`。
///
/// 提示：先使用 `is_ready` 檢查。
pub fn receive(&self) -> T {
    if !self.ready.swap(false, Acquire) {
        panic!("沒有可用的訊息！");
    }
    // 安全性：我們剛剛檢查（並重設）了`ready`旗標。
    unsafe { (*self.message.get()).assume_init_read() }
}

透過將load操作替換為swap(false, Acquire)，receive方法現在是完全安全的。使用者不再需要擔心呼叫此函式多次會導致未定義行為。

內容解密：

• swap(false, Acquire)操作：將ready旗標重設為false，並傳回其之前的值，確保訊息只被讀取一次。
• receive方法的安全性：透過檢查和重設ready旗標，保證了訊息的正確接收和消費。

處理多次傳送的情況

對於send方法，情況稍微複雜一些。為了防止多個send呼叫同時存取cell，我們需要知道是否已經有另一個send呼叫開始。ready旗標只能告訴我們是否有另一個send呼叫已經完成，因此這不夠。

我們新增一個名為in_use的第二個旗標，用於指示通道是否已被使用：

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    in_use: AtomicBool, // 新增！
    ready: AtomicBool,
}

impl<T> Channel<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            in_use: AtomicBool::new(false), // 新增！
            ready: AtomicBool::new(false),
        }
    }
    // ...
}

現在，我們只需要在send方法中，在存取cell之前將in_use設為true，並在已經被設定時panic：

/// 當嘗試傳送多於一個訊息時會`panic`。
pub fn send(&self, message: T) {
    if self.in_use.swap(true, Relaxed) {
        panic!("無法傳送多於一個訊息！");
    }
    unsafe { (*self.message.get()).write(message) };
    self.ready.store(true, Release);
}

對於in_use的原子交換操作，我們可以使用Relaxed記憶體順序，因為in_use的總修改順序保證了操作的正確性。

內容解密：

• in_use旗標的作用：用於防止多個send呼叫同時寫入cell，確保資料的完整性。
• Relaxed順序在in_use交換操作中的使用：由於in_use的總修改順序提供了必要的保證，因此使用Relaxed是安全的。

未來改進方向

1. 進一步最佳化同步機制：目前的實作已經透過執行階段檢查提升了安全性，但仍有最佳化的空間，例如探索使用更高效的同步原語。
2. 支援多訊息通道：目前的實作僅支援單一訊息的傳遞。未來可以考慮擴充套件到支援多訊息的佇列，以滿足更複雜的應用場景。
3. 錯誤處理機制的完善：雖然目前的實作透過panic來處理錯誤，但對於某些應用場景，傳回錯誤碼或使用Result型別可能更為合適。未來可以考慮加入更靈活的錯誤處理機制。

透過這些改進，我們的通道實作將變得更為強大和靈活，能夠滿足更多使用者的需求。

通道狀態轉換圖

@startuml
skinparam backgroundColor #FEFEFE
skinparam defaultTextAlignment center
skinparam rectangleBackgroundColor #F5F5F5
skinparam rectangleBorderColor #333333
skinparam arrowColor #333333

title 未來改進方向

rectangle "new" as node1
rectangle "send" as node2
rectangle "寫入訊息並設為ready" as node3
rectangle "receive" as node4
rectangle "讀取訊息並重設ready" as node5
rectangle "重複send" as node6
rectangle "重複receive" as node7

node1 --> node2
node2 --> node3
node3 --> node4
node4 --> node5
node5 --> node6
node6 --> node7

@enduml

圖表翻譯： 此圖示展示了通道的狀態轉換過程。通道初始狀態為「未初始化」，透過new方法初始化後進入「空閒」狀態。當呼叫send方法時，通道進入「傳送中」狀態，並在寫入訊息後進入「訊息就緒」狀態。如果在「傳送中」狀態再次呼叫send，則會觸發panic。當呼叫receive方法時，通道進入「接收中」狀態，並在讀取訊息後回到「空閒」狀態。如果在「訊息就緒」狀態再次呼叫receive，同樣會觸發panic。此圖清晰地展示了通道的生命週期和錯誤處理機制。

使用執行緒安全的通道實作與改進

在前面的章節中，我們實作了一個基本的通道（Channel），並且確保了其安全性與正確性。然而，在實際應用中，我們仍然需要考慮更多因素來進一步完善這個實作。

進一步的安全性檢查

為了確保通道的安全性，我們需要處理以下幾個問題：

1. 訊息遺失：當訊息傳送後未被接收，可能會導致資源洩漏。
2. 重複操作：避免對 send 或 receive 的重複呼叫。
3. 執行緒安全：確保通道在多執行緒環境下的正確性。

實作 Drop 特性確保資源釋放

為了避免資源洩漏，我們為 Channel 實作了 Drop 特性。這個實作確保了當通道被丟棄時，如果存在未被接收的訊息，將會正確地釋放資源。

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        if *self.ready.get_mut() {
            unsafe { self.message.get_mut().assume_init_drop() }
        }
    }
}

內容解密：

• 這個 Drop 實作首先檢查 ready 旗標，判斷是否有未接收的訊息。
• 如果存在未接收的訊息，將會呼叫 assume_init_drop 來釋放 message 中的資源。
• 這種實作方式確保了即使在異常情況下，資源也不會洩漏。

測試通道實作

為了驗證通道的正確性，我們編寫了一個測試程式。在這個測試中，我們建立了一個通道，並在另一個執行緒中傳送訊息，而主執行緒則等待並接收該訊息。

fn main() {
    let channel = Channel::new();
    let t = thread::current();
    thread::scope(|s| {
        s.spawn(|| {
            channel.send("hello world!");
            t.unpark();
        });
        while !channel.is_ready() {
            thread::park();
        }
        assert_eq!(channel.receive(), "hello world!");
    });
}

內容解密：

• 我們首先建立了一個 Channel 例項，並取得了當前執行緒的控制權。
• 在新的執行緒中，我們發送了字串 "hello world!" 並喚醒主執行緒。
• 主執行緒在等待訊息準備好後接收並驗證訊息內容。

使用單一原子變數最佳化通道狀態

為了最佳化記憶體使用，我們可以使用一個 AtomicU8 來表示通道的四種狀態，而不是使用兩個獨立的原子布林值。

const EMPTY: u8 = 0;
const WRITING: u8 = 1;
const READY: u8 = 2;
const READING: u8 = 3;

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    state: AtomicU8,
}

impl<T> Channel<T> {
    pub fn send(&self, message: T) {
        if self.state.compare_exchange(EMPTY, WRITING, Relaxed, Relaxed).is_err() {
            panic!("can't send more than one message!");
        }
        unsafe { (*self.message.get()).write(message) };
        self.state.store(READY, Release);
    }

    pub fn receive(&self) -> T {
        if self.state.compare_exchange(READY, READING, Acquire, Relaxed).is_err() {
            panic!("no message available!");
        }
        unsafe { (*self.message.get()).assume_init_read() }
    }
}

內容解密：

• 我們定義了四種狀態：EMPTY、WRITING、READY 和 READING，用於管理通道的狀態。
• 在 send 方法中，我們使用 compare_exchange 來確保訊息只被傳送一次。
• 在 receive 方法中，我們同樣使用 compare_exchange 來確保訊息只被接收一次。

透過型別系統保證安全性

為了進一步提高安全性，我們設計了一個新的介面，使用 Sender 和 Receiver 兩個獨立的型別來表示通道的兩端。

pub fn channel<T>() -> (Sender<T>, Receiver<T>) { ... }

pub struct Sender<T> { 
    channel: Arc<Channel<T>>,
}

pub struct Receiver<T> { 
    channel: Arc<Channel<T>>,
}

impl<T> Sender<T> {
    pub fn send(self, message: T) { ... }
}

impl<T> Receiver<T> {
    pub fn receive(self) -> T { ... }
}

內容解密：

• 我們使用 Arc 來分享 Channel 例項，使得 Sender 和 Receiver 可以在不同執行緒中安全地使用。
• 透過讓 send 和 receive 方法取得 self，我們確保了這些方法只能被呼叫一次。

未來，我們可以繼續最佳化通道的實作，例如增加更多的錯誤處理機制，或者支援更複雜的訊息傳遞模式。同時，我們也可以探索如何在不同的平行程式設計模型中使用我們的通道實作。

參考資料

• Rust 程式語言官方檔案
• 相關的平行程式設計文獻

附錄

通道狀態轉移圖

@startuml
skinparam backgroundColor #FEFEFE
skinparam defaultTextAlignment center
skinparam rectangleBackgroundColor #F5F5F5
skinparam rectangleBorderColor #333333
skinparam arrowColor #333333

title 附錄

rectangle "send" as node1
rectangle "完成寫入" as node2
rectangle "receive" as node3
rectangle "完成讀取" as node4

node1 --> node2
node2 --> node3
node3 --> node4

@enduml

圖表翻譯： 此圖表示通道的狀態轉移過程。通道從 EMPTY 狀態開始，當呼叫 send 方法時轉入 WRITING 狀態。完成寫入後，狀態轉為 READY。當呼叫 receive 方法時，狀態轉為 READING，並在完成讀取後回到 EMPTY 狀態。