ostd/sync/
rwarc.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
// SPDX-License-Identifier: MPL-2.0

use alloc::sync::Arc;
use core::sync::atomic::{fence, AtomicUsize, Ordering};

use super::{PreemptDisabled, RwLock, RwLockReadGuard, RwLockWriteGuard};

/// A reference-counting pointer with read-write capabilities.
///
/// This is essentially `Arc<RwLock<T>>`, so it can provide read-write capabilities through
/// [`RwArc::read`] and [`RwArc::write`].
///
/// In addition, this allows to derive another reference-counting pointer with read-only
/// capabilities ([`RoArc`]) via [`RwArc::clone_ro`].
///
/// The purpose of having this type is to allow lockless (read) access to the underlying data when
/// there is only one [`RwArc`] instance for the particular allocation (note that there can be any
/// number of [`RoArc`] instances for that allocation). See the [`RwArc::get`] method for more
/// details.
pub struct RwArc<T>(Arc<Inner<T>>);

/// A reference-counting pointer with read-only capabilities.
///
/// This type can be created from an existing [`RwArc`] using its [`RwArc::clone_ro`] method. See
/// the type and method documentation for more details.
pub struct RoArc<T>(Arc<Inner<T>>);

struct Inner<T> {
    data: RwLock<T>,
    num_rw: AtomicUsize,
}

impl<T> RwArc<T> {
    /// Creates a new `RwArc<T>`.
    pub fn new(data: T) -> Self {
        let inner = Inner {
            data: RwLock::new(data),
            num_rw: AtomicUsize::new(1),
        };
        Self(Arc::new(inner))
    }

    /// Acquires the read lock for immutable access.
    pub fn read(&self) -> RwLockReadGuard<T, PreemptDisabled> {
        self.0.data.read()
    }

    /// Acquires the write lock for mutable access.
    pub fn write(&self) -> RwLockWriteGuard<T, PreemptDisabled> {
        self.0.data.write()
    }

    /// Returns an immutable reference if no other `RwArc` points to the same allocation.
    ///
    /// This method is cheap because it does not acquire a lock.
    ///
    /// It's still sound because:
    /// - The mutable reference to `self` and the condition ensure that we are exclusively
    ///   accessing the unique `RwArc` instance for the particular allocation.
    /// - There may be any number of [`RoArc`]s pointing to the same allocation, but they may only
    ///   produce immutable references to the underlying data.
    pub fn get(&mut self) -> Option<&T> {
        if self.0.num_rw.load(Ordering::Relaxed) > 1 {
            return None;
        }

        // This will synchronize with `RwArc::drop` to make sure its changes are visible to us.
        fence(Ordering::Acquire);

        let data_ptr = self.0.data.as_ptr();

        // SAFETY: The data is valid. During the lifetime, no one will be able to create a mutable
        // reference to the data, so it's okay to create an immutable reference like the one below.
        Some(unsafe { &*data_ptr })
    }

    /// Clones a [`RoArc`] that points to the same allocation.
    pub fn clone_ro(&self) -> RoArc<T> {
        RoArc(self.0.clone())
    }
}

impl<T> Clone for RwArc<T> {
    fn clone(&self) -> Self {
        let inner = self.0.clone();

        // Note that overflowing the counter will make it unsound. But not to worry: the above
        // `Arc::clone` must have already aborted the kernel before this happens.
        inner.num_rw.fetch_add(1, Ordering::Relaxed);

        Self(inner)
    }
}

impl<T> Drop for RwArc<T> {
    fn drop(&mut self) {
        self.0.num_rw.fetch_sub(1, Ordering::Release);
    }
}

impl<T: Clone> RwArc<T> {
    /// Returns the contained value by cloning it.
    pub fn get_cloned(&self) -> T {
        let guard = self.read();
        guard.clone()
    }
}

impl<T> RoArc<T> {
    /// Acquires the read lock for immutable access.
    pub fn read(&self) -> RwLockReadGuard<T, PreemptDisabled> {
        self.0.data.read()
    }
}

#[cfg(ktest)]
mod test {
    use super::*;
    use crate::prelude::*;

    #[ktest]
    fn lockless_get() {
        let mut rw1 = RwArc::new(1u32);
        assert_eq!(rw1.get(), Some(1).as_ref());

        let _ro = rw1.clone_ro();
        assert_eq!(rw1.get(), Some(1).as_ref());

        let rw2 = rw1.clone();
        assert_eq!(rw1.get(), None);

        drop(rw2);
        assert_eq!(rw1.get(), Some(1).as_ref());
    }
}