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