Skip to main content

starnix_sync/
locks.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5// Use these crates so that we don't need to make the dependencies conditional.
6use fuchsia_sync as _;
7use lock_api as _;
8
9use lock_api::RawMutex;
10
11pub use fuchsia_sync::{
12    MappedMutexGuard, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
13};
14
15/// A trait for lock guards that can be temporarily unlocked asynchronously.
16/// This is useful for performing async operations while holding a lock, without
17/// causing deadlocks or holding the lock for an extended period.
18#[async_trait::async_trait(?Send)]
19pub trait AsyncUnlockable {
20    /// Temporarily unlocks the guard `s`, executes the async function `f`, and then
21    /// re-locks the guard.
22    /// The lock is guaranteed to be re-acquired before this function returns.
23    async fn unlocked_async<F, U>(s: &mut Self, f: F) -> U
24    where
25        F: AsyncFnOnce() -> U;
26}
27
28#[async_trait::async_trait(?Send)]
29impl<'a, T> crate::AsyncUnlockable for MutexGuard<'a, T> {
30    async fn unlocked_async<F, U>(s: &mut Self, f: F) -> U
31    where
32        F: AsyncFnOnce() -> U,
33    {
34        // SAFETY: The guard always have a lock mutex.
35        unsafe {
36            Self::mutex(s).raw().unlock();
37        }
38        scopeguard::defer!(
39            // SAFETY: The mutex has been unlocked previously.
40            unsafe { Self::mutex(s).raw().lock() }
41        );
42        f().await
43    }
44}
45
46/// A generic mutex for the ordered_lock operations.
47pub trait MutexLike {
48    type Guard<'a>
49    where
50        Self: 'a;
51    type Context;
52
53    fn context() -> Self::Context;
54
55    /// Lock the mutex. `level` is the index of the locked mutex in the lock ordering.
56    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_>;
57
58    #[inline(always)]
59    fn key(&self) -> *const ()
60    where
61        Self: Sized,
62    {
63        self as *const Self as *const ()
64    }
65}
66
67impl<T> MutexLike for Mutex<T> {
68    type Guard<'a>
69        = MutexGuard<'a, T>
70    where
71        T: 'a;
72    type Context = ();
73
74    #[inline(always)]
75    fn context() -> Self::Context {
76        ()
77    }
78
79    #[inline(always)]
80    fn lock(&self, _context: &mut Self::Context) -> Self::Guard<'_> {
81        return self.lock();
82    }
83}
84
85/// A generic rwlock for the ordered_lock operations.
86pub trait RwLockLike {
87    type ReadGuard<'a>
88    where
89        Self: 'a;
90    type WriteGuard<'a>
91    where
92        Self: 'a;
93    type Context;
94
95    fn context() -> Self::Context;
96
97    fn read(&self, context: &mut Self::Context) -> Self::ReadGuard<'_>;
98    fn write(&self, context: &mut Self::Context) -> Self::WriteGuard<'_>;
99
100    #[inline(always)]
101    fn key(&self) -> *const ()
102    where
103        Self: Sized,
104    {
105        self as *const Self as *const ()
106    }
107}
108
109impl<T> RwLockLike for RwLock<T> {
110    type ReadGuard<'a>
111        = RwLockReadGuard<'a, T>
112    where
113        T: 'a;
114    type WriteGuard<'a>
115        = RwLockWriteGuard<'a, T>
116    where
117        T: 'a;
118    type Context = ();
119
120    #[inline(always)]
121    fn context() -> Self::Context {
122        ()
123    }
124
125    #[inline(always)]
126    fn read(&self, _context: &mut Self::Context) -> Self::ReadGuard<'_> {
127        self.read()
128    }
129
130    #[inline(always)]
131    fn write(&self, _context: &mut Self::Context) -> Self::WriteGuard<'_> {
132        self.write()
133    }
134}
135
136/// Lock `m1` and `m2` in a consistent order (using the memory address of m1 and m2 and returns the
137/// associated guard. This ensure that `ordered_lock(m1, m2)` and `ordered_lock(m2, m1)` will not
138/// deadlock.
139pub fn ordered_lock<'a, M: MutexLike>(m1: &'a M, m2: &'a M) -> (M::Guard<'a>, M::Guard<'a>) {
140    let mut context = M::context();
141    if m1.key() < m2.key() {
142        let g1 = m1.lock(&mut context);
143        let g2 = m2.lock(&mut context);
144        (g1, g2)
145    } else {
146        let g2 = m2.lock(&mut context);
147        let g1 = m1.lock(&mut context);
148        (g1, g2)
149    }
150}
151
152/// Acquires multiple mutexes in a consistent order based on their memory addresses.
153/// This helps prevent deadlocks.
154pub fn ordered_lock_vec<'a, M: MutexLike>(mutexes: &[&'a M]) -> Vec<M::Guard<'a>> {
155    let mut context = M::context();
156
157    // Create a vector of tuples containing the mutex and its original index.
158    let mut indexed_mutexes = mutexes.iter().enumerate().map(|(i, m)| (i, *m)).collect::<Vec<_>>();
159
160    // Sort the indexed mutexes by their keys.
161    indexed_mutexes.sort_by_key(|(_, m)| m.key());
162
163    // Acquire the locks in the sorted order.
164    let mut guards =
165        indexed_mutexes.into_iter().map(|(i, m)| (i, m.lock(&mut context))).collect::<Vec<_>>();
166
167    // Reorder the guards to match the original order of the mutexes.
168    guards.sort_by_key(|(i, _)| *i);
169
170    guards.into_iter().map(|(_, g)| g).collect::<Vec<_>>()
171}
172
173/// Lock `r1` and `r2` in a consistent order (using the memory address of r1 and r2) for reading.
174pub fn ordered_read_lock<'a, R: RwLockLike>(
175    r1: &'a R,
176    r2: &'a R,
177) -> (R::ReadGuard<'a>, R::ReadGuard<'a>) {
178    let w1 = RwLockReadWrapper(r1);
179    let w2 = RwLockReadWrapper(r2);
180    ordered_lock(&w1, &w2)
181}
182
183/// Lock `r1` and `r2` in a consistent order (using the memory address of r1 and r2) for writing.
184pub fn ordered_write_lock<'a, R: RwLockLike>(
185    r1: &'a R,
186    r2: &'a R,
187) -> (R::WriteGuard<'a>, R::WriteGuard<'a>) {
188    let w1 = RwLockWriteWrapper(r1);
189    let w2 = RwLockWriteWrapper(r2);
190    ordered_lock(&w1, &w2)
191}
192
193/// Acquires multiple rwlocks in a consistent order based on their memory addresses for reading.
194pub fn ordered_read_lock_vec<'a, R: RwLockLike>(rwlocks: &[&'a R]) -> Vec<R::ReadGuard<'a>> {
195    let wrappers = rwlocks.iter().map(|r| RwLockReadWrapper(*r)).collect::<Vec<_>>();
196    let wrapper_refs = wrappers.iter().collect::<Vec<_>>();
197    ordered_lock_vec(&wrapper_refs)
198}
199
200/// Acquires multiple rwlocks in a consistent order based on their memory addresses for writing.
201pub fn ordered_write_lock_vec<'a, R: RwLockLike>(rwlocks: &[&'a R]) -> Vec<R::WriteGuard<'a>> {
202    let wrappers = rwlocks.iter().map(|r| RwLockWriteWrapper(*r)).collect::<Vec<_>>();
203    let wrapper_refs = wrappers.iter().collect::<Vec<_>>();
204    ordered_lock_vec(&wrapper_refs)
205}
206
207struct RwLockReadWrapper<'a, R>(&'a R);
208struct RwLockWriteWrapper<'a, R>(&'a R);
209
210impl<'a, R: RwLockLike> MutexLike for RwLockReadWrapper<'a, R> {
211    type Guard<'b>
212        = R::ReadGuard<'a>
213    where
214        Self: 'b;
215    type Context = R::Context;
216
217    #[inline(always)]
218    fn context() -> Self::Context {
219        R::context()
220    }
221
222    #[inline(always)]
223    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_> {
224        self.0.read(context)
225    }
226
227    #[inline(always)]
228    fn key(&self) -> *const () {
229        self.0.key()
230    }
231}
232
233impl<'a, R: RwLockLike> MutexLike for RwLockWriteWrapper<'a, R> {
234    type Guard<'b>
235        = R::WriteGuard<'a>
236    where
237        Self: 'b;
238    type Context = R::Context;
239
240    #[inline(always)]
241    fn context() -> Self::Context {
242        R::context()
243    }
244
245    #[inline(always)]
246    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_> {
247        self.0.write(context)
248    }
249
250    #[inline(always)]
251    fn key(&self) -> *const () {
252        self.0.key()
253    }
254}
255#[cfg(test)]
256mod test {
257    use super::*;
258
259    #[::fuchsia::test]
260    fn test_lock_ordering() {
261        let l1 = Mutex::new(1);
262        let l2 = Mutex::new(2);
263
264        {
265            let (g1, g2) = ordered_lock(&l1, &l2);
266            assert_eq!(*g1, 1);
267            assert_eq!(*g2, 2);
268        }
269        {
270            let (g2, g1) = ordered_lock(&l2, &l1);
271            assert_eq!(*g1, 1);
272            assert_eq!(*g2, 2);
273        }
274    }
275
276    #[::fuchsia::test]
277    fn test_vec_lock_ordering() {
278        let l1 = Mutex::new(1);
279        let l0 = Mutex::new(0);
280        let l2 = Mutex::new(2);
281
282        {
283            let guards = ordered_lock_vec(&[&l0, &l1, &l2]);
284            assert_eq!(*guards[0], 0);
285            assert_eq!(*guards[1], 1);
286            assert_eq!(*guards[2], 2);
287        }
288        {
289            let guards = ordered_lock_vec(&[&l2, &l1, &l0]);
290            assert_eq!(*guards[0], 2);
291            assert_eq!(*guards[1], 1);
292            assert_eq!(*guards[2], 0);
293        }
294    }
295
296    #[::fuchsia::test]
297    fn test_ordered_rwlock_wrappers() {
298        let l1: RwLock<u8> = RwLock::new(1);
299        let l2: RwLock<u8> = RwLock::new(2);
300
301        {
302            let (g1, g2) = ordered_read_lock(&l1, &l2);
303            assert_eq!(*g1, 1);
304            assert_eq!(*g2, 2);
305        }
306        {
307            let (g2, g1) = ordered_read_lock(&l2, &l1);
308            assert_eq!(*g1, 1);
309            assert_eq!(*g2, 2);
310        }
311        {
312            let (g1, g2) = ordered_write_lock(&l1, &l2);
313            assert_eq!(*g1, 1);
314            assert_eq!(*g2, 2);
315        }
316        {
317            let (g2, g1) = ordered_write_lock(&l2, &l1);
318            assert_eq!(*g1, 1);
319            assert_eq!(*g2, 2);
320        }
321    }
322
323    #[::fuchsia::test]
324    fn test_ordered_rwlock_vec() {
325        let l1: RwLock<u8> = RwLock::new(1);
326        let l0: RwLock<u8> = RwLock::new(0);
327        let l2: RwLock<u8> = RwLock::new(2);
328
329        {
330            let guards = ordered_read_lock_vec(&[&l0, &l1, &l2]);
331            assert_eq!(*guards[0], 0);
332            assert_eq!(*guards[1], 1);
333            assert_eq!(*guards[2], 2);
334        }
335        {
336            let guards = ordered_write_lock_vec(&[&l2, &l1, &l0]);
337            assert_eq!(*guards[0], 2);
338            assert_eq!(*guards[1], 1);
339            assert_eq!(*guards[2], 0);
340        }
341    }
342}