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