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::include::bindings::bindings::EOVERFLOW; 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, i32> { 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(-(EOVERFLOW as i32)); 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