Hedon Garden写作 · 构建 · 保持好奇
Rust 实战丨手写一个 Mutex 永久链接
Rust 实战丨手写一个 Mutex 的文章封面

Rust 实战集合

Rust 实战丨手写一个 Mutex

本文带你从零开始实现一个 Rust 中的 Mutex 锁，结合 Rust 原子操作和内存顺序的核心知识，逐步揭示 Mutex 背后的等待与唤醒机制。通过阅读本文，你不仅可以掌握如何使用 Rust 的原子 API 实现一个高效的互斥锁，还能深入理解原子操作背后的内存模型，为掌握更复杂的并发编程技巧打下坚实基础。

约 2,593 字

系列文章：

继上篇 Rust 原理丨操作系统并发原语，我们学习了不同操作系统下的并发原语实现，理解了它们最重要的贡献就是提供了一套 wait/wake_one/wake_all 的机制。本篇，我们将借助 Rust Atomics and Locks 的作者 Mara Bos 封装的 atomic-wait crate，来手写一个自己的 Mutex！

v1：基本实现

首先我们来思考一下如何定义数据结构：

  1. 我们需要 1 个原子变量 state 来记录锁的状态（0: unlocked, 1: locked），因为 atomic-wait 只支持 AtomicU32，所以这里我们的类型也定义为 AtomicU32。
  2. 另外我们需要一个 value 字段来保存数据，当抢到锁的时候，是可以对 value 进行修改的，但是这个时候只有共享引用，所以我们需要 UnsafeCell 来提供内部可变性。

同时，贯彻 *RAII（Resource Acquisition Is Initialization，资源获取即初始化）*原则：

  1. 我们在 lock(&self) 成功时返回一个 MutexGuard，它包含 &Mutex。
  2. 在 MutexGuard drop 的时候，我们将 state 重置为 0，表示释放锁，并唤醒一个潜在的阻塞线程。

为了让 Mutex 可以在线程之间共享，我们需要为其实现 Sync trait，而又因为 Mutex 实现的是独占访问，上锁成功的线程是拥有 T 的所有权的，即要求 T 可以在线程中转移，即要求 T 需要实现 Send trait。

综上，我们定义的 Mutex 和 MutexGuard 结构如下：

pub struct Mutex<T> {
    /// 0: unlocked
    /// 1: locked
    state: AtomicU32,
    value: UnsafeCell<T>,
}

pub struct MutexGuard<'a, T> {
    mutex: &'a Mutex<T>,
}

unsafe impl<T> Sync for Mutex<T> where T: Send {}

impl<T> Mutex<T> {
    pub fn new(value: T) -> Self {
        Self {
            state: AtomicU32::new(0),
            value: UnsafeCell::new(value),
        }
    }
}

为了方面访问内部数据，我们为 MutexGuard 实现 Deref 和 DerefMut 这 2 个 trait:

impl<T> Deref for MutexGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { &*self.mutex.value.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.mutex.value.get() }
    }
}

当 MutexGuard 离开作用域的时候，即被 drop 的时候，我们需要释放锁，并调用 wake_one 去唤醒一个潜在的阻塞线程：

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        self.mutex.state.store(0, Release);
        wake_one(&self.mutex.state);
    }
}

这里将 state 设置为 0，使用的内存顺序是 Release，是为了跟 lock 的时候使用 Acquire 建立 happens-before 原则，确保 state 的真实值在各个线程中都是可见的。

在 lock 的时候，我们需要将 state 从 0 替换为 1，如果成功，则说明上锁成功，直接返回 MutexGuard，如果失败，则说明锁已经被抢占了，这个时候我们使用 atomic-wait 的 wait() 陷入休眠，等待 wake_one 信号唤醒，再尝试抢锁。

impl<T> Mutex<T> {
  	// ...

    pub fn lock(&self) -> MutexGuard<T> {
        while self.state.swap(1, Acquire) == 1 {
            wait(&self.state, 1);
        }
        MutexGuard { mutex: self }
    }
}

至此，我们第一个版本的 Mutex 就完工了！是不是很简单！我们来写 2 个单元测试验证一下基本逻辑是否正确：

#[test]
fn one_thread_should_work() {
    let l = Mutex::new(vec![]);
    let mut guard = l.lock();
    guard.push(1);
    drop(guard);

    let guard = l.lock();
    assert_eq!(guard[0], 1);
}

#[test]
fn cross_thread_should_work() {
    let l = Mutex::new(vec![]);

    thread::scope(|s| {
        s.spawn(|| {
            let mut guard = l.lock();
            guard.push(1);
            sleep(Duration::from_millis(100)); // sleep for makeing the second thread to be blcoked.
        });

        sleep(Duration::from_millis(10)); // make sure the first thread get the lock
        s.spawn(|| {
            let mut guard = l.lock();
            guard.push(2);
        });
    });

    let guard = l.lock();
    assert_eq!(guard.len(), 2);
}

运行成功：

running 2 tests
test mutex::tests::one_thread_should_work ... ok
test mutex::tests::cross_thread_should_work ... ok

v2：减少系统调用

当 MutexGuard 的时候，我们将 state 置为 0，并调用 wake_one 唤醒一个潜在的线程，这个时候如果没有阻塞中的线程的话，那这个系统调用就比较浪费了。

所以在 v2 版本我们尝试来优化这一点。为此，我们需要扩展我们的 state 字段，新增表示是否有阻塞线程的能力。

pub struct Mutex<T> {
    /// 0: unlocked
    /// 1: locked, but no blocked thread
    /// 2: locked, but has blocked threads
    state: AtomicU32,
    value: UnsafeCell<T>,
}

修改了 state 的定义后我们需要修改上锁和解锁的逻辑，在上锁的时候，我们先尝试将 state 从 0 置为 1，如果成功了，说明抢到了锁，否则，我们将 state 置为 2，表示有线程被阻塞了。

这里书中的实现是这样的：

impl<T> Mutex<T> {
		// ...

    pub fn lock(&self) -> MutexGuard<T> {
        lock_contended(&self.state);
        MutexGuard { mutex: self }
    }
}

fn lock_contended(state: &AtomicU32) {
    if state.compare_exchange(0, 1, Acquire, Relaxed).is_err() {
        while state.swap(2, Acquire) != 0 {
            wait(&state, 2)
        }
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        if self.mutex.state.swap(0, Release) == 2 {
            println!("wake_one");
            wake_one(&self.mutex.state);
        }
    }
}

lock():

  • 如果成功将 state 从 0 变换为 1，则说明当前线程抢锁成功，直接返回 MutexGuard。
  • 如果失败了，就将 state 置为 2，然后调用 wait 进入休眠。

unlock():

  • 将 state 置为 0，如果之前是 2 的话，那就说明有线程被阻塞着，这个时候才调用 wake_one 唤醒一个阻塞的线程。

这个地方，笔者觉得有一些问题， state.swap(2, Acquire) 这一行代码会无条件将 state 置为 2，也就是说，当这个线程抢到锁后，它在 unlock() 的时候，无论有没有在阻塞的线程，这个时候 state 都是 2，所以都会调用 wake_one。

Mermaid
正在绘制图表…

我们可以运行上面的测试用例 cross_thread_should_work，可以看到输出了 2 个 wake_one，但是通过分析，应该只需要调用一次 wake_one 就足够了。

test mutex::tests::cross_thread_should_work ... ok

successes:

---- mutex::tests::cross_thread_should_work stdout ----
wake_one
wake_one

笔者的实现如下：

fn lock_contended(state: &AtomicU32) {
    while let Err(s) = state.compare_exchange(0, 1, Acquire, Relaxed) {
        if s == 1 {
            _ = state.compare_exchange(1, 2, Acquire, Relaxed);
        }
        wait(&state, 2)
    }
}
  1. 我们获取 state.compare_exchange(0,1) 的返回值，如果成功，说明抢到锁，直接返回。
  2. 如果失败了：
    • 原始值是 1，那我们就尝试将 state 从 1 交换为 2，然后调用 wait 陷入休眠。
    • 原始值是 2，说明已经有别的线程也被阻塞了，这个时候直接调用 wait 陷入休眠。
  3. 当被 wake_one 唤醒时，重新执行 state.compare_exchange(0,1) 抢占锁。

这个新的流程中，我们抢到锁的时候，state 会被正确的设置为 1 而不是 2，这个时候，在 drop 的时候就不会有不必要的 wake_one 的调用了。

Mermaid
正在绘制图表…

我们重新运行上面的测试用例 cross_thread_should_work，可以看到只输出了 1 个 wake_one：

test mutex::tests::cross_thread_should_work ... ok

successes:

---- mutex::tests::cross_thread_should_work stdout ----
wake_one

不过笔者在做 benchmark 后发现书中的实现性能其实更高，在 macbook m2max 机器上，书中的版本要比我的版本快 5~10% 左右，猜测大概率是 swap 的性能要比 compare_exchange 高。

v3：短暂自旋进一步避免系统调用

还有一种潜在的优化是，我们可以在抢锁失败且返回 state 为 1 的时候，进行短暂的自旋，如果实际场景中占用锁的时间非常短，那我们就可以再省略一次 wake 的系统调用了。

不过值得注意的是，这种优化未必是正向的，一方面，如果锁占用时间比较长，那前面的自旋就白白浪费了，另一方面，自旋的次数带来的性能消耗，未必就比系统调用要小（不同的平台表现可能很不一样）。

书中给出的经验值是自选 100 次。

优化后的 lock_contended 如下：

fn lock_contended(state: &AtomicU32) {
    let mut spin_count = 0;
    while let Err(s) = state.compare_exchange(0, 1, Acquire, Relaxed) {
        if s == 1 {
            if spin_count < 100 {
                spin_count += 1;
                std::hint::spin_loop();
                continue;
            }
            _ = state.compare_exchange(1, 2, Acquire, Relaxed);
        }
        wait(&state, 2)
    }
}

感兴趣的读者可以使用 criterion 做一个 benchmark 看看自旋与之前的版本的性能差异有多少。

完整代码可参考：conutils/mutex。

总结

在本篇中，我们从零开始，结合 Rust 原子操作和内存顺序的核心知识，实现一个 Rust 中的 Mutex 锁，逐步揭示 Mutex 背后的等待与唤醒机制，为更好理解标准库中的 Mutex 奠定了良好的基础。下篇，我们将尝试手写一个条件变量 Condition Variable！

Happy Coding! Peace~

#rust#并发编程#Rust 实战

来自开放网络的回应

引用、喜欢、转发和站外回复会通过 Webmention 回到这里。

接收与展示代码已经就绪，注册正式域名后即可开始收集回应。

评论

评论由 GitHub Discussions 保存和管理。