在“引用计数”中,我们看到了 std::sync::Arc<T> 类型,它允许通过引用计数进行共享所有权。Arc::new 函数创建一个新的内存分配,就像 Box::new 一样。然而,与 Box 不同,克隆后的 Arc 将共享原始的内存分配,而不会创建一个新的分配。只有当 Arc 及其所有克隆都被丢弃时,共享的内存才会被释放。
在实现这种类型时,涉及的内存顺序问题可能非常有趣。在本章中,我们将通过实现我们自己的 Arc<T> 来将更多的理论付诸实践。我们将从一个基础版本开始,然后扩展以支持循环结构的弱指针,最后以一个几乎与标准库中的实现相同的优化版本结束本章。
基础引用计数
我们第一个版本将使用一个 AtomicUsize 来计数共享一个分配的 Arc 对象数量。让我们从一个包含这个计数器和 T 对象的结构体开始:
struct ArcData<T> {
ref_count: AtomicUsize,
data: T,
}
注意,这个结构体不是公开的,它是我们 Arc 实现中的内部实现细节。
接下来是 Arc<T> 结构体本身,它实际上只是一个指向(共享的)ArcData<T> 对象的指针。
也许会有将其作为 Box<ArcData<T>> 的包装器的想法,使用标准的 Box 来处理 ArcData<T> 的分配。然而,Box 代表独占所有权,而不是共享所有权。我们也不能使用引用,因为我们并不是在借用其他对象拥有的数据,其生命周期(“直到最后一个克隆的 Arc 被丢弃”)无法直接用 Rust 的生命周期表示。
相反,我们需要使用一个指针并手动处理内存分配和所有权的概念。我们不会使用 *mut T 或 *const T,而是使用 std::ptr::NonNull<T>,它表示一个永远不为 null 的指向 T 的指针。这样,Option<Arc<T>> 的大小将与 Arc<T> 相同,使用空指针表示 None。
use std::ptr::NonNull;
pub struct Arc<T> {
ptr: NonNull<ArcData<T>>,
}
对于引用或 Box,编译器会自动理解对于哪个 T 需要使你的结构体实现 Send 或 Sync。但是使用原始指针或 NonNull 时,编译器会保守地认为它永远不是 Send 或 Sync,除非我们明确告诉它。
将 Arc<T> 跨线程发送意味着共享一个 T 对象,要求 T 是 Sync。类似地,将 Arc<T> 跨线程发送可能会导致另一个线程丢弃 T,这实际上将它转移到另一个线程中,这需要 T 是 Send。换句话说,只有当 T 同时是 Send 和 Sync 时,Arc<T> 才应该是 Send。对于 Sync,情况也是一样的,因为共享的 &Arc<T> 可以克隆为一个新的 Arc<T>。
unsafe impl<T: Send + Sync> Send for Arc<T> {}
unsafe impl<T: Send + Sync> Sync for Arc<T> {}
对于 Arc<T>::new,我们需要创建一个带有引用计数为 1 的 ArcData<T> 的新分配。我们将使用 Box::new 来创建一个新的分配,使用 Box::leak 放弃对该分配的独占所有权,并使用 NonNull::from 将其转换为一个指针:
impl<T> Arc<T> {
pub fn new(data: T) -> Arc<T> {
Arc {
ptr: NonNull::from(Box::leak(Box::new(ArcData {
ref_count: AtomicUsize::new(1),
data,
}))),
}
}
// ...
}
我们知道,只要 Arc 对象存在,指针将始终指向有效的 ArcData<T>。然而,这并不是编译器知道或检查的,因此通过指针访问 ArcData 需要使用不安全代码。我们将添加一个私有的辅助函数,以从 Arc 获取 ArcData,因为这是我们将多次执行的操作:
fn data(&self) -> &ArcData<T> {
unsafe { self.ptr.as_ref() }
}
通过这个函数,我们现在可以实现 Deref 特性,使我们的 Arc<T> 透明地表现得像 T 的引用:
impl<T> Deref for Arc<T> {
type Target = T;
fn deref(&self) -> &T {
&self.data().data
}
}
注意,我们没有实现 DerefMut。由于 Arc<T> 代表共享所有权,我们无法无条件地提供 &mut T。
接下来是 Clone 实现。克隆后的 Arc 将使用相同的指针,在此之前增加引用计数器:
impl<T> Clone for Arc<T> {
fn clone(&self) -> Self {
// TODO: 处理溢出。
self.data().ref_count.fetch_add(1, Relaxed);
Arc {
ptr: self.ptr,
}
}
}
我们可以使用 Relaxed 内存顺序来增加引用计数,因为在这个原子操作之前或之后没有对其他变量的操作需要严格发生。通过原始的 Arc,我们已经能够访问包含的 T,这在之后(现在通过至少两个 Arc 对象)保持不变。
一个 Arc 需要被克隆很多次,引用计数才有可能溢出,但在循环中运行 std::mem::forget(arc.clone()) 可以使其发生。我们可以使用“示例:ID 分配”和“示例:无溢出的 ID 分配”中讨论的任何技术来处理此问题。
为了在正常(未溢出)情况下尽可能保持高效,我们将保留原始的 fetch_add,并在接近溢出时终止整个进程:
if self.data().ref_count.fetch_add(1, Relaxed) > usize::MAX / 2 {
std::process::abort();
}
注意
终止进程并不是瞬时的,在这段时间内,另一个线程也可能调用 Arc::clone,进一步增加引用计数。因此,仅检查 usize::MAX - 1 是不够的。然而,使用 usize::MAX / 2 作为限制是可以的:假设每个线程至少占用几字节的内存,usize::MAX / 2 个线程不可能同时存在。
就像我们在克隆时增加计数器一样,我们在丢弃 Arc 时需要减少它。当看到计数器从 1 变为 0 的线程知道它丢弃了最后一个 Arc<T>,因此负责丢弃并释放 ArcData<T>。
我们将使用 Box::from_raw 来重新获得对分配的独占所有权,然后直接使用 drop() 将其丢弃:
impl<T> Drop for Arc<T> {
fn drop(&mut self) {
// TODO: 内存顺序。
if self.data().ref_count.fetch_sub(1, …) == 1 {
unsafe {
drop(Box::from_raw(self.ptr.as_ptr()));
}
}
}
}
对于这个操作,我们不能使用 Relaxed 顺序,因为我们需要确保在我们丢弃数据时没有其他东西仍在访问它。换句话说,所有以前的 Arc 克隆的每次丢弃都必须发生在最后一次丢弃之前。因此,最终的 fetch_sub 必须与之前的每次 fetch_sub 操作建立一个 happens-before 关系,我们可以使用释放和获取顺序来做到这一点:从 2 减到 1 有效地“释放”数据,而从 1 减到 0 则“获取”它的所有权。
我们可以使用 AcqRel 内存顺序来涵盖这两种情况,但只有最终减少到 0 时才需要 Acquire,而其他情况下只需要 Release。为了提高效率,我们只对 fetch_sub 操作使用 Release,并在必要时使用一个单独的 Acquire 栅栏:
if self.data().ref_count.fetch_sub(1, Release) == 1 {
fence(Acquire);
unsafe {
drop(Box::from_raw(self.ptr.as_ptr()));
}
}
测试我们的 Arc
为了测试我们的 Arc 是否按预期工作,我们可以编写一个单元测试,创建一个包含特殊对象的 Arc,当它被丢弃时可以告诉我们:
#[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);
}
}
// 创建两个共享对象的 Arc,其中包含一个字符串和一个 DetectDrop,用于检测何时被丢弃。
let x = Arc::new(("hello", DetectDrop));
let y = x.clone();
// 将 x 发送到另一个线程,并在该线程中使用它。
let t = std::thread::spawn(move || {
assert_eq!(x.0, "hello");
});
// 与此同时,y 在这里应该仍然可以使用。
assert_eq!(y.0, "hello");
// 等待线程完成。
t.join().unwrap();
// 一个 Arc(x)应该已经被丢弃。
// 我们仍然拥有 y,所以对象不应该被丢弃。
assert_eq!(NUM_DROPS.load(Relaxed), 0);
// 丢弃剩余的 Arc。
drop(y);
// 现在 y 也被丢弃了,对象应该已经被丢弃。
assert_eq!(NUM_DROPS.load(Relaxed), 1);
}
这个程序可以编译并运行良好,因此看起来我们的 Arc 确实按预期工作!虽然这是一个令人鼓舞的迹象,但这并不能证明实现是完全正确的。建议使用涉及多个线程的长时间压力测试来增加信心。
MIRI
使用 Miri 运行测试也非常有用。Miri 是一个实验性但非常强大且有用的工具,用于检查不安全代码中各种形式的未定义行为。
Miri 是 Rust 编译器中级中间表示的解释器。这意味着它不是通过将代码编译为本机处理器指令来运行程序,而是在类型和生命周期等信息仍然可用时解释代码。因此,Miri 运行程序的速度比正常编译和运行的程序慢得多,但它能够检测到许多会导致未定义行为的错误。
它包括检测数据竞争的实验性支持,这使它能够检测内存顺序问题。
有关更多详细信息和如何使用 Miri 的指南,请参阅其 GitHub 页面。
可变性
如前所述,我们无法为我们的 Arc 实现 DerefMut。我们无法无条件地承诺对数据的独占访问(&mut T),因为可能通过其他 Arc 对象访问该数据。
然而,我们可以有条件地允许这一点。我们可以创建一个方法,当且仅当引用计数为 1 时,才返回 &mut T,以证明没有其他 Arc 对象可以用来访问相同的数据。
这个函数,我们称之为 get_mut,需要使用 &mut Self,以确保没有其他东西可以使用同一个 Arc 访问 T。如果那个唯一的 Arc 仍然可以被共享,那么知道只有一个 Arc 就毫无意义。
我们需要使用获取(acquire)内存顺序来确保先前拥有 Arc 克隆的线程不再访问该数据。我们需要与导致引用计数变为 1 的每次丢弃操作建立一个 happens-before 关系。
这只有在引用计数确实为 1 时才重要;如果引用计数较高,我们将不会提供 &mut T,此时内存顺序无关紧要。因此,我们可以先使用一个松散的(relaxed)加载,然后在需要时使用一个条件获取栅栏,如下所示:
pub fn get_mut(arc: &mut Self) -> Option<&mut T> {
if arc.data().ref_count.load(Relaxed) == 1 {
fence(Acquire);
// Safety: 没有其他东西可以访问该数据,因为只有一个 Arc,
// 并且我们对此拥有独占访问权。
unsafe { Some(&mut arc.ptr.as_mut().data) }
} else {
None
}
}
这个函数不接受 self 参数,而是接受一个常规参数(名为 arc)。这意味着它只能被调用为 Arc::get_mut(&mut a),而不能调用为 a.get_mut()。对于实现了 Deref 的类型,建议使用这种方式,以避免与底层 T 上具有相似名称的方法产生歧义。
返回的可变引用隐式地借用自参数的生命周期,这意味着只要返回的 &mut T 仍然存在,原始的 Arc 就无法使用,从而允许安全地修改数据。
当 &mut T 的生命周期结束后,Arc 可以再次被使用并与其他线程共享。人们可能会怀疑是否需要担心随后线程访问数据的内存顺序。然而,这取决于用于与另一个线程共享 Arc(或其新克隆体)的机制。(例如,互斥锁(mutex)、通道(channel)或生成新线程)。
弱指针
引用计数在表示由多个对象组成的内存结构时非常有用。例如,树结构中的每个节点都可以包含一个指向其子节点的 Arc。这样,当一个节点被丢弃时,它的子节点如果不再被使用,也会递归地被丢弃。
然而,对于循环结构,这种做法就行不通了。如果子节点也包含一个指向父节点的 Arc,那么父节点和子节点将永远不会被丢弃,因为始终至少有一个 Arc 引用它们。
标准库中的 Arc 提供了此问题的解决方案:Weak<T>。Weak<T>,也称为弱指针,行为类似于 Arc<T>,但并不阻止对象被丢弃。一个 T 可以在多个 Arc<T> 和 Weak<T> 对象之间共享,但当所有 Arc<T> 对象都消失时,无论是否还有 Weak<T> 对象,T 都会被丢弃。
这意味着 Weak<T> 可以存在于没有 T 的情况下,因此不能像 Arc<T> 一样无条件地提供 &T。然而,要访问 Weak<T> 中的 T,可以通过其 upgrade() 方法将其升级为 Arc<T>。如果 T 已经被丢弃,则此方法返回 None。
在基于 Arc 的结构中,Weak 可以用来打破循环。例如,树结构中的子节点可以使用 Weak 而不是 Arc 来引用父节点。这样,父节点的丢弃不会因为其子节点的存在而被阻止。
让我们来实现这一点。
就像之前一样,当 Arc 的数量达到零时,我们可以丢弃包含的 T 对象。然而,我们不能立即丢弃并释放 ArcData,因为可能仍然有弱指针引用它。只有当最后一个弱指针也消失时,我们才可以丢弃并释放 ArcData。
因此,我们将使用两个计数器:一个用于“引用 T 的数量”,另一个用于“引用 ArcData<T> 的数量”。换句话说,第一个计数器和之前相同:它计数 Arc 对象,而第二个计数器计数 Arc 和 Weak 对象的总数。
我们还需要一些东西来允许我们在 ArcData<T> 仍被弱指针使用时丢弃包含的对象 (T)。我们将使用 Option<T>,以便在数据被丢弃时使用 None,并将其包装在 UnsafeCell 中以实现内部可变性(详见“内部可变性”),以允许在 ArcData<T> 不是独占拥有的情况下进行修改:
struct ArcData<T> {
/// `Arc` 的数量。
data_ref_count: AtomicUsize,
/// `Arc` 和 `Weak` 的总数。
alloc_ref_count: AtomicUsize,
/// 数据。当只剩下弱指针时为 `None`。
data: UnsafeCell<Option<T>>,
}
如果我们将 Weak<T> 视为负责保持 ArcData<T> 活着的对象,那么将 Arc<T> 实现为包含一个 Weak<T> 的结构体是有意义的,因为 Arc<T> 需要做相同的事情,甚至更多。
pub struct Arc<T> {
weak: Weak<T>,
}
pub struct Weak<T> {
ptr: NonNull<ArcData<T>>,
}
unsafe impl<T: Sync + Send> Send for Weak<T> {}
unsafe impl<T: Sync + Send> Sync for Weak<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)),
}))),
},
}
}
}
和之前一样,我们假设 ptr 字段总是指向有效的 ArcData<T>。这次我们将该假设编码为 Weak<T> 上的私有 data() 辅助方法:
impl<T> Weak<T> {
fn data(&self) -> &ArcData<T> {
unsafe { self.ptr.as_ref() }
}
}
在 Arc<T> 的 Deref 实现中,我们现在必须使用 UnsafeCell::get() 来获取指向单元内容的指针,并使用不安全代码来保证它在此时可以安全共享。我们还需要 as_ref().unwrap() 来获取 Option<T> 中的引用。我们不必担心这会导致 panic,因为只有在没有 Arc 对象时,Option 才会是 None。
impl<T> Deref for Arc<T> {
type Target = T;
fn deref(&self) -> &T {
let ptr = self.weak.data().data.get();
// 安全性:由于有 `Arc` 引用数据,所以数据存在且可以共享。
unsafe { (*ptr).as_ref().unwrap() }
}
}
Weak<T> 的 Clone 实现相对直接,与之前的 Arc<T> 实现几乎相同:
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 }
}
}
新 Arc<T> 的 Clone 实现需要同时增加两个计数器。我们可以简单地使用 self.weak.clone() 来重用上面的代码来增加第一个计数器,因此只需手动增加第二个计数器:
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 }
}
}
丢弃 Weak 应该递减其计数器,并在计数器从 1 变为 0 时丢弃并释放 ArcData,这与之前 Arc 的 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()));
}
}
}
}
丢弃 Arc 应该递减两个计数器。注意,其中一个已经被自动处理了,因为每个 Arc 都包含一个 Weak,所以丢弃 Arc 也会导致丢弃 Weak。我们只需要处理另一个计数器:
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();
// 安全性:数据引用计数为零,因此没有东西会访问它。
unsafe {
(*ptr) = None;
}
}
}
}
get_mut 方法中的检查大部分保持不变,只是现在需要考虑弱指针。虽然似乎可以在检查独占性时忽略弱指针,但 Weak<T> 可以随时升级为 Arc<T>。因此,get_mut 必须检查是否没有其他 Arc<T> 或 Weak<T> 指针,才能提供 &mut T:
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);
// 安全性:没有其他东西可以访问数据,因为只有一个 Arc,
// 并且没有 Weak 指针。
let arcdata = unsafe { arc.weak.ptr.as_mut() };
let option = arcdata.data.get_mut();
// 我们知道数据仍然可用,因为我们有一个 Arc 引用它,所以不会 panic。
let data = option.as_mut().unwrap();
Some(data)
} else {
None
}
}
}
接下来是升级弱指针。将 Weak 升级为 Arc 只有在数据仍然存在时才有可能。如果只剩下弱指针,则没有可通过 Arc 共享的数据。因此,我们需要增加 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() });
}
}
}
提示
注意这次我们可以检查 n < usize::MAX,因为如果断言失败,它会在我们修改 data_ref_count 之前 panic。
与之相对,从 Arc<T> 获取 Weak<T> 则要简单得多:
impl<T> Arc<T> {
pub fn downgrade(arc: &Self) -> Weak<T> {
arc.weak.clone()
}
}
测试
为了快速测试我们的实现,我们将修改之前的单元测试来使用弱指针,并验证它们是否能在预期的情况下被升级:
#[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);
}
}
// 创建一个 Arc,并创建两个弱指针。
let x = Arc::new(("hello", DetectDrop));
let y = Arc::downgrade(&x);
let z = Arc::downgrade(&x);
let t = std::thread::spawn(move || {
// 此时弱指针应该可以被升级。
let y = y.upgrade().unwrap();
assert_eq!(y.0, "hello");
});
assert_eq!(x.0, "hello");
t.join().unwrap();
// 数据还不应该被丢弃,弱指针应该可以被升级。
assert_eq!(NUM_DROPS.load(Relaxed), 0);
assert!(z.upgrade().is_some());
drop(x);
// 现在,数据应该被丢弃,弱指针不再可升级。
assert_eq!(NUM_DROPS.load(Relaxed), 1);
assert!(z.upgrade().is_none());
}
这段代码也可以编译并正常运行,这让我们得到了一个非常实用的手工实现的 Arc。
优化
这个内容涉及到使用 Arc 和 Weak 智能指针进行内存管理和优化的实现,特别是如何通过减少原子操作来提高性能。以下是译文:
优化 虽然弱指针有时非常有用,但 Arc 类型通常在不使用任何弱指针的情况下使用。我们上一个实现的一个缺点是,克隆和删除 Arc 现在都需要两个原子操作,因为它们必须增加或减少两个计数器。这意味着即使不使用弱指针,所有 Arc 用户也要为它们的开销“买单”。
似乎解决方案是将 Arc<T> 和 Weak<T> 的指针分别计数,但这样我们就无法原子地检查两个计数器是否都为零。为了理解这个问题,假设有一个线程在执行以下令人头疼的函数:
fn annoying(mut arc: Arc<Something>) {
loop {
let weak = Arc::downgrade(&arc);
drop(arc);
println!("I have no Arc!");
arc = weak.upgrade().unwrap();
drop(weak);
println!("I have no Weak!");
}
}
这个线程不断地将 Arc 降级为 Weak 并再次升级为 Arc,这样它就不断循环通过持有没有 Arc 的时刻,以及没有 Weak 的时刻。如果我们检查两个计数器,查看是否还有线程在使用这个分配,那么在这个线程的第一个打印语句期间检查 Arc 计数器,而在第二个打印语句期间检查 Weak 计数器,它可能会掩盖其存在。
在我们之前的实现中,我们通过将每个 Arc 也计为一个 Weak 来解决这个问题。一种更巧妙的方法是将所有 Arc 指针合并计为一个单独的 Weak 指针。这样,只要还有至少一个 Arc 对象存在,弱指针计数器(alloc_ref_count)就不会变为零,和我们之前的实现一样,但克隆 Arc 时不需要碰这个计数器。只有删除最后一个 Arc 时,才会减少弱指针计数器。
让我们尝试一下。
这次我们不能再简单地将 Arc<T> 实现为 Weak<T> 的包装器,因此两者都会包装一个指向分配内存的非空指针:
pub struct Arc<T> {
ptr: NonNull<ArcData<T>>,
}
unsafe impl<T: Sync + Send> Send for Arc<T> {}
unsafe impl<T: Sync + Send> Sync for Arc<T> {}
pub struct Weak<T> {
ptr: NonNull<ArcData<T>>,
}
unsafe impl<T: Sync + Send> Send for Weak<T> {}
unsafe impl<T: Sync + Send> Sync for Weak<T> {}
既然我们在优化实现,不妨通过使用 std::mem::ManuallyDrop<T> 来让 ArcData<T> 变得稍微小一些,而不是使用 Option<T>。我们使用 Option<T> 来替换在删除数据时的 Some(T) 为 None,但我们实际上不需要单独的 None 状态来告诉我们数据已经消失,因为 Arc<T> 的存在与否已经告诉我们这一点。ManuallyDrop<T> 占用的空间与 T 相同,但允许我们通过不安全调用 ManuallyDrop::drop() 在任意时刻手动删除它:
use std::mem::ManuallyDrop;
struct ArcData<T> {
/// `Arc` 的数量。
data_ref_count: AtomicUsize,
/// `Weak` 的数量,如果有任何 `Arc`,则加一。
alloc_ref_count: AtomicUsize,
/// 数据。当只有弱指针剩下时将被删除。
data: UnsafeCell<ManuallyDrop<T>>,
}
Arc::new 函数基本保持不变,像以前一样将两个计数器初始化为 1,但现在使用 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)),
}))),
}
}
…
}
Deref 实现不再使用 Weak 类型上的私有 data 方法,因此我们将在 Arc<T> 上添加相同的私有辅助函数:
impl<T> Arc<T> {
…
fn data(&self) -> &ArcData<T> {
unsafe { self.ptr.as_ref() }
}
…
}
impl<T> Deref for Arc<T> {
type Target = T;
fn deref(&self) -> &T {
// 安全性:由于存在一个 `Arc` 指向数据,
// 数据存在且可以共享。
unsafe { &*self.data().data.get() }
}
}
Weak<T> 的 Clone 和 Drop 实现与我们之前的实现完全相同。以下是完整实现,包括私有的 Weak::data 辅助函数:
impl<T> Weak<T> {
fn data(&self) -> &ArcData<T> {
unsafe { self.ptr.as_ref() }
}
…
}
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> 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()));
}
}
}
}
接下来是新优化实现的关键部分——克隆 Arc<T> 现在只需要操作一个计数器:
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 }
}
}
类似地,删除 Arc<T> 现在也只需要减少一个计数器,除了当计数器从 1 变为 0 的最后一次删除。在这种情况下,弱指针计数器也需要减少,以便在没有弱指针存在时达到零。我们通过简单地从空中创建一个 Weak<T> 并立即删除它来完成这一操作:
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());
}
// 现在没有 `Arc<T>` 剩下了,
// 删除代表所有 `Arc<T>` 的隐式弱指针。
drop(Weak { ptr: self.ptr });
}
}
}
Weak<T> 的 upgrade 方法保持大致相同,只是它不再克隆一个弱指针,因为它不需要再增加弱指针计数器。升级只有在分配内存中至少有一个 Arc<T> 时才会成功,这意味着 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 { ptr: self.ptr });
}
}
}
到目前为止,这个实现与我们之前的实现之间的差异非常小。不过,事情变得复杂的是我们还需要实现的最后两个方法:downgrade 和 get_mut。
与之前不同的是,get_mut 方法现在需要检查两个计数器是否都为一,以确定是否只有一个 Arc<T> 并且没有 Weak<T> 剩余,因为一个弱指针计数器的值为 1 现在可能代表多个 Arc<T> 指针。读取计数器是两个独立的操作,它们发生在(略有)不同的时间,因此我们必须非常小心,以免错过任何并发的降级操作,比如我们在“优化”部分开头看到的示例情况。
如果我们首先检查 data_ref_count 是否为 1,那么我们可能会在检查另一个计数器之前错过随后的 upgrade() 操作。但是,如果我们首先检查 alloc_ref_count 是否为 1,那么我们可能会在检查另一个计数器之前错过随后的 downgrade() 操作。
解决这种两难困境的方法是通过“锁定”弱指针计数器来暂时阻止 downgrade() 操作。为此,我们不需要像互斥锁这样的东西。我们可以使用一个特殊的值,比如 usize::MAX,来表示弱指针计数器的特殊“锁定”状态。它只会被锁定非常短的时间,仅用于加载另一个计数器,因此在极不可能的情况下,downgrade 方法可以一直循环,直到计数器解锁,若它恰好在与 get_mut 相同的时刻运行。
因此,在 get_mut 中,我们首先需要检查 alloc_ref_count 是否为 1,并同时将其替换为 usize::MAX,如果它确实是 1。这需要使用 compare_exchange。
然后我们需要检查另一个计数器是否也是 1,之后我们可以立即解锁弱指针计数器。如果第二个计数器也是 1,我们就知道我们对分配和数据有独占访问权限,因此可以返回一个 &mut T。
pub fn get_mut(arc: &mut Self) -> Option<&mut T> {
// Acquire 与 Weak::drop 的 Release 匹配,以确保在下次 data_ref_count.load 之前可以看到任何升级的指针。
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;
// Release 与 `downgrade` 中的 Acquire 增量匹配,以确保 `downgrade` 之后对 data_ref_count 的任何更改不会更改上面的 is_unique 结果。
arc.data().alloc_ref_count.store(1, Release);
if !is_unique {
return None;
}
// Acquire 与 Arc::drop 的 Release 匹配,以确保没有其他内容在访问数据。
fence(Acquire);
unsafe { Some(&mut *arc.data().data.get()) }
}
正如你可能预料到的,锁定操作(compare_exchange)必须使用 Acquire 内存排序,而解锁操作(store)必须使用 Release 内存排序。
如果我们对 compare_exchange 使用 Relaxed,那么在随后从 data_ref_count 的加载中可能看不到新升级的弱指针的值,即使 compare_exchange 已经确认所有弱指针都被删除。
如果我们对 store 使用 Relaxed,那么前面的加载可能会观察到未来某个 Arc::drop 的结果,而这个 Arc 仍然可以被降级。
获取屏障(fence(Acquire))与之前一样:它与 Arc::Drop 中的释放递减操作同步,以确保通过之前的所有 Arc 克隆的访问在新的独占访问之前完成。
最后一部分是 downgrade 方法,它必须检查特殊的 usize::MAX 值,以查看弱指针计数器是否被锁定,并一直循环直到它解锁。就像在 upgrade 实现中一样,我们将使用一个比较并交换的循环来检查特殊值并在计数器增加之前防止溢出:
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);
// Acquire 与 get_mut 的 release-store 同步。
if let Err(e) =
arc.data()
.alloc_ref_count
.compare_exchange_weak(n, n + 1, Acquire, Relaxed)
{
n = e;
continue;
}
return Weak { ptr: arc.ptr };
}
}
我们对 compare_exchange_weak 使用 Acquire 内存排序,它与 get_mut 函数中的释放存储同步。否则,在 get_mut 解锁计数器之前,可能会有一个后续的 Arc::drop 的效果对运行 get_mut 的线程可见。
换句话说,这里的获取比较和交换操作实际上“锁定”了 get_mut,防止它成功。它可以通过稍后将计数器减回 1 的 Weak::drop 再次“解锁”,使用释放内存排序。
注意
我们刚刚制作的 Arc<T> 和 Weak<T> 的优化实现几乎与 Rust 标准库中包含的实现相同。如果我们运行之前相同的测试(“测试”),我们会看到这个优化的实现也能编译并通过我们的测试。
提示
如果你觉得为这个优化实现做出正确的内存排序决策很困难,不要担心。许多并发数据结构比这个更容易正确实现。这个 Arc 实现被包含在这一章中特别是因为它在内存排序方面存在棘手的微妙之处。
总结
Arc<T>提供对引用计数分配的共享所有权。- 通过检查引用计数器是否正好为 1,
Arc<T>可以有条件地提供独占访问(&mut T)。 - 可以使用松弛操作来增加原子引用计数器,但最终的递减必须与之前的所有递减操作同步。
- 弱指针(
Weak<T>)可以用来避免循环引用。 NonNull<T>类型表示一个从不为空的指向T的指针。ManuallyDrop<T>类型可以用来使用不安全代码手动决定何时销毁T。- 一旦涉及多个原子变量,事情就会变得更加复杂。
- 实现一个特定(自旋)锁有时可以作为对多个原子变量同时操作的有效策略。