aptos_infallible_link/
rwlock.rs

1// Copyright (c) Aptos
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::RwLock as StdRwLock;
5
6pub use std::sync::{RwLockReadGuard, RwLockWriteGuard};
7
8/// A simple wrapper around the lock() function of a std::sync::RwLock
9/// The only difference is that you don't need to call unwrap() on it.
10#[derive(Debug, Default)]
11pub struct RwLock<T>(StdRwLock<T>);
12
13impl<T> RwLock<T> {
14    /// creates a read-write lock
15    pub fn new(t: T) -> Self {
16        Self(StdRwLock::new(t))
17    }
18
19    /// lock the rwlock in read mode
20    pub fn read(&self) -> RwLockReadGuard<'_, T> {
21        self.0
22            .read()
23            .expect("Cannot currently handle a poisoned lock")
24    }
25
26    /// lock the rwlock in write mode
27    pub fn write(&self) -> RwLockWriteGuard<'_, T> {
28        self.0
29            .write()
30            .expect("Cannot currently handle a poisoned lock")
31    }
32
33    /// return the owned type consuming the lock
34    pub fn into_inner(self) -> T {
35        self.0
36            .into_inner()
37            .expect("Cannot currently handle a poisoned lock")
38    }
39}
40
41#[cfg(test)]
42mod tests {
43
44    use super::*;
45    use std::{sync::Arc, thread};
46
47    #[test]
48    fn test_aptos_rwlock() {
49        let a = 7u8;
50        let rwlock = Arc::new(RwLock::new(a));
51        let rwlock2 = rwlock.clone();
52        let rwlock3 = rwlock.clone();
53
54        let thread1 = thread::spawn(move || {
55            let mut b = rwlock2.write();
56            *b = 8;
57        });
58        let thread2 = thread::spawn(move || {
59            let mut b = rwlock3.write();
60            *b = 9;
61        });
62
63        let _ = thread1.join();
64        let _ = thread2.join();
65
66        let _read = rwlock.read();
67    }
68}