ostd/sync/rwarc.rs
1// SPDX-License-Identifier: MPL-2.0
2
3use alloc::sync::Arc;
4use core::sync::atomic::{AtomicUsize, Ordering, fence};
5
6use super::{PreemptDisabled, RwLock, RwLockReadGuard, RwLockWriteGuard};
7
8/// A reference-counting pointer with read-write capabilities.
9///
10/// This is essentially `Arc<RwLock<T>>`, so it can provide read-write capabilities through
11/// [`RwArc::read`] and [`RwArc::write`].
12///
13/// In addition, this allows to derive another reference-counting pointer with read-only
14/// capabilities ([`RoArc`]) via [`RwArc::clone_ro`].
15///
16/// The purpose of having this type is to allow lockless (read) access to the underlying data when
17/// there is only one [`RwArc`] instance for the particular allocation (note that there can be any
18/// number of [`RoArc`] instances for that allocation). See the [`RwArc::get`] method for more
19/// details.
20pub struct RwArc<T>(Arc<Inner<T>>);
21
22/// A reference-counting pointer with read-only capabilities.
23///
24/// This type can be created from an existing [`RwArc`] using its [`RwArc::clone_ro`] method. See
25/// the type and method documentation for more details.
26pub struct RoArc<T>(Arc<Inner<T>>);
27
28struct Inner<T> {
29 data: RwLock<T>,
30 num_rw: AtomicUsize,
31}
32
33impl<T> RwArc<T> {
34 /// Creates a new `RwArc<T>`.
35 pub fn new(data: T) -> Self {
36 let inner = Inner {
37 data: RwLock::new(data),
38 num_rw: AtomicUsize::new(1),
39 };
40 Self(Arc::new(inner))
41 }
42
43 /// Acquires the read lock for immutable access.
44 pub fn read(&self) -> RwLockReadGuard<'_, T, PreemptDisabled> {
45 self.0.data.read()
46 }
47
48 /// Acquires the write lock for mutable access.
49 pub fn write(&self) -> RwLockWriteGuard<'_, T, PreemptDisabled> {
50 self.0.data.write()
51 }
52
53 /// Returns a raw pointer to the contained value.
54 pub fn as_ptr(&self) -> *const T {
55 self.0.data.as_ptr()
56 }
57
58 /// Returns an immutable reference if no other `RwArc` points to the same allocation.
59 ///
60 /// This method is cheap because it does not acquire a lock.
61 ///
62 /// It's still sound because:
63 /// - The mutable reference to `self` and the condition ensure that we are exclusively
64 /// accessing the unique `RwArc` instance for the particular allocation.
65 /// - There may be any number of [`RoArc`]s pointing to the same allocation, but they may only
66 /// produce immutable references to the underlying data.
67 pub fn get(&mut self) -> Option<&T> {
68 if self.0.num_rw.load(Ordering::Relaxed) > 1 {
69 return None;
70 }
71
72 // This will synchronize with `RwArc::drop` to make sure its changes are visible to us.
73 fence(Ordering::Acquire);
74
75 let data_ptr = self.0.data.as_ptr();
76
77 // SAFETY: The data is valid. During the lifetime, no one will be able to create a mutable
78 // reference to the data, so it's okay to create an immutable reference like the one below.
79 Some(unsafe { &*data_ptr })
80 }
81
82 /// Clones a [`RoArc`] that points to the same allocation.
83 pub fn clone_ro(&self) -> RoArc<T> {
84 RoArc(self.0.clone())
85 }
86}
87
88impl<T> Clone for RwArc<T> {
89 fn clone(&self) -> Self {
90 let inner = self.0.clone();
91
92 // Note that overflowing the counter will make it unsound. But not to worry: the above
93 // `Arc::clone` must have already aborted the kernel before this happens.
94 inner.num_rw.fetch_add(1, Ordering::Relaxed);
95
96 Self(inner)
97 }
98}
99
100impl<T> Drop for RwArc<T> {
101 fn drop(&mut self) {
102 self.0.num_rw.fetch_sub(1, Ordering::Release);
103 }
104}
105
106impl<T: Clone> RwArc<T> {
107 /// Returns the contained value by cloning it.
108 pub fn get_cloned(&self) -> T {
109 let guard = self.read();
110 guard.clone()
111 }
112}
113
114impl<T> RoArc<T> {
115 /// Acquires the read lock for immutable access.
116 pub fn read(&self) -> RwLockReadGuard<'_, T, PreemptDisabled> {
117 self.0.data.read()
118 }
119}
120
121#[cfg(ktest)]
122mod test {
123 use super::*;
124 use crate::prelude::*;
125
126 #[ktest]
127 fn lockless_get() {
128 let mut rw1 = RwArc::new(1u32);
129 assert_eq!(rw1.get(), Some(1).as_ref());
130
131 let _ro = rw1.clone_ro();
132 assert_eq!(rw1.get(), Some(1).as_ref());
133
134 let rw2 = rw1.clone();
135 assert_eq!(rw1.get(), None);
136
137 drop(rw2);
138 assert_eq!(rw1.get(), Some(1).as_ref());
139 }
140}