xref: /DragonOS/kernel/src/libs/rwlock.rs (revision 8cb2e9b344230227fe5f3ab3ebeb2522f1c5e289)
1 #![allow(dead_code)]
2 use core::{
3     cell::UnsafeCell,
4     hint::spin_loop,
5     mem::{self, ManuallyDrop},
6     ops::{Deref, DerefMut},
7     sync::atomic::{AtomicU32, Ordering},
8 };
9 
10 use system_error::SystemError;
11 
12 use crate::{
13     arch::CurrentIrqArch,
14     exception::{InterruptArch, IrqFlagsGuard},
15     process::ProcessManager,
16 };
17 
18 ///RwLock读写锁
19 
20 /// @brief READER位占据从右往左数第三个比特位
21 const READER: u32 = 1 << 2;
22 
23 /// @brief UPGRADED位占据从右到左数第二个比特位
24 const UPGRADED: u32 = 1 << 1;
25 
26 /// @brief WRITER位占据最右边的比特位
27 const WRITER: u32 = 1;
28 
29 const READER_BIT: u32 = 2;
30 
31 /// @brief 读写锁的基本数据结构
32 /// @param lock 32位原子变量,最右边的两位从左到右分别是UPGRADED,WRITER (标志位)
33 ///             剩下的bit位存储READER数量(除了MSB)
34 ///             对于标志位,0代表无, 1代表有
35 ///             对于剩下的比特位表征READER的数量的多少
36 ///             lock的MSB必须为0,否则溢出
37 #[derive(Debug)]
38 pub struct RwLock<T> {
39     lock: AtomicU32,
40     data: UnsafeCell<T>,
41 }
42 
43 /// @brief  READER守卫的数据结构
44 /// @param lock 是对RwLock的lock属性值的只读引用
45 pub struct RwLockReadGuard<'a, T: 'a> {
46     data: *const T,
47     lock: &'a AtomicU32,
48     irq_guard: Option<IrqFlagsGuard>,
49 }
50 
51 /// @brief UPGRADED是介于READER和WRITER之间的一种锁,它可以升级为WRITER,
52 ///        UPGRADED守卫的数据结构,注册UPGRADED锁只需要查看UPGRADED和WRITER的比特位
53 ///        但是当UPGRADED守卫注册后,不允许有新的读者锁注册
54 /// @param inner    是对RwLock数据结构的只读引用
55 pub struct RwLockUpgradableGuard<'a, T: 'a> {
56     data: *const T,
57     inner: &'a RwLock<T>,
58     irq_guard: Option<IrqFlagsGuard>,
59 }
60 
61 /// @brief WRITER守卫的数据结构
62 /// @param data     RwLock的data的可变引用
63 /// @param inner    是对RwLock数据结构的只读引用
64 pub struct RwLockWriteGuard<'a, T: 'a> {
65     data: *mut T,
66     inner: &'a RwLock<T>,
67     irq_guard: Option<IrqFlagsGuard>,
68 }
69 
70 unsafe impl<T: Send> Send for RwLock<T> {}
71 unsafe impl<T: Send + Sync> Sync for RwLock<T> {}
72 
73 /// @brief RwLock的API
74 impl<T> RwLock<T> {
75     #[inline]
76     /// @brief  RwLock的初始化
77     pub const fn new(data: T) -> Self {
78         return RwLock {
79             lock: AtomicU32::new(0),
80             data: UnsafeCell::new(data),
81         };
82     }
83 
84     #[allow(dead_code)]
85     #[inline]
86     /// @brief 将读写锁的皮扒掉,返回内在的data,返回的是一个真身而非引用
87     pub fn into_inner(self) -> T {
88         let RwLock { data, .. } = self;
89         return data.into_inner();
90     }
91 
92     #[allow(dead_code)]
93     #[inline]
94     /// @brief 返回data的raw pointer,
95     /// unsafe
96     pub fn as_mut_ptr(&self) -> *mut T {
97         return self.data.get();
98     }
99 
100     #[allow(dead_code)]
101     #[inline]
102     /// @brief 获取实时的读者数并尝试加1,如果增加值成功则返回增加1后的读者数,否则panic
103     fn current_reader(&self) -> Result<u32, SystemError> {
104         const MAX_READERS: u32 = core::u32::MAX >> READER_BIT >> 1; //右移3位
105 
106         let value = self.lock.fetch_add(READER, Ordering::Acquire);
107         //value二进制形式的MSB不能为1, 否则导致溢出
108 
109         if value > MAX_READERS << READER_BIT {
110             self.lock.fetch_sub(READER, Ordering::Release);
111             //panic!("Too many lock readers, cannot safely proceed");
112             return Err(SystemError::EOVERFLOW);
113         } else {
114             return Ok(value);
115         }
116     }
117 
118     #[allow(dead_code)]
119     #[inline]
120     /// @brief 尝试获取READER守卫
121     pub fn try_read(&self) -> Option<RwLockReadGuard<T>> {
122         ProcessManager::preempt_disable();
123         let r = self.inner_try_read();
124         if r.is_none() {
125             ProcessManager::preempt_enable();
126         }
127         return r;
128     }
129 
130     fn inner_try_read(&self) -> Option<RwLockReadGuard<T>> {
131         let reader_value = self.current_reader();
132         //得到自增后的reader_value, 包括了尝试获得READER守卫的进程
133         let value;
134 
135         if reader_value.is_err() {
136             return None; //获取失败
137         } else {
138             value = reader_value.unwrap();
139         }
140 
141         //判断有没有writer和upgrader
142         //注意, 若upgrader存在,已经存在的读者继续占有锁,但新读者不允许获得锁
143         if value & (WRITER | UPGRADED) != 0 {
144             self.lock.fetch_sub(READER, Ordering::Release);
145             return None;
146         } else {
147             return Some(RwLockReadGuard {
148                 data: unsafe { &*self.data.get() },
149                 lock: &self.lock,
150                 irq_guard: None,
151             });
152         }
153     }
154 
155     #[allow(dead_code)]
156     #[inline]
157     /// @brief 获得READER的守卫
158     pub fn read(&self) -> RwLockReadGuard<T> {
159         loop {
160             match self.try_read() {
161                 Some(guard) => return guard,
162                 None => spin_loop(),
163             }
164         } //忙等待
165     }
166 
167     /// 关中断并获取读者守卫
168     pub fn read_irqsave(&self) -> RwLockReadGuard<T> {
169         loop {
170             let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
171             match self.try_read() {
172                 Some(mut guard) => {
173                     guard.irq_guard = Some(irq_guard);
174                     return guard;
175                 }
176                 None => spin_loop(),
177             }
178         }
179     }
180 
181     /// 尝试关闭中断并获取读者守卫
182     pub fn try_read_irqsave(&self) -> Option<RwLockReadGuard<T>> {
183         let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
184         if let Some(mut guard) = self.try_read() {
185             guard.irq_guard = Some(irq_guard);
186             return Some(guard);
187         } else {
188             return None;
189         }
190     }
191 
192     #[allow(dead_code)]
193     #[inline]
194     /// @brief 获取读者+UPGRADER的数量, 不能保证能否获得同步值
195     pub fn reader_count(&self) -> u32 {
196         let state = self.lock.load(Ordering::Relaxed);
197         return state / READER + (state & UPGRADED) / UPGRADED;
198     }
199 
200     #[allow(dead_code)]
201     #[inline]
202     /// @brief 获取写者数量,不能保证能否获得同步值
203     pub fn writer_count(&self) -> u32 {
204         return (self.lock.load(Ordering::Relaxed) & WRITER) / WRITER;
205     }
206 
207     #[cfg(any(target_arch = "x86_64", target_arch = "riscv64"))]
208     #[allow(dead_code)]
209     #[inline]
210     /// @brief 尝试获得WRITER守卫
211     pub fn try_write(&self) -> Option<RwLockWriteGuard<T>> {
212         ProcessManager::preempt_disable();
213         let r = self.inner_try_write();
214         if r.is_none() {
215             ProcessManager::preempt_enable();
216         }
217 
218         return r;
219     } //当架构为arm时,有些代码需要作出调整compare_exchange=>compare_exchange_weak
220 
221     #[cfg(any(target_arch = "x86_64", target_arch = "riscv64"))]
222     #[allow(dead_code)]
223     #[inline]
224     pub fn try_write_irqsave(&self) -> Option<RwLockWriteGuard<T>> {
225         ProcessManager::preempt_disable();
226         let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
227         let r = self.inner_try_write().map(|mut g| {
228             g.irq_guard = Some(irq_guard);
229             g
230         });
231         if r.is_none() {
232             ProcessManager::preempt_enable();
233         }
234 
235         return r;
236     }
237 
238     #[cfg(any(target_arch = "x86_64", target_arch = "riscv64"))]
239     #[allow(dead_code)]
240     fn inner_try_write(&self) -> Option<RwLockWriteGuard<T>> {
241         let res: bool = self
242             .lock
243             .compare_exchange(0, WRITER, Ordering::Acquire, Ordering::Relaxed)
244             .is_ok();
245         //只有lock大小为0的时候能获得写者守卫
246         if res {
247             return Some(RwLockWriteGuard {
248                 data: unsafe { &mut *self.data.get() },
249                 inner: self,
250                 irq_guard: None,
251             });
252         } else {
253             return None;
254         }
255     }
256 
257     #[allow(dead_code)]
258     #[inline]
259     /// @brief 获得WRITER守卫
260     pub fn write(&self) -> RwLockWriteGuard<T> {
261         loop {
262             match self.try_write() {
263                 Some(guard) => return guard,
264                 None => spin_loop(),
265             }
266         }
267     }
268 
269     #[allow(dead_code)]
270     #[inline]
271     /// @brief 获取WRITER守卫并关中断
272     pub fn write_irqsave(&self) -> RwLockWriteGuard<T> {
273         loop {
274             let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
275             match self.try_write() {
276                 Some(mut guard) => {
277                     guard.irq_guard = Some(irq_guard);
278                     return guard;
279                 }
280                 None => spin_loop(),
281             }
282         }
283     }
284 
285     #[allow(dead_code)]
286     #[inline]
287     /// @brief 尝试获得UPGRADER守卫
288     pub fn try_upgradeable_read(&self) -> Option<RwLockUpgradableGuard<T>> {
289         ProcessManager::preempt_disable();
290         let r = self.inner_try_upgradeable_read();
291         if r.is_none() {
292             ProcessManager::preempt_enable();
293         }
294 
295         return r;
296     }
297 
298     #[allow(dead_code)]
299     pub fn try_upgradeable_read_irqsave(&self) -> Option<RwLockUpgradableGuard<T>> {
300         let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
301         ProcessManager::preempt_disable();
302         let mut r = self.inner_try_upgradeable_read();
303         if r.is_none() {
304             ProcessManager::preempt_enable();
305         } else {
306             r.as_mut().unwrap().irq_guard = Some(irq_guard);
307         }
308 
309         return r;
310     }
311 
312     fn inner_try_upgradeable_read(&self) -> Option<RwLockUpgradableGuard<T>> {
313         // 获得UPGRADER守卫不需要查看读者位
314         // 如果获得读者锁失败,不需要撤回fetch_or的原子操作
315         if self.lock.fetch_or(UPGRADED, Ordering::Acquire) & (WRITER | UPGRADED) == 0 {
316             return Some(RwLockUpgradableGuard {
317                 inner: self,
318                 data: unsafe { &mut *self.data.get() },
319                 irq_guard: None,
320             });
321         } else {
322             return None;
323         }
324     }
325 
326     #[allow(dead_code)]
327     #[inline]
328     /// @brief 获得UPGRADER守卫
329     pub fn upgradeable_read(&self) -> RwLockUpgradableGuard<T> {
330         loop {
331             match self.try_upgradeable_read() {
332                 Some(guard) => return guard,
333                 None => spin_loop(),
334             }
335         }
336     }
337 
338     #[inline]
339     /// @brief 获得UPGRADER守卫
340     pub fn upgradeable_read_irqsave(&self) -> RwLockUpgradableGuard<T> {
341         loop {
342             let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };
343             match self.try_upgradeable_read() {
344                 Some(mut guard) => {
345                     guard.irq_guard = Some(irq_guard);
346                     return guard;
347                 }
348                 None => spin_loop(),
349             }
350         }
351     }
352 
353     #[allow(dead_code)]
354     #[inline]
355     //extremely unsafe behavior
356     /// @brief 强制减少READER数
357     pub unsafe fn force_read_decrement(&self) {
358         debug_assert!(self.lock.load(Ordering::Relaxed) & !WRITER > 0);
359         self.lock.fetch_sub(READER, Ordering::Release);
360     }
361 
362     #[allow(dead_code)]
363     #[inline]
364     //extremely unsafe behavior
365     /// @brief 强制给WRITER解锁
366     pub unsafe fn force_write_unlock(&self) {
367         debug_assert_eq!(self.lock.load(Ordering::Relaxed) & !(WRITER | UPGRADED), 0);
368         self.lock.fetch_and(!(WRITER | UPGRADED), Ordering::Release);
369     }
370 
371     #[allow(dead_code)]
372     pub unsafe fn get_mut(&mut self) -> &mut T {
373         unsafe { &mut *self.data.get() }
374     }
375 }
376 
377 impl<T: Default> Default for RwLock<T> {
378     fn default() -> Self {
379         Self::new(Default::default())
380     }
381 }
382 
383 /// @brief 由原有的值创建新的锁
384 impl<T> From<T> for RwLock<T> {
385     fn from(data: T) -> Self {
386         return Self::new(data);
387     }
388 }
389 
390 impl<'rwlock, T> RwLockReadGuard<'rwlock, T> {
391     /// @brief 释放守卫,获得保护的值的不可变引用
392     ///
393     /// ## Safety
394     ///
395     /// 由于这样做可能导致守卫在另一个线程中被释放,从而导致pcb的preempt count不正确,
396     /// 因此必须小心的手动维护好preempt count。
397     ///
398     /// 并且,leak还可能导致锁的状态不正确。因此请仔细考虑是否真的需要使用这个函数。
399     #[allow(dead_code)]
400     #[inline]
401     pub unsafe fn leak(this: Self) -> &'rwlock T {
402         let this = ManuallyDrop::new(this);
403         return unsafe { &*this.data };
404     }
405 }
406 
407 impl<'rwlock, T> RwLockUpgradableGuard<'rwlock, T> {
408     #[allow(dead_code)]
409     #[inline]
410     /// @brief 尝试将UPGRADER守卫升级为WRITER守卫
411     pub fn try_upgrade(mut self) -> Result<RwLockWriteGuard<'rwlock, T>, Self> {
412         let res = self.inner.lock.compare_exchange(
413             UPGRADED,
414             WRITER,
415             Ordering::Acquire,
416             Ordering::Relaxed,
417         );
418         //当且仅当只有UPGRADED守卫时可以升级
419 
420         if res.is_ok() {
421             let inner = self.inner;
422             let irq_guard = self.irq_guard.take();
423             mem::forget(self);
424 
425             Ok(RwLockWriteGuard {
426                 data: unsafe { &mut *inner.data.get() },
427                 inner,
428                 irq_guard,
429             })
430         } else {
431             Err(self)
432         }
433     }
434 
435     #[allow(dead_code)]
436     #[inline]
437     /// @brief 将upgrader升级成writer
438     pub fn upgrade(mut self) -> RwLockWriteGuard<'rwlock, T> {
439         loop {
440             self = match self.try_upgrade() {
441                 Ok(writeguard) => return writeguard,
442                 Err(former) => former,
443             };
444 
445             spin_loop();
446         }
447     }
448 
449     #[allow(dead_code)]
450     #[inline]
451     /// @brief UPGRADER降级为READER
452     pub fn downgrade(mut self) -> RwLockReadGuard<'rwlock, T> {
453         while self.inner.current_reader().is_err() {
454             spin_loop();
455         }
456 
457         let inner: &RwLock<T> = self.inner;
458         let irq_guard = self.irq_guard.take();
459         // 自动移去UPGRADED比特位
460         mem::drop(self);
461 
462         RwLockReadGuard {
463             data: unsafe { &*inner.data.get() },
464             lock: &inner.lock,
465             irq_guard,
466         }
467     }
468 
469     #[allow(dead_code)]
470     #[inline]
471     /// @brief 返回内部数据的引用,消除守卫
472     ///
473     /// ## Safety
474     ///
475     /// 由于这样做可能导致守卫在另一个线程中被释放,从而导致pcb的preempt count不正确,
476     /// 因此必须小心的手动维护好preempt count。
477     ///
478     /// 并且,leak还可能导致锁的状态不正确。因此请仔细考虑是否真的需要使用这个函数。
479     pub unsafe fn leak(this: Self) -> &'rwlock T {
480         let this: ManuallyDrop<RwLockUpgradableGuard<'_, T>> = ManuallyDrop::new(this);
481 
482         unsafe { &*this.data }
483     }
484 }
485 
486 impl<'rwlock, T> RwLockWriteGuard<'rwlock, T> {
487     #[allow(dead_code)]
488     #[inline]
489     /// @brief 返回内部数据的引用,消除守卫
490     ///
491     /// ## Safety
492     ///
493     /// 由于这样做可能导致守卫在另一个线程中被释放,从而导致pcb的preempt count不正确,
494     /// 因此必须小心的手动维护好preempt count。
495     ///
496     /// 并且,leak还可能导致锁的状态不正确。因此请仔细考虑是否真的需要使用这个函数。
497     pub unsafe fn leak(this: Self) -> &'rwlock T {
498         let this = ManuallyDrop::new(this);
499 
500         return unsafe { &*this.data };
501     }
502 
503     #[allow(dead_code)]
504     #[inline]
505     /// @brief 将WRITER降级为READER
506     pub fn downgrade(mut self) -> RwLockReadGuard<'rwlock, T> {
507         while self.inner.current_reader().is_err() {
508             spin_loop();
509         }
510         //本质上来说绝对保证没有任何读者
511 
512         let inner = self.inner;
513         let irq_guard = self.irq_guard.take();
514         mem::drop(self);
515 
516         return RwLockReadGuard {
517             data: unsafe { &*inner.data.get() },
518             lock: &inner.lock,
519             irq_guard,
520         };
521     }
522 
523     #[allow(dead_code)]
524     #[inline]
525     /// @brief 将WRITER降级为UPGRADER
526     pub fn downgrade_to_upgradeable(mut self) -> RwLockUpgradableGuard<'rwlock, T> {
527         debug_assert_eq!(
528             self.inner.lock.load(Ordering::Acquire) & (WRITER | UPGRADED),
529             WRITER
530         );
531 
532         self.inner.lock.store(UPGRADED, Ordering::Release);
533 
534         let inner = self.inner;
535 
536         let irq_guard = self.irq_guard.take();
537         mem::forget(self);
538 
539         return RwLockUpgradableGuard {
540             inner,
541             data: unsafe { &*inner.data.get() },
542             irq_guard,
543         };
544     }
545 }
546 
547 impl<'rwlock, T> Deref for RwLockReadGuard<'rwlock, T> {
548     type Target = T;
549 
550     fn deref(&self) -> &Self::Target {
551         return unsafe { &*self.data };
552     }
553 }
554 
555 impl<'rwlock, T> Deref for RwLockUpgradableGuard<'rwlock, T> {
556     type Target = T;
557 
558     fn deref(&self) -> &Self::Target {
559         return unsafe { &*self.data };
560     }
561 }
562 
563 impl<'rwlock, T> Deref for RwLockWriteGuard<'rwlock, T> {
564     type Target = T;
565 
566     fn deref(&self) -> &Self::Target {
567         return unsafe { &*self.data };
568     }
569 }
570 
571 impl<'rwlock, T> DerefMut for RwLockWriteGuard<'rwlock, T> {
572     fn deref_mut(&mut self) -> &mut Self::Target {
573         return unsafe { &mut *self.data };
574     }
575 }
576 
577 impl<'rwlock, T> Drop for RwLockReadGuard<'rwlock, T> {
578     fn drop(&mut self) {
579         debug_assert!(self.lock.load(Ordering::Relaxed) & !(WRITER | UPGRADED) > 0);
580         self.lock.fetch_sub(READER, Ordering::Release);
581         ProcessManager::preempt_enable();
582     }
583 }
584 
585 impl<'rwlock, T> Drop for RwLockUpgradableGuard<'rwlock, T> {
586     fn drop(&mut self) {
587         debug_assert_eq!(
588             self.inner.lock.load(Ordering::Relaxed) & (WRITER | UPGRADED),
589             UPGRADED
590         );
591         self.inner.lock.fetch_sub(UPGRADED, Ordering::AcqRel);
592         ProcessManager::preempt_enable();
593         //这里为啥要AcqRel? Release应该就行了?
594     }
595 }
596 
597 impl<'rwlock, T> Drop for RwLockWriteGuard<'rwlock, T> {
598     fn drop(&mut self) {
599         debug_assert_eq!(self.inner.lock.load(Ordering::Relaxed) & WRITER, WRITER);
600         self.inner
601             .lock
602             .fetch_and(!(WRITER | UPGRADED), Ordering::Release);
603         self.irq_guard.take();
604         ProcessManager::preempt_enable();
605     }
606 }
607