Rust 弱指標 Weak Pointer 實作解析

Rust 的 Arc<T> 提供了弱指標 Weak<T> 來解決迴圈參照問題。Arc<T> 內部使用 ArcData<T> 結構儲存資料和計數器，包含 data_ref_count 追蹤 Arc 數量和 alloc_ref_count 追蹤 Arc 與 Weak 總數。Arc<T> 持有一個 Weak<T>，而 Weak<T> 則持有指向 ArcData<T> 的裸指標，允許 Arc 和 Weak 分享底層資料。Clone 實作中，Arc 複製會增加兩個計數器，而 Weak 只增加 alloc_ref_count。Drop 實作中，Arc 計數歸零時資料設為 None，Weak 計數歸零時釋放 ArcData。get_mut 確保只有單一 Arc 且無 Weak 時才允許可變借用。升級 Weak 指標使用比較並交換迴圈，避免資料已被釋放時錯誤升級。文章提供的單元測試程式碼驗證了 Weak 指標的升級機制和資料釋放的時機。

弱指標（Weak Pointers）實作解析

在 Rust 的 Arc<T>（原子參照計數）實作中，弱指標（Weak Pointers）扮演著關鍵角色，用於解決迴圈參照（cyclic references）問題。本章節將探討弱指標的實作原理及其在 Arc<T> 中的應用。

弱指標的基本概念

弱指標是一種特殊的智慧指標，它允許我們參照某個物件而不增加其參照計數。當我們需要打破物件之間的迴圈參照時，弱指標就顯得尤為重要。

`ArcData` 結構設計

在實作 Arc<T> 時，我們定義了一個名為 ArcData 的結構來儲存相關的資料：

struct ArcData<T> {
    /// 計算 `Arc` 的數量
    data_ref_count: AtomicUsize,
    /// 計算 `Arc` 和 `Weak` 的總數量
    alloc_ref_count: AtomicUsize,
    /// 儲存的資料，當僅剩弱指標時為 `None`
    data: UnsafeCell<Option<T>>,
}

這個結構包含了兩個原子計數器：data_ref_count 用於追蹤 Arc 的數量，而 alloc_ref_count 則用於追蹤 Arc 和 Weak 的總數量。此外，data 欄位用於儲存實際的資料，並使用 UnsafeCell 來實作內部可變性。

內容解密：

data_ref_count：記錄有多少個 Arc 指標參照著 T 物件，當其歸零時，表示沒有 Arc 指標參照，資料可以被丟棄。
alloc_ref_count：記錄有多少個 Arc 和 Weak 指標參照著 ArcData 本身，當其歸零時，表示 ArcData 可以被釋放。
data：儲存實際的資料，使用 Option<T> 來表示資料是否存在。當 data_ref_count 歸零時，資料會被設為 None。

`Arc` 和 `Weak` 的定義

pub struct Arc<T> {
    weak: Weak<T>,
}

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

在這個實作中，Arc<T> 被定義為包含一個 Weak<T>，而 Weak<T> 則持有指向 ArcData<T> 的裸指標。這樣的設計使得 Arc<T> 和 Weak<T> 可以共用相同的底層資料結構。

內容解密：

Arc<T>：包含一個 Weak<T>，用於管理 ArcData<T> 的生命週期。
Weak<T>：持有指向 ArcData<T> 的裸指標，用於弱參照 ArcData<T>。

`new` 方法實作

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

當建立一個新的 Arc<T> 時，我們會初始化 ArcData<T>，並將 alloc_ref_count 和 data_ref_count 都設為 1，表示目前有一個 Arc<T> 參照著資料。

內容解密：

初始化 ArcData<T>，並將計數器設為 1。
使用 Box::leak 將 ArcData<T> 組態在堆積疊上，並取得其裸指標。

`Deref` 實作

impl<T> Deref for Arc<T> {
    type Target = T;
    fn deref(&self) -> &T {
        let ptr = self.weak.data().data.get();
        // Safety: Since there's an Arc to the data,
        // the data exists and may be shared.
        unsafe { (*ptr).as_ref().unwrap() }
    }
}

在 Deref 實作中，我們透過 Weak 指標取得 ArcData<T> 的參照，並進一步取得資料的參照。由於我們知道目前有一個 Arc<T> 參照著資料，因此資料是有效的，可以安全地被參照。

內容解密：

透過 Weak 指標取得 ArcData<T> 的參照。
使用 UnsafeCell::get 取得資料的指標，並進行解參照。

`Clone` 實作

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

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

在 Clone 實作中，Weak<T> 的複製僅需增加 alloc_ref_count，而 Arc<T> 的複製則需要同時增加 data_ref_count 和 alloc_ref_count。

內容解密：

Weak<T> 的複製：增加 alloc_ref_count。
Arc<T> 的複製：增加 data_ref_count 和 alloc_ref_count。

`Drop` 實作

impl<T> Drop for Weak<T> {
    fn drop(&mut self) {
        if self.data().alloc_ref_count.fetch_sub(1, Release) == 1 {
            fence(Acquire);
            unsafe {
                drop(Box::from_raw(self.ptr.as_ptr()));
            }
        }
    }
}

impl<T> Drop for Arc<T> {
    fn drop(&mut self) {
        if self.weak.data().data_ref_count.fetch_sub(1, Release) == 1 {
            fence(Acquire);
            let ptr = self.weak.data().data.get();
            // Safety: The data reference counter is zero,
            // so nothing will access it.
            unsafe {
                (*ptr) = None;
            }
        }
    }
}

在 Drop 實作中，當 Weak<T> 被丟棄時，如果 alloc_ref_count 歸零，則會釋放 ArcData<T>。當 Arc<T> 被丟棄時，如果 data_ref_count 歸零，則會將資料設為 None。

內容解密：

Weak<T> 的丟棄：減少 alloc_ref_count，並在歸零時釋放 ArcData<T>。
Arc<T> 的丟棄：減少 data_ref_count，並在歸零時將資料設為 None。

`get_mut` 方法實作

impl<T> Arc<T> {
    pub fn get_mut(arc: &mut Self) -> Option<&mut T> {
        if arc.weak.data().alloc_ref_count.load(Relaxed) == 1 {
            fence(Acquire);
            // Safety: Nothing else can access the data, since
            // there's only one Arc, to which we have exclusive access,
            // and no Weak pointers.
            let arcdata = unsafe { arc.weak.ptr.as_mut() };
            let option = arcdata.data.get_mut();
            // We know the data is still available since we
            // have an Arc to it, so this won't panic.
            let data = option.as_mut().unwrap();
            Some(data)
        } else {
            None
        }
    }
}

get_mut 方法檢查是否只有一個 Arc<T> 參照著資料，並且沒有 Weak<T> 參照。如果滿足這些條件，則傳回一個可變參照給資料。

內容解密：

檢查 alloc_ref_count 是否為 1，以確保只有一個 Arc<T>。
使用 UnsafeCell::get_mut 取得資料的可變參照。

升級弱指標（Weak Pointer）

在 Rust 的 Arc（原子參照計數）實作中，升級弱指標（Weak）至 Arc 只有在資料仍然存在時才有可能。如果只剩下弱指標，那麼就沒有可以透過 Arc 分享的資料。因此，我們需要增加 Arc 計數器，但只有在計數器尚未達到零時才能進行此操作。我們將使用比較並交換迴圈（compare-and-exchange loop）來實作這一點。

為什麼使用比較並交換迴圈？

比較並交換迴圈是一種原子操作，可以用來更新分享變數。在這種情況下，我們使用它來增加 Arc 的計數器，同時確保計數器不會從零開始增加（因為那意味著資料已經被丟棄）。

程式碼實作：升級弱指標

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 { weak: self.clone() });
        }
    }
}

內容解密：

let mut n = self.data().data_ref_count.load(Relaxed);：首先，我們載入目前的 data_ref_count 值，使用 Relaxed 記憶體排序。Relaxed 排序表示該操作不受其他原子操作的順序限制，僅保證操作的原子性。
if n == 0 { return None; }：如果 data_ref_count 為零，表示資料已經被丟棄，因此無法升級 Weak 至 Arc，傳回 None。
assert!(n < usize::MAX);：確保 n 小於 usize::MAX，因為我們即將對 n 進行加一操作。如果 n 已經是 usize::MAX，那麼加一操作將導致溢位。
compare_exchange_weak：嘗試將 data_ref_count 從 n 更新為 n + 1。如果操作成功，表示我們成功升級了 Weak 至 Arc，並傳回 Some(Arc { weak: self.clone() })。如果操作失敗（因為其他執行緒可能同時修改了 data_ref_count），我們將使用最新的 data_ref_count 值重試。
Relaxed 記憶體排序：在這裡使用 Relaxed 排序是足夠的，因為我們只關心 data_ref_count 的值，而不涉及其他變數的順序。

從 Arc 取得 Weak

將 Arc<T> 轉換為 Weak<T> 的過程相對簡單：

impl<T> Arc<T> {
    pub fn downgrade(arc: &Self) -> Weak<T> {
        arc.weak.clone()
    }
}

內容解密：

arc.weak.clone()：直接複製 Arc 中的 Weak 指標。這裡假設 Arc 結構中包含一個 Weak 欄位。

測試實作

為了測試我們的實作，我們修改了之前的單元測試，使用弱指標並驗證它們在預期時可以被升級：

#[test]
fn test() {
    static NUM_DROPS: AtomicUsize = AtomicUsize::new(0);
    struct DetectDrop;
    impl Drop for DetectDrop {
        fn drop(&mut self) {
            NUM_DROPS.fetch_add(1, Relaxed);
        }
    }

    // Create an Arc with two weak pointers.
    let x = Arc::new(("hello", DetectDrop));
    let y = Arc::downgrade(&x);
    let z = Arc::downgrade(&x);
    let t = std::thread::spawn(move || {
        // Weak pointer should be upgradable at this point.
        let y = y.upgrade().unwrap();
        assert_eq!(y.0, "hello");
    });
    assert_eq!(x.0, "hello");
    t.join().unwrap();
    // The data shouldn't be dropped yet,
    // and the weak pointer should be upgradable.
    assert_eq!(NUM_DROPS.load(Relaxed), 0);
    assert!(z.upgrade().is_some());
    drop(x);
    // Now, the data should be dropped, and the
    // weak pointer should no longer be upgradable.
    assert_eq!(NUM_DROPS.load(Relaxed), 1);
    assert!(z.upgrade().is_none());
}

內容解密：

DetectDrop 結構：用於檢測何時資料被丟棄。透過實作 Drop 特性，我們可以在 DetectDrop 被丟棄時增加 NUM_DROPS 計數器。
Arc::new(("hello", DetectDrop))：建立一個包含 ("hello", DetectDrop) 的 Arc，並建立兩個弱指標 y 和 z。
y.upgrade().unwrap()：在另一個執行緒中升級 y，並驗證資料是否仍然存在。
assert_eq!(NUM_DROPS.load(Relaxed), 0)：驗證在丟棄 x 之前，資料尚未被丟棄。
assert!(z.upgrade().is_some())：驗證在丟棄 x 之前，z 仍然可以被升級。
drop(x)：丟棄 x 後，資料應該被丟棄，且 z 不再可以被升級。

玄貓

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。

Rust 弱指標 Weak Pointer 實作解析

弱指標（Weak Pointers）實作解析

弱指標的基本概念

ArcData 結構設計

內容解密：

Arc 和 Weak 的定義

內容解密：

new 方法實作

內容解密：

Deref 實作

內容解密：

Clone 實作

內容解密：

Drop 實作

內容解密：

get_mut 方法實作

內容解密：

升級弱指標（Weak Pointer）

為什麼使用比較並交換迴圈？

程式碼實作：升級弱指標

內容解密：

從 Arc 取得 Weak

內容解密：

測試實作

內容解密：

玄貓

`ArcData` 結構設計

`Arc` 和 `Weak` 的定義

`new` 方法實作

`Deref` 實作

`Clone` 實作

`Drop` 實作

`get_mut` 方法實作