os/sync/
mutex.rs

1#![allow(dead_code)]
2use core::cell::UnsafeCell;
3use core::sync::atomic::{AtomicBool, Ordering};
4
5use crate::kernel::{WaitQueue, current_cpu, yield_task};
6use crate::sync::SpinLock;
7use crate::sync::{raw_spin_lock::RawSpinLock, raw_spin_lock::RawSpinLockGuard};
8
9/// 互斥锁
10pub struct Mutex<T> {
11    locked: AtomicBool,
12    guard: RawSpinLock,
13    queue: SpinLock<WaitQueue>,
14    data: UnsafeCell<T>,
15}
16
17unsafe impl<T: Send> Send for Mutex<T> {}
18unsafe impl<T: Send> Sync for Mutex<T> {}
19
20impl<T> Mutex<T> {
21    pub fn new(data: T) -> Self {
22        Self {
23            locked: AtomicBool::new(false),
24            guard: RawSpinLock::new(),
25            queue: SpinLock::new(WaitQueue::new()),
26            data: UnsafeCell::new(data),
27        }
28    }
29
30    pub fn lock(&self) -> MutexGuard<'_, T> {
31        loop {
32            // 先自旋获取内部短锁
33            let spin = self.guard.lock();
34            // 尝试占用
35            if !self.locked.swap(true, Ordering::Acquire) {
36                // 成功,占用了
37                let data_ref = unsafe { &mut *self.data.get() };
38                return MutexGuard {
39                    mutex: self as *const _,
40                    _spin: spin,
41                    data: data_ref,
42                };
43            }
44            // 已被占用:把当前任务挂到等待队列
45            let current = current_cpu()
46                .lock()
47                .current_task
48                .as_ref()
49                .expect("mutex::lock: no current task")
50                .clone();
51            self.queue.lock().sleep(current);
52            // 释放短锁后让调度器切走
53            drop(spin);
54            yield_task();
55            // 醒来后重试
56        }
57    }
58}
59
60pub struct MutexGuard<'a, T> {
61    mutex: *const Mutex<T>,
62    _spin: RawSpinLockGuard<'a>, // 保持内部互斥到 drop
63    data: &'a mut T,
64}
65
66impl<T> core::ops::Deref for MutexGuard<'_, T> {
67    type Target = T;
68    fn deref(&self) -> &Self::Target {
69        self.data
70    }
71}
72
73impl<T> core::ops::DerefMut for MutexGuard<'_, T> {
74    fn deref_mut(&mut self) -> &mut Self::Target {
75        self.data
76    }
77}
78
79impl<T> Drop for MutexGuard<'_, T> {
80    fn drop(&mut self) {
81        // 仍持有 _spin,自然是互斥的
82        let m = unsafe { &*self.mutex };
83        m.locked.store(false, Ordering::Release);
84        m.queue.lock().wake_up_all();
85    }
86}