Tokio 源码学习01——Mutex

清醒的土土土
105 阅读10分钟

一、为什么需要异步 Mutex？

在开始深入 Tokio Mutex 之前，我们先思考一个问题：Rust 标准库不是已经有 std::sync::Mutex 了吗？为什么 Tokio 还要重新实现一个？

答案在于 异步编程与同步编程的本质差异。

1.1 核心问题：.await 期间持有锁

// 使用 std::sync::Mutex 的错误示例
use std::sync::Mutex;

async fn bad_example(mutex: &Mutex<Vec<u32>>) {
    let guard = mutex.lock().unwrap();  // 阻塞获取锁
    // 警告：持有守卫时跨越 .await
    some_async_function().await;         // 编译器警告！
    drop(guard);
}

问题分析：

  • std::sync::Mutex::lock() 是阻塞式的，会阻塞当前线程
  • 如果在持有 MutexGuard 时调用 .await，守卫的生命周期可能跨越 .await 点
  • 当任务被挂起时，锁仍被持有，但其他任务可能被调度到同一个线程上运行
  • 如果这些任务也尝试获取同一个锁，由于锁还被上一个任务持有，就会导致死锁

1.2 Tokio Mutex 的解决方案

// 使用 tokio::sync::Mutex 的正确示例
use tokio::sync::Mutex;

async fn good_example(mutex: &Mutex<Vec<u32>>) {
    let mut guard = mutex.lock().await;  // 异步获取锁
    // 安全：可以跨越 .await 点
    some_async_function().await;
    drop(guard);
}

关键优势：

  • lock() 返回 Future，不会阻塞线程
  • Guard 可以安全地跨越 .await 点
  • 如果锁不可用，任务会让出执行权，而不是阻塞线程
  • 被挂起的任务会被放入等待队列，锁可用时会被自动唤醒
  • 在该场景下，对比std的Mutex，tokio的Mutex可以有效的利用CPU资源

二、核心设计原理

2.1 基于信号量的实现

Tokio Mutex 的核心是一个计数为 1 的信号量（Semaphore）：

classDiagram
    class Mutex~T~ {
        +Semaphore s
        +UnsafeCell~T~ c
    }
    class Semaphore {
        +AtomicUsize permits (1 → 锁可用; 0 → 锁被占用)
        +LinkedList~Waiter~ waiters
    }
    class UnsafeCell~T~ {
        +T value (提供内部可变性; 只有引用时，可以修改内部数据)
    }
    class Waiter {
        +Option~Waker~ waker
        +AtomicUsize state
    }

    Mutex~T~ --> Semaphore : s (信号量，管理锁的获取/释放)
    Mutex~T~ --> UnsafeCell~T~ : c (数据)
    Semaphore --> Waiter : 管理等待队列

为什么使用信号量？

  1. FIFO 公平性：信号量维护一个先进先出的等待队列
  2. 异步友好：内置 Waker 机制，支持任务挂起和唤醒

2.2 锁获取流程

获取锁的核心逻辑（伪代码）：

async fn lock(&self) -> MutexGuard<'_, T> {
    // 步骤 1: 尝试从信号量获取许可
    if self.semaphore.acquire(1) {
        // 成功：permits 从 1 变为 0，立即获得锁
        return MutexGuard { lock: self };
    }

    // 失败：permits 已经是 0，需要等待
    // 步骤 2: 创建等待节点，保存当前任务的 Waker
    let waiter = Waiter::new(current_task_waker());

    // 步骤 3: 加入 FIFO 等待队列
    self.semaphore.waiters.push_back(waiter);

    // 步骤 4: 返回 Pending，挂起当前任务
    // 任务会被放入运行时的等待队列，线程可以执行其他任务
    Pending.yield_now().await;
}

释放锁的流程：

释放锁时，会唤醒等待的任务获取锁。

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        // 步骤 1: 释放信号量许可
        self.lock.semaphore.release(1);  // permits: 0 → 1

        // 步骤 2: 检查等待队列
        if let Some(waiter) = self.lock.semaphore.waiters.pop_front() {
            // 步骤 3: 唤醒队列中的第一个等待者
            // permits: 1 → 0（被唤醒的等待者获得锁）
            waiter.waker.wake();
        }
    }
}

关键点说明：

  1. 原子操作：try_acquire 使用原子指令（如 fetch_sub），确保多线程安全
  2. FIFO 公平性：等待队列是先进先出，按请求顺序分配锁
  3. 零成本挂起：任务挂起时不阻塞 OS 线程，线程可以执行其他任务
  4. Waker 机制：每个等待者保存一个 Waker，锁释放时通过 wake() 唤醒对应的任务

2.3 三种 Guard 类型

Tokio Mutex 提供了三种不同的 Guard 类型，分别适用于不同的使用场景：

类型特性生命周期使用场景
MutexGuard<'a, T>借用 &Mutex'a一般场景，最常用
OwnedMutexGuard<T>持有 Arc<Mutex>'static需要跨任务传递 Guard
MappedMutexGuard<'a, T>映射到子字段'a只需访问部分数据

内存布局对比

classDiagram
    class Mutex~T~ {
        +Semaphore s
        +UnsafeCell~T~ c
    }

    class MutexGuard {
        + &'a Mutex~T~ lock
    }

    class OwnedMutexGuard {
        +Arc~Mutex~T~~ lock
    }

    class MappedMutexGuard {
        +&'a Semaphore s
        +*mut T data
    }

    MutexGuard --> Mutex~T~ : 借用 (生命周期 'a)
    OwnedMutexGuard --> Mutex~T~ : 拥有 (Arc, 生命周期 'static)
    MappedMutexGuard --> T : 限制访问范围，指向子字段 (生命周期 'a)

2.4 MutexGuard 与 OwnedMutexGuard的详细对比

特性MutexGuardOwnedMutexGuard
持有方式借用 &Mutex持有 Arc<Mutex>
生命周期'a (受 Mutex 限制)'static (无限制)
创建方法mutex.lock()mutex.clone().lock_owned()
Send✅ 是✅ 是
存储需求Mutex 必须比 Guard 长活无限制
跨线程受生命周期限制完全自由
典型场景临时访问、作用域内存储、传递、跨任务

MutexGuard - 借用模式

pub struct MutexGuard<'a, T: ?Sized> {
    lock: &'a Mutex<T>,  // ← 借用引用
    //     ^^^
    //     这是关键！Guard 持有的是 Mutex 的引用
}

// 创建
async fn example1() {
    let mutex = Mutex::new(42);
    
    {
        let guard = mutex.lock().await;
        //        ^^^^^ 借用 &mutex
        *guard = 100;
    } // guard 释放，mutex 仍然存在
}

// 生命周期限制
async fn example2() -> MutexGuard<'static, i32> {  // ❌ 编译错误！
    let mutex = Mutex::new(42);
    mutex.lock().await  // 返回的生命周期不够长
}
  • 调用lock后返回的guard是MutexGuard类型，guard修改时，是通过MutexGuard类型的deref_mut等方法操作了内部数据实现的。
  • 当 } 结束后，guard的生命周期释放，这时会自动调用MutexGuard类型的drop方法实现释放锁。
  • 整个过程看似，我们都在操作带锁的数据，实际上，我们操作的都是包含锁的MutexGuard类型。

OwnedMutexGuard - 拥有模式

pub struct OwnedMutexGuard<T: ?Sized> {
    lock: Arc<Mutex<T>>,  // ← 拥有 Arc
    //     ^^^
    //     这是关键！Guard 持有的是 Mutex 的所有权
}

// 创建
async fn example3() {
    let mutex = Arc::new(Mutex::new(42));
    
    {
        let guard = mutex.clone().lock_owned().await;
        //        ^^^^^^ clone Arc
        *guard = 100;
    } // guard 释放，但 Arc 引用可能还在
}

// 无生命周期限制
async fn example4() -> OwnedMutexGuard<i32> {  // ✅ 可以！
    let mutex = Arc::new(Mutex::new(42));
    mutex.clone().lock_owned().await
}

应用场景对比

场景 1: 临时访问（MutexGuard）
async fn temp_access(mutex: &Mutex<i32>) -> i32 {
    // 场景：只需要临时访问数据
    let guard = mutex.lock().await;
    let value = *guard;
    drop(guard);  // 显式释放
    value
}
// MutexGuard 最适合这种场景
场景 2: 存储 Guard（需要 OwnedMutexGuard）
struct MyServer {
    // ❌ MutexGuard 不能存储！
    // guard: MutexGuard<'a, i32>,  // 生命周期 'a 无法确定
    
    // ✅ OwnedMutexGuard 可以存储！
    guard: OwnedMutexGuard<i32>,
}

impl MyServer {
    async fn new(mutex: Arc<Mutex<i32>>) -> Self {
        let guard = mutex.clone().lock_owned().await;
        //                   ^^^^^^ 使用 lock_owned
        Self { guard }
    }
}
场景 3: 跨任务传递（需要 OwnedMutexGuard）
async fn cross_task() {
    let mutex = Arc::new(Mutex::new(42));
    
    // ✅ OwnedMutexGuard 可以跨任务
    let mutex2 = mutex.clone();
    tokio::spawn(async move {
        let guard = mutex2.lock_owned().await;
        // Guard 可以在这个新任务中使用
        *guard = 100;
    });
}
场景 4: 返回 Guard（需要 OwnedMutexGuard）
// ❌ MutexGuard - 无法返回
async fn get_guard_bad(mutex: &Mutex<i32>) -> MutexGuard<i32> {
    mutex.lock().await  // 编译错误：生命周期不够长
}

// ✅ OwnedMutexGuard - 可以返回
async fn get_guard_good(mutex: Arc<Mutex<i32>>) -> OwnedMutexGuard<i32> {
    mutex.clone().lock_owned().await  // 'static 生命周期
}

三、与 std::sync::Mutex 的本质区别

3.1 对比表

特性std::sync::Mutextokio::sync::Mutex
获取锁方式阻塞式 lock()异步式 lock().await
线程阻塞会阻塞 OS 线程不会阻塞，任务让出
跨 .await 持有❌ 不安全✅ 安全
性能开销低（仅原子操作）高（信号量 + Waker）
使用场景同步代码、数据保护异步代码、IO 资源

3.2 深入理解：为什么会有性能差异？

std::sync::Mutex：

// 伪代码
fn lock(&self) -> LockGuard<T> {
    loop {
        if self.compare_and_swap(0, 1) == 0 {
            return Guard;  // 获得锁
        }
        // 使用 futex 等待，阻塞 OS 线程
        syscall(futex_wait, &self.state);
    }
}

tokio::sync::Mutex：

// 伪代码
async fn lock(&self) -> MutexGuard<T> {
    if self.semaphore.try_acquire(1) {
        return Guard;  // 获得锁
    }
    // 创建 Waiter，加入队列
    let waiter = Waiter::new(current_task_waker());
    self.queue.push_back(waiter);
    Pending // 挂起任务，不阻塞线程
}

关键点：

  • std Mutex 阻塞 OS 线程，线程无法执行其他任务
  • Tokio Mutex 挂起任务，线程可以执行其他任务
  • Tokio Mutex 额外维护了等待队列和 Waker，所以开销更大

四、应用场景与实践

4.1 何时使用 Tokio Mutex？

✅ 适合场景

  1. IO 资源共享
use tokio::sync::Mutex;
use tokio::net::TcpStream;

struct DatabaseConnection {
    stream: TcpStream,
    // ...
}

async fn query(db: &Mutex<DatabaseConnection>, sql: &str) {
    let mut conn = db.lock().await;
    // 持有锁期间可能跨越多个 .await
    conn.stream.write_all(sql.as_bytes()).await?;
    let response = conn.stream.read_buf(&mut buf).await?;
    Ok(())
}
  1. 需要跨 .await 持有锁
async fn process_with_lock(mutex: &Mutex<Data>) {
    let mut data = mutex.lock().await;
    // 第一个异步操作
    async_step1(&mut data).await;
    // 第二个异步操作（仍然持有锁）
    async_step2(&mut data).await;
}

❌ 不适合场景

  1. 纯数据保护（优先用 std Mutex）
// 推荐：使用 std::sync::Mutex
use std::sync::Mutex;

struct Counter {
    count: Mutex<i32>,
}

impl Counter {
    fn increment(&self) {
        let mut count = self.count.lock().unwrap();
        *count += 1;
        // 纯内存操作，无需异步
    }
}
  1. 短时间锁定（优先用 std Mutex）
// 如果只是快速读写，不需要跨越 .await
// std Mutex 性能更好
let value = std_mutex.lock().unwrap();
let result = *value * 2;
drop(value);

4.2 最佳实践

实践 1：限制锁的持有时间

// ❌ 不推荐：持有锁时间过长
async fn bad(mutex: &Mutex<Data>) {
    let mut data = mutex.lock().await;
    // 执行耗时操作
    heavy_computation(&data).await;  // 锁被持有
    another_async_operation().await; // 锁仍被持有
}

// ✅ 推荐：最小化锁持有时间
async fn good(mutex: &Mutex<Data>) {
    // 1. 先获取需要的数据
    let (key, value) = {
        let data = mutex.lock().await;
        (data.key.clone(), data.value.clone())
    }; // 锁在这里释放

    // 2. 执行耗时操作（不持有锁）
    let result = heavy_computation(&key, &value).await;

    // 3. 最后再获取锁更新数据
    let mut data = mutex.lock().await;
    data.result = result;
}

实践 2：使用消息传递模式

对于 IO 资源，更好的模式是使用消息传递：

use tokio::sync::mpsc;

struct DatabaseManager {
    rx: mpsc::Receiver<Query>,
    connection: DatabaseConnection,
}

impl DatabaseManager {
    async fn run(mut self) {
        while let Some(query) = self.rx.recv().await {
            let result = self.execute_query(&query).await;
            query.response.send(result).ok();
        }
    }
}

// 使用者通过通道发送请求，无需直接加锁
manager.tx.send(Query::new("SELECT * FROM users")).await?;
let result = response.recv().await?;

优势：

  • 无需直接管理锁
  • 自然避免死锁
  • 更好的并发性（管理者可以内部优化）

实践 3：使用 OwnedMutexGuard 跨任务传递

use tokio::sync::{Mutex, mpsc};
use std::sync::Arc;

#[tokio::main]
async fn main() {
    let mutex = Arc::new(Mutex::new(Vec::new()));
    let (tx, mut rx) = mpsc::channel(100);

    // 生产者任务
    let producer = tokio::spawn(async move {
        for i in 0..10 {
            // 获取拥有所有权的 Guard
            let guard = mutex.clone().lock_owned().await;
            tx.send(guard).await.unwrap();
        }
    });

    // 消费者任务
    let consumer = tokio::spawn(async move {
        while let Some(mut vec) = rx.recv().await {
            vec.push(42);
            // Guard 在这里 drop，自动释放锁
        }
    });

    producer.await.unwrap();
    consumer.await.unwrap();
}

实践 4：使用 MappedMutexGuard 限制访问

use tokio::sync::{Mutex, MutexGuard};

struct Config {
    sensitive_data: String,
    public_setting: u32,
}

// 只暴露公共字段给调用者
fn get_public_setting(
    config: &Mutex<Config>
) -> impl Future<Output = MappedMutexGuard<'_, u32>> {
    let guard = config.lock();
    async move {
        MutexGuard::map(guard.await, |config| &mut config.public_setting)
    }
}

// 调用者只能访问 public_setting，无法访问敏感数据
async fn update_setting(config: &Mutex<Config>) {
    let mut setting = get_public_setting(config).await;
    *setting += 1;
    // setting.sensitive_data // 编译错误！无法访问
}

五、常见陷阱与解决方案

5.1 死锁风险

虽然 Tokio Mutex 比 std Mutex 更安全，但仍然可能死锁：

// ❌ 死锁：两个锁以不同顺序获取
async fn deadlock(mutex1: &Mutex<()>, mutex2: &Mutex<()>) {
    let lock1 = mutex1.lock().await;
    // 如果另一个任务先获取了 mutex2，就会死锁
    let lock2 = mutex2.lock().await;  // 可能永远阻塞
}

// ✅ 解决：按固定顺序获取锁
async fn no_deadlock(mutex1: &Mutex<()>, mutex2: &Mutex<()>) {
    // 始终先获取 mutex1，再获取 mutex2
    let lock1 = mutex1.lock().await;
    let lock2 = mutex2.lock().await;
}

5.2 性能陷阱

// ❌ 不推荐：在循环中频繁获取锁
async fn bad_loop(mutex: &Mutex<Vec<i32>>) {
    for i in 0..1000 {
        let mut data = mutex.lock().await;
        data.push(i);
    }  // 1000 次异步锁操作
}

// ✅ 推荐：批量操作
async fn good_loop(mutex: &Mutex<Vec<i32>>) {
    let mut new_items = Vec::new();
    for i in 0..1000 {
        new_items.push(i);
    }
    // 只获取一次锁
    let mut data = mutex.lock().await;
    data.append(&mut new_items);
}

六、性能对比与选择建议

6.1 性能测试

以下是简单的性能对比（仅供参考）：

操作std MutexTokio Mutex性能差异
无竞争 lock/unlock~50ns~200ns4x 慢
单线程简单操作~100ns~400ns4x 慢
多线程高并发~500ns~800ns1.6x 慢

结论：Tokio Mutex 有明显的性能开销，但在异步场景下是必要的。

6.2 选择决策

  • 需要跨越 .await 持有锁，使用 tokio::sync::Mutex
  • IO 资源（如数据库连接），考虑消息传递模式 tokio::sync::mpsc
  • 其他简单的数据保护，使用 std::sync::Mutex（性能最优）

七、总结

关键要点

  1. 使用场景：仅在需要跨越 .await 持有锁时使用
  2. 性能考虑：有明显的性能开销，优先考虑 std Mutex 或消息传递
  3. 三种 Guard：根据需求选择合适的 Guard 类型
  4. 最佳实践：最小化锁持有时间，优先使用消息传递模式

学习路径

  1. 理解 Rust 异步编程基础（Future、.await、运行时）
  2. 掌握 Tokio 的任务调度和 Waker 机制
  3. 学习信号量原理（tokio/src/sync/batch_semaphore.rs）
  4. 深入 Mutex 实现（tokio/src/sync/mutex.rs）
  5. 可以结合 AI 的输出理解难点

相关资源

avatar
文章
阅读
粉丝