Rust Arc 弱指標實作與效能最佳化

Rust 的 Arc<T> 型別在多執行緒環境中提供安全的分享所有權，但引入弱指標機制後，效能開銷成為一個考量因素。本文提出的最佳化方案，將 Arc 和 Weak 指標都指向同一個 ArcData 結構，其中包含兩個原子計數器：data_ref_count 追蹤 Arc 的數量，alloc_ref_count 追蹤 Weak 的數量並在存在 Arc 時加一。此設計允許 clone 操作只需修改 data_ref_count，提升效能。drop 方法則在 data_ref_count 降至 1 時，釋放資料並遞減 alloc_ref_count。Weak 的 upgrade 方法嘗試增加 data_ref_count，成功則傳回 Arc，否則傳回 None。get_mut 方法則需確保 Arc 唯一性才能傳回可變參照。downgrade 方法則增加 alloc_ref_count 並傳回 Weak，同時處理了並發衝突的可能性。

最佳化

雖然弱指標很有用，但 Arc 型別經常在沒有弱指標的情況下使用。我們之前的實作的一個缺點是，複製和丟棄 Arc 現在都需要兩個原子操作，因為它們必須增加或減少兩個計數器。這使得所有 Arc 使用者都必須為弱指標付出代價，即使他們沒有使用它們。

新的最佳化實作

這次，我們不能簡單地將 Arc<T> 實作為 Weak<T> 的包裝器，因此兩者都將包裝一個非空指標指向分配：

pub struct Arc<T> {
    ptr: NonNull<ArcData<T>>,
}

pub struct Weak<T> {
    ptr: NonNull<ArcData<T>>,
}

內容解密：

1. NonNull<ArcData<T>>：使用 NonNull 來表示一個非空指標，指向 ArcData<T>。

ArcData 結構

struct ArcData<T> {
    /// Number of `Arc`s.
    data_ref_count: AtomicUsize,
    /// Number of `Weak`s, plus one if there are any `Arc`s.
    alloc_ref_count: AtomicUsize,
    /// The data. Dropped if there are only weak pointers left.
    data: UnsafeCell<ManuallyDrop<T>>,
}

內容解密：

1. data_ref_count：Arc 的計數器。
2. alloc_ref_count：Weak 的計數器，如果存在任何 Arc，則加一。
3. data: UnsafeCell<ManuallyDrop<T>>：儲存資料，使用 ManuallyDrop 以便手動丟棄資料。

新的實作細節

在這個新的實作中，Arc::new 函式幾乎保持不變，但使用 ManuallyDrop::new() 而不是 Some()：

impl<T> Arc<T> {
    pub fn new(data: T) -> Arc<T> {
        Arc {
            ptr: NonNull::from(Box::leak(Box::new(ArcData {
                alloc_ref_count: AtomicUsize::new(1),
                data_ref_count: AtomicUsize::new(1),
                data: UnsafeCell::new(ManuallyDrop::new(data)),
            }))),
        }
    }
}

內容解密：

1. ManuallyDrop::new(data)：使用 ManuallyDrop 包裝資料，以便手動控制資料的丟棄。

最佳化 Arc 實作：效能與安全性兼顧

在前面的章節中，我們探討了 Arc<T> 的基本實作原理。本章節將深入最佳化 Arc<T> 的實作，著重於提升效能和安全性。

最佳化 Clone 和 Drop 實作

最佳化後的 Arc<T> 實作中，clone 方法只需操作一個計數器，大大簡化了邏輯並提升了效能：

impl<T> Clone for Arc<T> {
    fn clone(&self) -> Self {
        if self.data().data_ref_count.fetch_add(1, Relaxed) > usize::MAX / 2 {
            std::process::abort();
        }
        Arc { ptr: self.ptr }
    }
}

內容解密：

1. data_ref_count.fetch_add(1, Relaxed)：使用 Relaxed 記憶體順序進行原子加一操作，快速增加參照計數。
2. 檢查溢位：若計數超過 usize::MAX / 2，則終止程式，防止溢位。
3. 傳回新的 Arc 例項：直接複製 ptr 指標，建立新的 Arc 例項。

同樣地，drop 方法的最佳化減少了操作次數，提升了效能：

impl<T> Drop for Arc<T> {
    fn drop(&mut self) {
        if self.data().data_ref_count.fetch_sub(1, Release) == 1 {
            fence(Acquire);
            unsafe {
                ManuallyDrop::drop(&mut *self.data().data.get());
            }
            drop(Weak { ptr: self.ptr });
        }
    }
}

內容解密：

1. data_ref_count.fetch_sub(1, Release)：使用 Release 記憶體順序減少參照計數，確保記憶體操作的正確性。
2. fence(Acquire)：確保在釋放資料前，所有對資料的存取都已完成。
3. 銷毀資料與弱參照計數：當最後一個 Arc 例項被丟棄時，釋放資料並減少弱參照計數。

最佳化 Weak 的 upgrade 方法

Weak<T> 的 upgrade 方法最佳化後，不再需要增加弱參照計數，直接嘗試升級為 Arc<T>：

impl<T> Weak<T> {
    pub fn upgrade(&self) -> Option<Arc<T>> {
        let mut n = self.data().data_ref_count.load(Relaxed);
        loop {
            if n == 0 {
                return None;
            }
            assert!(n < usize::MAX);
            if let Err(e) = self.data().data_ref_count.compare_exchange_weak(n, n + 1, Relaxed, Relaxed) {
                n = e;
                continue;
            }
            return Some(Arc { ptr: self.ptr });
        }
    }
}

內容解密：

1. data_ref_count.load(Relaxed)：讀取當前參照計數，檢查是否可以升級。
2. compare_exchange_weak：使用 CAS 操作嘗試增加參照計數，若失敗則重試。
3. 傳回 Arc 或 None：成功則傳回新的 Arc<T>，否則傳回 None。

實作 get_mut 方法

get_mut 方法用於取得 Arc<T> 的可變參照，需確保沒有其他 Arc<T> 或 Weak<T> 例項存在：

pub fn get_mut(arc: &mut Self) -> Option<&mut T> {
    if arc.data().alloc_ref_count.compare_exchange(1, usize::MAX, Acquire, Relaxed).is_err() {
        return None;
    }
    let is_unique = arc.data().data_ref_count.load(Relaxed) == 1;
    arc.data().alloc_ref_count.store(1, Release);
    if !is_unique {
        return None;
    }
    fence(Acquire);
    unsafe { Some(&mut *arc.data().data.get()) }
}

內容解密：

1. compare_exchange(1, usize::MAX, Acquire, Relaxed)：鎖定弱參照計數，防止並發操作。
2. 檢查唯一性：確認 data_ref_count 是否為 1，確保沒有其他 Arc<T> 存在。
3. 解鎖並傳回可變參照：若唯一，則傳回 &mut T，否則傳回 None。

實作 downgrade 方法

downgrade 方法將 Arc<T> 降級為 Weak<T>，需處理可能的並發衝突：

pub fn downgrade(arc: &Self) -> Weak<T> {
    let mut n = arc.data().alloc_ref_count.load(Relaxed);
    loop {
        if n == usize::MAX {
            std::hint::spin_loop();
            n = arc.data().alloc_ref_count.load(Relaxed);
            continue;
        }
        assert!(n < usize::MAX - 1);
        if let Err(e) = arc.data().alloc_ref_count.compare_exchange_weak(n, n + 1, Relaxed, Relaxed) {
            n = e;
            continue;
        }
        return Weak { ptr: arc.ptr };
    }
}

內容解密：

1. 檢查鎖定狀態：若 alloc_ref_count 為 usize::MAX，表示被鎖定，進行 spin wait。
2. compare_exchange_weak：嘗試增加弱參照計數，若失敗則重試。
3. 傳回 Weak：成功後傳回新的 Weak<T> 例項。

瞭解處理器架構與原子操作

在前面的章節中，我們已經學習瞭如何撰寫正確的平行程式碼。不過，若能對處理器層級的運作有基本的瞭解，將有助於我們撰寫出更有效率的程式碼。本章節將探討原子操作在處理器層級的實作方式、不同處理器架構之間的差異、以及快取與記憶體排序之間的關係。

處理器指令與組合語言簡介

當我們使用 Rust 或 C 等編譯語言撰寫軟體時，原始碼會被編譯成機器指令，這些指令能夠被處理器直接執行。這些機器指令是針對特定的處理器架構而設計的。

組合語言的基本概念

機器指令是以二進位形式編碼的，這對於人類來說是難以閱讀的。組合語言（Assembly）是機器指令的人類可讀表示形式。每一條指令通常由一個單詞或縮寫開頭，後面跟隨著運算元或運算元。

組合語言的範例如下：

mov rax, qword ptr [rbp - 8]
add rax, 1
mov qword ptr [rbp - 8], rax

在這個範例中，我們看到三條組合語言指令，分別是 mov（行動資料）和 add（加法運算）。這些指令操作的是暫存器（如 rax）和記憶體位置（如 [rbp - 8]）。

從原始碼到機器指令的轉換

編譯過程會將原始碼轉換成機器指令。在這個過程中，大部分的原始碼結構都會被破壞。根據最佳化程度的不同，函式和函式呼叫可能仍然可被辨識。然而，資料型別（如結構體或列舉）會被簡化為位元組和位址，而迴圈和條件判斷則會被轉換為基本的跳轉指令。

兩種主要的處理器架構：x86-64 與 ARM64

本章節將重點介紹兩種主要的 64 位元處理器架構：x86-64 和 ARM64。

x86-64 架構

x86-64 是由 Intel 和 AMD 實作的 64 位元 x86 架構，廣泛應用於筆記型電腦、桌上型電腦、伺服器和部分遊戲主機中。原本的 16 位元 x86 架構及其 32 位元擴充套件由 Intel 開發，而 64 位元版本則最初由 AMD 以 AMD64 的名義開發。Intel 也開發了自己的 64 位元架構 IA-64，但最終採用了更為流行的 AMD x86 擴充套件（以 IA-32e、EM64T 和後來的 Intel64 等名稱）。

ARM64 架構

ARM64 是由 ARM 開發的 64 位元架構，幾乎所有的現代行動裝置、高效能嵌入式系統，以及越來越多的筆記型電腦和桌上型電腦都採用了這一架構。它也被稱為 AArch64，是 ARMv8 架構的一部分。早期的 32 位元 ARM 架構在更多種類別的應用中被使用，從汽車到電子 COVID 測試裝置，許多流行的微控制器都是根據 ARMv6 和 ARMv7。

兩種架構的差異

這兩種架構在許多方面存在差異，其中最重要的是它們對原子操作的處理方式不同。瞭解這兩種架構下原子操作的實作方式，可以幫助我們對其他架構有更廣泛的理解。

處理器指令與原子操作

透過檢視編譯器的輸出結果，我們可以瞭解處理器實際執行的指令，從而對處理器層級的運作有更深入的瞭解。

使用編譯器輸出了解處理器指令

use std::sync::atomic::{AtomicUsize, Ordering};

fn main() {
    let counter = AtomicUsize::new(0);
    counter.fetch_add(1, Ordering::SeqCst);
}

編譯上述 Rust 程式碼後，我們可以檢視其對應的組合語言指令。不同的處理器架構會產生不同的指令。

x86-64 架構下的原子操作

在 x86-64 架構下，fetch_add 操作會被編譯成 lock xadd 指令：

lock xadd qword ptr [rbp - 8], 1

lock 字首用於確保操作的原子性。

ARM64 架構下的原子操作

在 ARM64 架構下，相同的操作可能會被編譯成一系列指令，使用 Load-Exclusive 和 Store-Conditional 指令來實作原子操作：

ldxr x8, [x9]
add x8, x8, #1
stxr w10, x8, [x9]
cbnz w10, loop

這個範例展示了 ARM64 如何使用 Load-Exclusive (ldxr) 和 Store-Conditional (stxr) 指令來實作原子操作。如果 Store-Conditional 失敗（即 stxr 傳回非零值），程式會重試操作。

記憶體排序與快取

記憶體排序（Memory Ordering）是指處理器執行記憶體存取操作的順序。不同的記憶體排序模型會影響程式的執行結果和效能。

記憶體排序模型

不同的處理器架構採用不同的記憶體排序模型。x86-64 採用強記憶體模型（Strong Memory Model），而 ARM64 則採用弱記憶體模型（Weak Memory Model）。

快取一致性

現代處理器通常具有多層快取。快取一致性協定（如 MESI 協定）確保了多個處理器核心之間的快取保持一致。

內容解密：

• 本章節探討了處理器層級的原子操作和記憶體排序。
• 瞭解不同處理器架構的差異有助於撰寫更有效的平行程式碼。
• 組合語言和處理器指令的瞭解對於深入理解電腦系統的運作至關重要。

隨著處理器架構的不斷演進，瞭解不同架構之間的差異和最佳實踐將變得越來越重要。未來，我們可以期待看到更多針對特定架構的最佳化和改進。

程式碼範例：

use std::sync::atomic::{AtomicUsize, Ordering};

fn main() {
    let counter = AtomicUsize::new(0);
    let result = counter.fetch_add(1, Ordering::SeqCst);
    println!("Counter: {}", result);
}

處理器架構比較

  graph LR
    A[x86-64] -->|強記憶體模型|> B[高效能運算]
    C[ARM64] -->|弱記憶體模型|> D[行動裝置]
    A --> E[桌上型電腦]
    C --> F[嵌入式系統]

圖表翻譯： 本圖表比較了 x86-64 和 ARM64 兩種處理器架構的主要特點和應用領域。x86-64 主要應用於桌上型電腦和高效能運算，而 ARM64 則廣泛應用於行動裝置和嵌入式系統。

參考資料

• 《深入理解電腦系統》（Computer Systems: A Programmer’s Perspective）
• 《Rust 程式設計語言》（The Rust Programming Language）

透過本章節的學習，我們對處理器層級的原子操作和記憶體排序有了更深入的瞭解。這些知識將有助於我們撰寫出更有效率、更正確的平行程式碼。未來，我們將繼續探討更多關於平行程式設計的高階主題。