xref: /DragonOS/kernel/src/libs/rwlock.rs (revision d8ad0a5e7724469abd5cc3cf271993538878033e)
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     #[allow(dead_code)]
296     #[inline]
297     /// @brief 释放守卫,获得保护的值的不可变引用
298     pub fn leak(this: Self) -> &'rwlock T {
299         let Self { data, .. } = this;
300         return unsafe { &*data };
301     }
302 }
303 
304 impl<'rwlock, T> RwLockUpgradableGuard<'rwlock, T> {
305     #[allow(dead_code)]
306     #[inline]
307     /// @brief 尝试将UPGRADER守卫升级为WRITER守卫
308     pub fn try_upgrade(self) -> Result<RwLockWriteGuard<'rwlock, T>, Self> {
309         let res = self.inner.lock.compare_exchange(
310             UPGRADED,
311             WRITER,
312             Ordering::Acquire,
313             Ordering::Relaxed,
314         );
315         //当且仅当只有UPGRADED守卫时可以升级
316 
317         if res.is_ok() {
318             let inner = self.inner;
319 
320             mem::forget(self);
321 
322             Ok(RwLockWriteGuard {
323                 data: unsafe { &mut *inner.data.get() },
324                 inner,
325             })
326         } else {
327             Err(self)
328         }
329     }
330 
331     #[allow(dead_code)]
332     #[inline]
333     /// @brief 将upgrader升级成writer
334     pub fn upgrade(mut self) -> RwLockWriteGuard<'rwlock, T> {
335         loop {
336             self = match self.try_upgrade() {
337                 Ok(writeguard) => return writeguard,
338                 Err(former) => former,
339             };
340 
341             spin_loop();
342         }
343     }
344 
345     #[allow(dead_code)]
346     #[inline]
347     /// @brief UPGRADER降级为READER
348     pub fn downgrade(self) -> RwLockReadGuard<'rwlock, T> {
349         while self.inner.current_reader().is_err() {
350             spin_loop();
351         }
352 
353         let inner: &RwLock<T> = self.inner;
354 
355         // 自动移去UPGRADED比特位
356         mem::drop(self);
357 
358         RwLockReadGuard {
359             data: unsafe { &*inner.data.get() },
360             lock: &inner.lock,
361         }
362     }
363 
364     #[allow(dead_code)]
365     #[inline]
366     /// @brief 返回内部数据的引用,消除锁
367     pub fn leak(this: Self) -> &'rwlock T {
368         let this = ManuallyDrop::new(this);
369 
370         unsafe { &*this.data }
371     }
372 }
373 
374 impl<'rwlock, T> RwLockWriteGuard<'rwlock, T> {
375     #[allow(dead_code)]
376     #[inline]
377     /// @brief 返回内部数据的引用,消除锁
378     pub fn leak(this: Self) -> &'rwlock T {
379         let this = ManuallyDrop::new(this);
380 
381         return unsafe { &*this.data };
382     }
383 
384     #[allow(dead_code)]
385     #[inline]
386     /// @brief 将WRITER降级为READER
387     pub fn downgrade(self) -> RwLockReadGuard<'rwlock, T> {
388         while self.inner.current_reader().is_err() {
389             spin_loop();
390         }
391         //本质上来说绝对保证没有任何读者
392 
393         let inner = self.inner;
394 
395         mem::drop(self);
396 
397         return RwLockReadGuard {
398             data: unsafe { &*inner.data.get() },
399             lock: &inner.lock,
400         };
401     }
402 
403     #[allow(dead_code)]
404     #[inline]
405     /// @brief 将WRITER降级为UPGRADER
406     pub fn downgrade_to_upgradeable(self) -> RwLockUpgradableGuard<'rwlock, T> {
407         debug_assert_eq!(
408             self.inner.lock.load(Ordering::Acquire) & (WRITER | UPGRADED),
409             WRITER
410         );
411 
412         self.inner.lock.store(UPGRADED, Ordering::Release);
413 
414         let inner = self.inner;
415 
416         mem::forget(self);
417 
418         return RwLockUpgradableGuard {
419             inner,
420             data: unsafe { &*inner.data.get() },
421         };
422     }
423 }
424 
425 impl<'rwlock, T> Deref for RwLockReadGuard<'rwlock, T> {
426     type Target = T;
427 
428     fn deref(&self) -> &Self::Target {
429         return unsafe { &*self.data };
430     }
431 }
432 
433 impl<'rwlock, T> Deref for RwLockUpgradableGuard<'rwlock, T> {
434     type Target = T;
435 
436     fn deref(&self) -> &Self::Target {
437         return unsafe { &*self.data };
438     }
439 }
440 
441 impl<'rwlock, T> Deref for RwLockWriteGuard<'rwlock, T> {
442     type Target = T;
443 
444     fn deref(&self) -> &Self::Target {
445         return unsafe { &*self.data };
446     }
447 }
448 
449 impl<'rwlock, T> DerefMut for RwLockWriteGuard<'rwlock, T> {
450     fn deref_mut(&mut self) -> &mut Self::Target {
451         return unsafe { &mut *self.data };
452     }
453 }
454 
455 impl<'rwlock, T> Drop for RwLockReadGuard<'rwlock, T> {
456     fn drop(&mut self) {
457         debug_assert!(self.lock.load(Ordering::Relaxed) & !(WRITER | UPGRADED) > 0);
458         self.lock.fetch_sub(READER, Ordering::Release);
459         preempt_enable();
460     }
461 }
462 
463 impl<'rwlock, T> Drop for RwLockUpgradableGuard<'rwlock, T> {
464     fn drop(&mut self) {
465         debug_assert_eq!(
466             self.inner.lock.load(Ordering::Relaxed) & (WRITER | UPGRADED),
467             UPGRADED
468         );
469         self.inner.lock.fetch_sub(UPGRADED, Ordering::AcqRel);
470         preempt_enable();
471         //这里为啥要AcqRel? Release应该就行了?
472     }
473 }
474 
475 impl<'rwlock, T> Drop for RwLockWriteGuard<'rwlock, T> {
476     fn drop(&mut self) {
477         debug_assert_eq!(self.inner.lock.load(Ordering::Relaxed) & WRITER, WRITER);
478         self.inner
479             .lock
480             .fetch_and(!(WRITER | UPGRADED), Ordering::Release);
481 
482         preempt_enable();
483     }
484 }
485