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 crate::{
10    LockAfter, LockBefore, LockDepMutex, LockDepRwLock, LockFor, Locked, RwLockFor,
11    UninterruptibleLock,
12};
13
14use lock_api::RawMutex;
15use std::{any, fmt};
16
17pub use fuchsia_sync::{
18    MappedMutexGuard, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
19};
20
21/// A trait for lock guards that can be temporarily unlocked asynchronously.
22/// This is useful for performing async operations while holding a lock, without
23/// causing deadlocks or holding the lock for an extended period.
24#[async_trait::async_trait(?Send)]
25pub trait AsyncUnlockable {
26    /// Temporarily unlocks the guard `s`, executes the async function `f`, and then
27    /// re-locks the guard.
28    /// The lock is guaranteed to be re-acquired before this function returns.
29    async fn unlocked_async<F, U>(s: &mut Self, f: F) -> U
30    where
31        F: AsyncFnOnce() -> U;
32}
33
34#[async_trait::async_trait(?Send)]
35impl<'a, T> crate::AsyncUnlockable for MutexGuard<'a, T> {
36    async fn unlocked_async<F, U>(s: &mut Self, f: F) -> U
37    where
38        F: AsyncFnOnce() -> U,
39    {
40        // SAFETY: The guard always have a lock mutex.
41        unsafe {
42            Self::mutex(s).raw().unlock();
43        }
44        scopeguard::defer!(
45            // SAFETY: The mutex has been unlocked previously.
46            unsafe { Self::mutex(s).raw().lock() }
47        );
48        f().await
49    }
50}
51
52/// A generic mutex for the ordered_lock operations.
53pub trait MutexLike {
54    type Guard<'a>
55    where
56        Self: 'a;
57    type Context;
58
59    fn context() -> Self::Context;
60
61    /// Lock the mutex. `level` is the index of the locked mutex in the lock ordering.
62    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_>;
63
64    #[inline(always)]
65    fn key(&self) -> *const ()
66    where
67        Self: Sized,
68    {
69        self as *const Self as *const ()
70    }
71}
72
73impl<T> MutexLike for Mutex<T> {
74    type Guard<'a>
75        = MutexGuard<'a, T>
76    where
77        T: 'a;
78    type Context = ();
79
80    #[inline(always)]
81    fn context() -> Self::Context {
82        ()
83    }
84
85    #[inline(always)]
86    fn lock(&self, _context: &mut Self::Context) -> Self::Guard<'_> {
87        return self.lock();
88    }
89}
90
91/// A generic rwlock for the ordered_lock operations.
92pub trait RwLockLike {
93    type ReadGuard<'a>
94    where
95        Self: 'a;
96    type WriteGuard<'a>
97    where
98        Self: 'a;
99    type Context;
100
101    fn context() -> Self::Context;
102
103    fn read(&self, context: &mut Self::Context) -> Self::ReadGuard<'_>;
104    fn write(&self, context: &mut Self::Context) -> Self::WriteGuard<'_>;
105
106    #[inline(always)]
107    fn key(&self) -> *const ()
108    where
109        Self: Sized,
110    {
111        self as *const Self as *const ()
112    }
113}
114
115impl<T> RwLockLike for RwLock<T> {
116    type ReadGuard<'a>
117        = RwLockReadGuard<'a, T>
118    where
119        T: 'a;
120    type WriteGuard<'a>
121        = RwLockWriteGuard<'a, T>
122    where
123        T: 'a;
124    type Context = ();
125
126    #[inline(always)]
127    fn context() -> Self::Context {
128        ()
129    }
130
131    #[inline(always)]
132    fn read(&self, _context: &mut Self::Context) -> Self::ReadGuard<'_> {
133        self.read()
134    }
135
136    #[inline(always)]
137    fn write(&self, _context: &mut Self::Context) -> Self::WriteGuard<'_> {
138        self.write()
139    }
140}
141
142/// Lock `m1` and `m2` in a consistent order (using the memory address of m1 and m2 and returns the
143/// associated guard. This ensure that `ordered_lock(m1, m2)` and `ordered_lock(m2, m1)` will not
144/// deadlock.
145pub fn ordered_lock<'a, M: MutexLike>(m1: &'a M, m2: &'a M) -> (M::Guard<'a>, M::Guard<'a>) {
146    let mut context = M::context();
147    if m1.key() < m2.key() {
148        let g1 = m1.lock(&mut context);
149        let g2 = m2.lock(&mut context);
150        (g1, g2)
151    } else {
152        let g2 = m2.lock(&mut context);
153        let g1 = m1.lock(&mut context);
154        (g1, g2)
155    }
156}
157
158/// Acquires multiple mutexes in a consistent order based on their memory addresses.
159/// This helps prevent deadlocks.
160pub fn ordered_lock_vec<'a, M: MutexLike>(mutexes: &[&'a M]) -> Vec<M::Guard<'a>> {
161    let mut context = M::context();
162
163    // Create a vector of tuples containing the mutex and its original index.
164    let mut indexed_mutexes = mutexes.iter().enumerate().map(|(i, m)| (i, *m)).collect::<Vec<_>>();
165
166    // Sort the indexed mutexes by their keys.
167    indexed_mutexes.sort_by_key(|(_, m)| m.key());
168
169    // Acquire the locks in the sorted order.
170    let mut guards =
171        indexed_mutexes.into_iter().map(|(i, m)| (i, m.lock(&mut context))).collect::<Vec<_>>();
172
173    // Reorder the guards to match the original order of the mutexes.
174    guards.sort_by_key(|(i, _)| *i);
175
176    guards.into_iter().map(|(_, g)| g).collect::<Vec<_>>()
177}
178
179/// Lock `r1` and `r2` in a consistent order (using the memory address of r1 and r2) for reading.
180pub fn ordered_read_lock<'a, R: RwLockLike>(
181    r1: &'a R,
182    r2: &'a R,
183) -> (R::ReadGuard<'a>, R::ReadGuard<'a>) {
184    let w1 = RwLockReadWrapper(r1);
185    let w2 = RwLockReadWrapper(r2);
186    ordered_lock(&w1, &w2)
187}
188
189/// Lock `r1` and `r2` in a consistent order (using the memory address of r1 and r2) for writing.
190pub fn ordered_write_lock<'a, R: RwLockLike>(
191    r1: &'a R,
192    r2: &'a R,
193) -> (R::WriteGuard<'a>, R::WriteGuard<'a>) {
194    let w1 = RwLockWriteWrapper(r1);
195    let w2 = RwLockWriteWrapper(r2);
196    ordered_lock(&w1, &w2)
197}
198
199/// Acquires multiple rwlocks in a consistent order based on their memory addresses for reading.
200pub fn ordered_read_lock_vec<'a, R: RwLockLike>(rwlocks: &[&'a R]) -> Vec<R::ReadGuard<'a>> {
201    let wrappers = rwlocks.iter().map(|r| RwLockReadWrapper(*r)).collect::<Vec<_>>();
202    let wrapper_refs = wrappers.iter().collect::<Vec<_>>();
203    ordered_lock_vec(&wrapper_refs)
204}
205
206/// Acquires multiple rwlocks in a consistent order based on their memory addresses for writing.
207pub fn ordered_write_lock_vec<'a, R: RwLockLike>(rwlocks: &[&'a R]) -> Vec<R::WriteGuard<'a>> {
208    let wrappers = rwlocks.iter().map(|r| RwLockWriteWrapper(*r)).collect::<Vec<_>>();
209    let wrapper_refs = wrappers.iter().collect::<Vec<_>>();
210    ordered_lock_vec(&wrapper_refs)
211}
212
213/// A wrapper for mutex that requires a `Locked` context to acquire.
214/// This context must be of a level that precedes `L` in the lock ordering graph
215/// where `L` is a level associated with this mutex.
216pub struct OrderedMutex<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> {
217    mutex: LockDepMutex<T, L>,
218}
219
220impl<T: Default, L: LockAfter<UninterruptibleLock> + crate::LockLevel> Default
221    for OrderedMutex<T, L>
222{
223    fn default() -> Self {
224        Self { mutex: LockDepMutex::new(T::default()) }
225    }
226}
227
228impl<T: fmt::Debug, L: LockAfter<UninterruptibleLock> + crate::LockLevel> fmt::Debug
229    for OrderedMutex<T, L>
230{
231    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
232        write!(f, "OrderedMutex({:?}, {})", self.mutex, any::type_name::<L>())
233    }
234}
235
236impl<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> LockFor<L> for OrderedMutex<T, L> {
237    type Data = T;
238    type Guard<'a>
239        = crate::LockDepGuard<'a, T>
240    where
241        T: 'a,
242        L: 'a;
243    fn lock(&self) -> Self::Guard<'_> {
244        self.mutex.lock()
245    }
246}
247
248impl<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> OrderedMutex<T, L> {
249    pub const fn new(t: T) -> Self {
250        Self { mutex: LockDepMutex::new(t) }
251    }
252
253    pub fn lock<'a, P>(&'a self, locked: &'a mut Locked<P>) -> <Self as LockFor<L>>::Guard<'a>
254    where
255        P: LockBefore<L>,
256    {
257        locked.lock(self)
258    }
259
260    pub fn lock_and<'a, P>(
261        &'a self,
262        locked: &'a mut Locked<P>,
263    ) -> (<Self as LockFor<L>>::Guard<'a>, &'a mut Locked<L>)
264    where
265        P: LockBefore<L>,
266    {
267        locked.lock_and(self)
268    }
269}
270
271/// Lock two OrderedMutex of the same level in the consistent order. Returns both
272/// guards and a new locked context.
273pub fn lock_both<'a, T, L: LockAfter<UninterruptibleLock> + crate::LockLevel, P>(
274    locked: &'a mut Locked<P>,
275    m1: &'a OrderedMutex<T, L>,
276    m2: &'a OrderedMutex<T, L>,
277) -> (crate::LockDepGuard<'a, T>, crate::LockDepGuard<'a, T>, &'a mut Locked<L>)
278where
279    P: LockBefore<L>,
280{
281    locked.lock_both_and(m1, m2)
282}
283
284/// A wrapper for an RwLock that requires a `Locked` context to acquire.
285/// This context must be of a level that precedes `L` in the lock ordering graph
286/// where `L` is a level associated with this RwLock.
287pub struct OrderedRwLock<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> {
288    rwlock: LockDepRwLock<T, L>,
289}
290
291impl<T: Default, L: LockAfter<UninterruptibleLock> + crate::LockLevel> Default
292    for OrderedRwLock<T, L>
293{
294    fn default() -> Self {
295        Self { rwlock: LockDepRwLock::new(T::default()) }
296    }
297}
298
299impl<T: fmt::Debug, L: LockAfter<UninterruptibleLock> + crate::LockLevel> fmt::Debug
300    for OrderedRwLock<T, L>
301{
302    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
303        write!(f, "OrderedRwLock({:?}, {})", self.rwlock, any::type_name::<L>())
304    }
305}
306
307impl<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> RwLockFor<L> for OrderedRwLock<T, L> {
308    type Data = T;
309    type ReadGuard<'a>
310        = crate::LockDepReadGuard<'a, T>
311    where
312        T: 'a,
313        L: 'a;
314    type WriteGuard<'a>
315        = crate::LockDepWriteGuard<'a, T>
316    where
317        T: 'a,
318        L: 'a;
319    fn read_lock(&self) -> Self::ReadGuard<'_> {
320        self.rwlock.read()
321    }
322    fn write_lock(&self) -> Self::WriteGuard<'_> {
323        self.rwlock.write()
324    }
325}
326
327impl<T, L: LockAfter<UninterruptibleLock> + crate::LockLevel> OrderedRwLock<T, L> {
328    pub const fn new(t: T) -> Self {
329        Self { rwlock: LockDepRwLock::new(t) }
330    }
331
332    pub fn read<'a, P>(&'a self, locked: &'a mut Locked<P>) -> <Self as RwLockFor<L>>::ReadGuard<'a>
333    where
334        P: LockBefore<L>,
335    {
336        locked.read_lock(self)
337    }
338
339    pub fn write<'a, P>(
340        &'a self,
341        locked: &'a mut Locked<P>,
342    ) -> <Self as RwLockFor<L>>::WriteGuard<'a>
343    where
344        P: LockBefore<L>,
345    {
346        locked.write_lock(self)
347    }
348
349    pub fn read_and<'a, P>(
350        &'a self,
351        locked: &'a mut Locked<P>,
352    ) -> (<Self as RwLockFor<L>>::ReadGuard<'a>, &'a mut Locked<L>)
353    where
354        P: LockBefore<L>,
355    {
356        locked.read_lock_and(self)
357    }
358
359    pub fn write_and<'a, P>(
360        &'a self,
361        locked: &'a mut Locked<P>,
362    ) -> (<Self as RwLockFor<L>>::WriteGuard<'a>, &'a mut Locked<L>)
363    where
364        P: LockBefore<L>,
365    {
366        locked.write_lock_and(self)
367    }
368}
369
370struct RwLockReadWrapper<'a, R>(&'a R);
371struct RwLockWriteWrapper<'a, R>(&'a R);
372
373impl<'a, R: RwLockLike> MutexLike for RwLockReadWrapper<'a, R> {
374    type Guard<'b>
375        = R::ReadGuard<'a>
376    where
377        Self: 'b;
378    type Context = R::Context;
379
380    #[inline(always)]
381    fn context() -> Self::Context {
382        R::context()
383    }
384
385    #[inline(always)]
386    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_> {
387        self.0.read(context)
388    }
389
390    #[inline(always)]
391    fn key(&self) -> *const () {
392        self.0.key()
393    }
394}
395
396impl<'a, R: RwLockLike> MutexLike for RwLockWriteWrapper<'a, R> {
397    type Guard<'b>
398        = R::WriteGuard<'a>
399    where
400        Self: 'b;
401    type Context = R::Context;
402
403    #[inline(always)]
404    fn context() -> Self::Context {
405        R::context()
406    }
407
408    #[inline(always)]
409    fn lock(&self, context: &mut Self::Context) -> Self::Guard<'_> {
410        self.0.write(context)
411    }
412
413    #[inline(always)]
414    fn key(&self) -> *const () {
415        self.0.key()
416    }
417}
418#[cfg(test)]
419mod test {
420    use super::*;
421    use crate::Unlocked;
422
423    #[::fuchsia::test]
424    fn test_lock_ordering() {
425        let l1 = Mutex::new(1);
426        let l2 = Mutex::new(2);
427
428        {
429            let (g1, g2) = ordered_lock(&l1, &l2);
430            assert_eq!(*g1, 1);
431            assert_eq!(*g2, 2);
432        }
433        {
434            let (g2, g1) = ordered_lock(&l2, &l1);
435            assert_eq!(*g1, 1);
436            assert_eq!(*g2, 2);
437        }
438    }
439
440    #[::fuchsia::test]
441    fn test_vec_lock_ordering() {
442        let l1 = Mutex::new(1);
443        let l0 = Mutex::new(0);
444        let l2 = Mutex::new(2);
445
446        {
447            let guards = ordered_lock_vec(&[&l0, &l1, &l2]);
448            assert_eq!(*guards[0], 0);
449            assert_eq!(*guards[1], 1);
450            assert_eq!(*guards[2], 2);
451        }
452        {
453            let guards = ordered_lock_vec(&[&l2, &l1, &l0]);
454            assert_eq!(*guards[0], 2);
455            assert_eq!(*guards[1], 1);
456            assert_eq!(*guards[2], 0);
457        }
458    }
459
460    mod lock_levels {
461        //! Lock ordering tree:
462        //! Unlocked -> A -> B -> C
463        //!          -> D -> E -> F
464        use crate::{LockAfter, UninterruptibleLock, Unlocked};
465        use lock_ordering_macro::lock_ordering;
466        lock_ordering! {
467            Unlocked => A,
468            A => B,
469            B => C,
470            Unlocked => D,
471            D => E,
472            E => F,
473        }
474
475        impl LockAfter<UninterruptibleLock> for A {}
476        impl LockAfter<UninterruptibleLock> for B {}
477        impl LockAfter<UninterruptibleLock> for C {}
478        impl LockAfter<UninterruptibleLock> for D {}
479        impl LockAfter<UninterruptibleLock> for E {}
480        impl LockAfter<UninterruptibleLock> for F {}
481    }
482
483    use lock_levels::{A, B, C, D, E, F};
484
485    #[test]
486    fn test_ordered_mutex() {
487        let a: OrderedMutex<u8, A> = OrderedMutex::new(15);
488        let _b: OrderedMutex<u16, B> = OrderedMutex::new(30);
489        let c: OrderedMutex<u32, C> = OrderedMutex::new(45);
490
491        #[allow(
492            clippy::undocumented_unsafe_blocks,
493            reason = "Force documented unsafe blocks in Starnix"
494        )]
495        let locked = unsafe { Unlocked::new() };
496
497        let (a_data, mut next_locked) = a.lock_and(locked);
498        let c_data = c.lock(&mut next_locked);
499
500        // This won't compile
501        //let _b_data = _b.lock(locked);
502        //let _b_data = _b.lock(&mut next_locked);
503
504        assert_eq!(&*a_data, &15);
505        assert_eq!(&*c_data, &45);
506    }
507    #[test]
508    fn test_ordered_rwlock() {
509        let d: OrderedRwLock<u8, D> = OrderedRwLock::new(15);
510        let _e: OrderedRwLock<u16, E> = OrderedRwLock::new(30);
511        let f: OrderedRwLock<u32, F> = OrderedRwLock::new(45);
512
513        #[allow(
514            clippy::undocumented_unsafe_blocks,
515            reason = "Force documented unsafe blocks in Starnix"
516        )]
517        let locked = unsafe { Unlocked::new() };
518        {
519            let (d_data, mut next_locked) = d.write_and(locked);
520            let f_data = f.read(&mut next_locked);
521
522            // This won't compile
523            //let _e_data = _e.read(locked);
524            //let _e_data = _e.read(&mut next_locked);
525
526            assert_eq!(&*d_data, &15);
527            assert_eq!(&*f_data, &45);
528        }
529        {
530            let (d_data, mut next_locked) = d.read_and(locked);
531            let f_data = f.write(&mut next_locked);
532
533            // This won't compile
534            //let _e_data = _e.write(locked);
535            //let _e_data = _e.write(&mut next_locked);
536
537            assert_eq!(&*d_data, &15);
538            assert_eq!(&*f_data, &45);
539        }
540    }
541
542    #[test]
543    fn test_lock_both() {
544        let a1: OrderedMutex<u8, A> = OrderedMutex::new(15);
545        let a2: OrderedMutex<u8, A> = OrderedMutex::new(30);
546        #[allow(
547            clippy::undocumented_unsafe_blocks,
548            reason = "Force documented unsafe blocks in Starnix"
549        )]
550        let locked = unsafe { Unlocked::new() };
551        {
552            let (a1_data, a2_data, _) = lock_both(locked, &a1, &a2);
553            assert_eq!(&*a1_data, &15);
554            assert_eq!(&*a2_data, &30);
555        }
556        {
557            let (a2_data, a1_data, _) = lock_both(locked, &a2, &a1);
558            assert_eq!(&*a1_data, &15);
559            assert_eq!(&*a2_data, &30);
560        }
561    }
562
563    #[::fuchsia::test]
564    fn test_ordered_rwlock_wrappers() {
565        let l1: LockDepRwLock<u8, A> = LockDepRwLock::new(1);
566        let l2: LockDepRwLock<u8, A> = LockDepRwLock::new(2);
567
568        {
569            let (g1, g2) = ordered_read_lock(&l1, &l2);
570            assert_eq!(*g1, 1);
571            assert_eq!(*g2, 2);
572        }
573        {
574            let (g2, g1) = ordered_read_lock(&l2, &l1);
575            assert_eq!(*g1, 1);
576            assert_eq!(*g2, 2);
577        }
578        {
579            let (g1, g2) = ordered_write_lock(&l1, &l2);
580            assert_eq!(*g1, 1);
581            assert_eq!(*g2, 2);
582        }
583        {
584            let (g2, g1) = ordered_write_lock(&l2, &l1);
585            assert_eq!(*g1, 1);
586            assert_eq!(*g2, 2);
587        }
588    }
589
590    #[::fuchsia::test]
591    fn test_ordered_rwlock_vec() {
592        let l1: LockDepRwLock<u8, A> = LockDepRwLock::new(1);
593        let l0: LockDepRwLock<u8, A> = LockDepRwLock::new(0);
594        let l2: LockDepRwLock<u8, A> = LockDepRwLock::new(2);
595
596        {
597            let guards = ordered_read_lock_vec(&[&l0, &l1, &l2]);
598            assert_eq!(*guards[0], 0);
599            assert_eq!(*guards[1], 1);
600            assert_eq!(*guards[2], 2);
601        }
602        {
603            let guards = ordered_write_lock_vec(&[&l2, &l1, &l0]);
604            assert_eq!(*guards[0], 2);
605            assert_eq!(*guards[1], 1);
606            assert_eq!(*guards[2], 0);
607        }
608    }
609}