os/sync/
raw_spin_lock.rs

1use crate::sync::intr_guard::IntrGuard;
2use core::{
3    hint,
4    sync::atomic::{AtomicBool, Ordering},
5};
6
7/// 自旋锁结构体，提供互斥访问临界区的能力。
8/// 基于原子操作实现自旋锁机制，结合 IntrGuard 实现中断保护。
9/// 不可重入 (即不能嵌套调用 RawSpinLock::lock())。
10/// 使用示例：
11/// ```ignore
12/// let lock = RawSpinLock::new();
13/// {
14///   let guard = lock.lock(); // 获取锁，禁用中断
15///   // 临界区代码
16/// } // 离开作用域，自动释放锁并恢复中断状态
17/// ```
18#[derive(Debug)]
19pub struct RawSpinLock {
20    lock: AtomicBool,
21}
22
23impl RawSpinLock {
24    pub const fn new() -> Self {
25        RawSpinLock {
26            lock: AtomicBool::new(false),
27        }
28    }
29
30    /// 尝试获取自旋锁，并返回一个 RAII 保护器。
31    ///
32    /// 内部原子地获取锁，并在当前 CPU 禁用本地中断。
33    pub fn lock(&self) -> RawSpinLockGuard<'_> {
34        let guard = IntrGuard::new();
35
36        while self
37            .lock
38            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
39            .is_err()
40        {
41            hint::spin_loop();
42        }
43
44        RawSpinLockGuard {
45            lock: self,
46            intr_guard: guard,
47        }
48    }
49
50    /// 尝试获取自旋锁，如果成功则返回 RAII 保护器，否则返回 None。
51    ///
52    /// 内部原子地尝试获取锁，并在当前 CPU 禁用本地中断。
53    /// 如果获取失败，会立即恢复中断状态（通过 Drop IntrGuard）。
54    pub fn try_lock(&self) -> Option<RawSpinLockGuard<'_>> {
55        let guard = IntrGuard::new();
56
57        if self
58            .lock
59            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
60            .is_ok()
61        {
62            Some(RawSpinLockGuard {
63                lock: self,
64                intr_guard: guard,
65            })
66        } else {
67            None
68        }
69    }
70
71    /// 仅释放锁标志。
72    fn unlock(&self) {
73        self.lock.store(false, Ordering::Release);
74    }
75
76    /// 检查锁是否被占用 (仅用于调试/测试)
77    /// 返回值：锁是否被占用
78    #[cfg(test)]
79    pub fn is_locked(&self) -> bool {
80        self.lock.load(Ordering::Relaxed)
81    }
82}
83
84/// 自动释放自旋锁和恢复中断状态的 RAII 结构体
85pub struct RawSpinLockGuard<'a> {
86    lock: &'a RawSpinLock,
87    intr_guard: IntrGuard,
88}
89
90use core::ops::Drop;
91
92impl Drop for RawSpinLockGuard<'_> {
93    /// 退出作用域时自动执行，顺序如下：
94    /// 1. 释放自旋锁标志。
95    /// 2. IntrGuard 被 Drop，恢复中断状态。
96    fn drop(&mut self) {
97        self.lock.unlock();
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::{
105        arch::intr::{are_interrupts_enabled, read_and_disable_interrupts, restore_interrupts},
106        kassert, test_case,
107    };
108
109    // 模拟一个共享资源，必须用 RawSpinLock 保护
110    static COUNTER: AtomicBool = AtomicBool::new(false);
111
112    // 测试锁的初始化状态和基本锁定/解锁功能
113    test_case!(test_raw_spin_lock_basic_lock_unlock, {
114        let lock = RawSpinLock::new();
115        kassert!(!lock.is_locked());
116
117        let guard = lock.lock();
118        kassert!(lock.is_locked());
119
120        // 手动释放 (Drop)
121        drop(guard);
122        kassert!(!lock.is_locked());
123    });
124
125    // 测试 RAII 行为 (自动释放)
126    test_case!(test_raw_spin_lock_raii_release, {
127        let lock = RawSpinLock::new();
128
129        {
130            let _guard = lock.lock();
131            kassert!(lock.is_locked());
132        } // <- _guard 在此离开作用域，Drop 被自动调用
133
134        kassert!(!lock.is_locked());
135    });
136
137    // 测试互斥性 (只能获取一次)
138    test_case!(test_raw_spin_lock_mutual_exclusion, {
139        let lock = RawSpinLock::new();
140
141        let guard1 = lock.lock();
142        kassert!(lock.is_locked());
143
144        // 尝试第二次获取 (理论上会进入无限自旋，但测试中我们只检查状态)
145        // NOTE: 在实际运行环境中，第二次调用会死循环，测试环境通常需要模拟并发
146        // 在这里我们依赖测试框架的单线程执行来简单检查 is_locked 状态
147
148        // 模拟多线程获取失败的场景：
149        let second_lock_failed;
150
151        // 临时释放，让第二次获取成功
152        drop(guard1);
153
154        let guard2 = lock.lock();
155        if lock.is_locked() {
156            // 第二次获取成功
157            second_lock_failed = false;
158        } else {
159            second_lock_failed = true;
160        }
161
162        kassert!(!second_lock_failed);
163        drop(guard2);
164    });
165
166    // -----------------------------------------------------------
167    // 中断保护测试
168    // -----------------------------------------------------------
169
170    // 测试 lock() 是否禁用了中断
171    test_case!(test_interrupt_disable, {
172        // 1. 确保中断最初是启用的
173        let initial_flags = unsafe { read_and_disable_interrupts() };
174        unsafe { restore_interrupts(initial_flags | (1 << 1)) }; // 确保 SIE 启用
175        kassert!(are_interrupts_enabled());
176
177        let lock = RawSpinLock::new();
178        let guard = lock.lock();
179
180        // 2. 检查中断是否被禁用
181        kassert!(!are_interrupts_enabled());
182        kassert!(guard.intr_guard.was_enabled());
183
184        // 3. 检查 Drop 后中断是否恢复
185        drop(guard);
186        kassert!(are_interrupts_enabled());
187
188        // 恢复测试前的环境
189        unsafe { restore_interrupts(initial_flags) };
190    });
191}
os/sync/raw_spin_lock.rs

os/sync/
raw_spin_lock.rs
