Skip to main content

kernel_api/sync/
rwlock.rs

1use core::fmt::Formatter;
2use core::panic::Location;
3use core::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
4use log::warn;
5
6/// A reader-writer lock.
7pub type RwSpinlock<T> = lock_api::RwLock<RwCount, T>;
8
9/// RAII structure used to release the shared read access of a lock when dropped.
10pub type RwReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, RwCount, T>;
11
12/// RAII structure used to release upgradable read access of a lock when dropped.
13pub type RwUpgradableReadGuard<'a, T> = lock_api::RwLockUpgradableReadGuard<'a, RwCount, T>;
14
15/// RAII structure used to release the exclusive write access of a lock when dropped.
16pub type RwWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, RwCount, T>;
17
18// INVARIANT: `self.1` is always a valid 'static ref to a `Location<'static>` unless null
19#[doc(hidden)]
20pub struct RwCount(AtomicUsize, #[cfg(debug_assertions)] AtomicPtr<Location<'static>>);
21
22// FIXME: Deadlocks due to interrupts
23impl RwCount {
24    const WRITE_BIT_MASK: usize = 1<<(usize::BITS as usize - 1);
25    const UPGRADEABLE_BIT_MASK: usize = 1<<(usize::BITS as usize - 2);
26    const READ_COUNT_MASK: usize = !(Self::WRITE_BIT_MASK | Self::UPGRADEABLE_BIT_MASK);
27
28    #[cfg(debug_assertions)]
29    fn lock_location(&self) -> Option<&'static Location<'static>> {
30        let location = self.1.load(Ordering::Relaxed);
31	    // SAFETY: invariant of type that `location` is always valid unless null
32        unsafe { location.as_ref::<'static>() }
33    }
34}
35
36// SAFETY: number of readers/writers checked on each attempt at locking
37unsafe impl lock_api::RawRwLock for RwCount {
38    const INIT: Self = Self(AtomicUsize::new(0), #[cfg(debug_assertions)] AtomicPtr::new(core::ptr::null_mut()));
39    type GuardMarker = lock_api::GuardNoSend; // Can always be relaxed in future
40
41    #[track_caller]
42    fn lock_shared(&self) {
43        while !self.try_lock_shared() {
44            core::hint::spin_loop();
45        }
46    }
47
48    #[track_caller]
49    fn try_lock_shared(&self) -> bool {
50        let mut old_value = self.0.load(Ordering::Relaxed);
51
52        loop {
53            let old_normal_count = old_value & Self::READ_COUNT_MASK;
54
55            assert_ne!(old_normal_count, Self::READ_COUNT_MASK, "Reader count overflowed");
56            if (old_value & Self::WRITE_BIT_MASK) != 0 { return false; }
57
58            let new_value = (old_normal_count + 1) | (old_value & Self::UPGRADEABLE_BIT_MASK);
59
60            match self.0.compare_exchange_weak(old_value, new_value, Ordering::Acquire, Ordering::Relaxed) {
61                Ok(_) => {
62                    #[cfg(debug_assertions)] self.1.store(core::ptr::from_ref(Location::caller()).cast_mut(), Ordering::Relaxed);
63                    return true;
64                },
65                Err(new_old_value) => {
66                    #[cfg(debug_assertions)] warn!("locked at {:?}", self.lock_location());
67                    old_value = new_old_value;
68                }
69            }
70        }
71    }
72
73    unsafe fn unlock_shared(&self) {
74        let mut old_value = self.0.load(Ordering::Relaxed);
75        loop {
76            let old_normal_count = old_value & Self::READ_COUNT_MASK;
77
78            debug_assert_eq!(old_value & Self::WRITE_BIT_MASK, 0, "BUG: RwLock reader dropped while writer was active");
79
80            let new_value = (old_normal_count - 1) | (old_value & !Self::READ_COUNT_MASK);
81            match self.0.compare_exchange_weak(old_value, new_value, Ordering::Release, Ordering::Relaxed) {
82                Ok(_) => return,
83                Err(new_old_value) => old_value = new_old_value
84            }
85        }
86    }
87
88    #[track_caller]
89    fn lock_exclusive(&self) {
90        while !self.try_lock_exclusive() {
91            core::hint::spin_loop();
92        }
93    }
94
95    #[track_caller]
96    fn try_lock_exclusive(&self) -> bool {
97        let res = self.0.compare_exchange_weak(0, Self::WRITE_BIT_MASK, Ordering::Acquire, Ordering::Relaxed)
98            .is_ok();
99
100        #[cfg(debug_assertions)] if res {
101            self.1.store(core::ptr::from_ref(Location::caller()).cast_mut(), Ordering::Relaxed);
102        } else {
103            warn!("locked at {:?}", self.lock_location());
104        }
105
106        res
107    }
108
109    unsafe fn unlock_exclusive(&self) {
110        if cfg!(debug_assertions) {
111            self.0.compare_exchange(Self::WRITE_BIT_MASK, 0, Ordering::Release, Ordering::Relaxed)
112                .expect("BUG: RwLock writer dropped while readers were active");
113        } else {
114            self.0.store(0, Ordering::Release);
115        }
116    }
117}
118
119// SAFETY: See comment on base RawRwLock impl
120unsafe impl lock_api::RawRwLockDowngrade for RwCount {
121    unsafe fn downgrade(&self) {
122        if cfg!(debug_assertions) {
123            self.0.compare_exchange(Self::WRITE_BIT_MASK, 1, Ordering::SeqCst, Ordering::Relaxed)
124                .expect("BUG: RwLock writer downgraded while readers were active");
125        } else {
126            // No existing readers should exist therefore can unconditionally set read count to 1
127            self.0.store(1, Ordering::SeqCst); // FIXME: what order to use here
128        }
129    }
130}
131
132// SAFETY: See comment on base RawRwLock impl
133unsafe impl lock_api::RawRwLockUpgrade for RwCount {
134    fn lock_upgradable(&self) {
135        while !self.try_lock_upgradable() {
136            core::hint::spin_loop();
137        }
138    }
139
140    fn try_lock_upgradable(&self) -> bool {
141        let mut old_value = self.0.load(Ordering::Relaxed);
142
143        loop {
144            if (old_value & Self::WRITE_BIT_MASK) != 0 { return false; }
145            if (old_value & Self::UPGRADEABLE_BIT_MASK) != 0 { return false; }
146
147            let new_value = old_value | Self::UPGRADEABLE_BIT_MASK;
148
149            match self.0.compare_exchange_weak(old_value, new_value, Ordering::Acquire, Ordering::Relaxed) {
150                Ok(_) => return true,
151                Err(new_old_value) => old_value = new_old_value
152            }
153        }
154    }
155
156    unsafe fn unlock_upgradable(&self) {
157        let mut old_value = self.0.load(Ordering::Relaxed);
158        loop {
159            debug_assert_eq!(old_value & Self::WRITE_BIT_MASK, 0, "BUG: RwLock upgradable reader dropped while writer was active");
160
161            let new_value = old_value & !Self::UPGRADEABLE_BIT_MASK;
162            match self.0.compare_exchange_weak(old_value, new_value, Ordering::Release, Ordering::Relaxed) {
163                Ok(_) => return,
164                Err(new_old_value) => old_value = new_old_value
165            }
166        }
167    }
168
169    unsafe fn upgrade(&self) {
170        // SAFETY: `RawRwLockUpgrade::upgrade()` and `RawRwLockUpgrade::try_upgrade()` have same safety requirements
171        while unsafe { !self.try_upgrade() } {
172            core::hint::spin_loop();
173        }
174    }
175
176    unsafe fn try_upgrade(&self) -> bool {
177        self.0.compare_exchange_weak(Self::UPGRADEABLE_BIT_MASK, Self::WRITE_BIT_MASK, Ordering::Relaxed, Ordering::Relaxed)
178            .is_ok()
179    }
180}
181
182// SAFETY: See comment on base RawRwLock impl
183unsafe impl lock_api::RawRwLockUpgradeDowngrade for RwCount {
184    unsafe fn downgrade_upgradable(&self) {
185        let mut old_value = self.0.load(Ordering::Relaxed);
186        loop {
187            debug_assert_eq!(old_value & Self::WRITE_BIT_MASK, 0, "BUG: RwLock upgradable reader downgraded while writer was active");
188
189            let old_normal_count = old_value & Self::READ_COUNT_MASK;
190
191            assert_ne!(old_normal_count, Self::READ_COUNT_MASK, "Reader count overflowed");
192
193            let new_value = old_normal_count + 1;
194            match self.0.compare_exchange_weak(old_value, new_value, Ordering::Relaxed, Ordering::Relaxed) {
195                Ok(_) => return,
196                Err(new_old_value) => old_value = new_old_value
197            }
198        }
199    }
200
201    unsafe fn downgrade_to_upgradable(&self) {
202        if cfg!(debug_assertions) {
203            self.0.compare_exchange(Self::WRITE_BIT_MASK, Self::UPGRADEABLE_BIT_MASK, Ordering::SeqCst, Ordering::Relaxed)
204                .expect("BUG: RwLock writer downgraded while readers were active");
205        } else {
206            // No existing readers should exist therefore can unconditionally set upgradable bit
207            self.0.store(Self::UPGRADEABLE_BIT_MASK, Ordering::SeqCst); // FIXME: what ordering to use here?
208        }
209    }
210}
211
212impl core::fmt::Debug for RwCount {
213    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
214        let val = self.0.load(Ordering::Relaxed);
215        let write = val & Self::WRITE_BIT_MASK != 0;
216        let read = val & Self::READ_COUNT_MASK;
217        let upgradeable_reader = val & Self::UPGRADEABLE_BIT_MASK != 0;
218
219        f.debug_struct("RwCount")
220         .field("write", &write)
221         .field("read", &(read + usize::from(upgradeable_reader)))
222         .finish_non_exhaustive()
223    }
224}