Skip to main content

relibc/sync/
rwlock.rs

1use core::{
2    cell::UnsafeCell,
3    fmt, ops,
4    sync::atomic::{AtomicU32, Ordering},
5};
6
7use crate::{
8    error::{Errno, Result},
9    header::{
10        errno::{EINVAL, ETIMEDOUT},
11        time::{CLOCK_MONOTONIC, CLOCK_REALTIME, timespec, timespec_realtime_to_monotonic},
12    },
13    platform::types::clockid_t,
14    pthread::Pshared,
15};
16
17pub struct InnerRwLock {
18    state: AtomicU32,
19}
20// PTHREAD_RWLOCK_INITIALIZER is defined as "all zeroes".
21
22const WAITING_WR: u32 = 1 << (u32::BITS - 1);
23const COUNT_MASK: u32 = WAITING_WR - 1;
24const EXCLUSIVE: u32 = COUNT_MASK;
25
26// TODO: Optimize for short waits and long waits, using AtomicLock::wait_until, but still
27// supporting timeouts.
28// TODO: Add futex ops that use bitmasks.
29
30impl InnerRwLock {
31    pub const fn new(_pshared: Pshared) -> Self {
32        Self {
33            state: AtomicU32::new(0),
34        }
35    }
36    fn translate_timeout(deadline: Option<(&timespec, i32)>) -> Result<Option<timespec>, Errno> {
37        let relative = match deadline {
38            // FUTEX expect monotonic clock
39            Some((abstime, CLOCK_MONOTONIC)) => Some(abstime.clone()),
40            Some((abstime, CLOCK_REALTIME)) => Some(timespec_realtime_to_monotonic(abstime)?),
41            None => None,
42            _ => {
43                return Err(Errno(EINVAL));
44            }
45        };
46        Ok(relative)
47    }
48    pub fn acquire_write_lock(
49        &self,
50        deadline: Option<(&timespec, clockid_t)>,
51    ) -> Result<(), Errno> {
52        let relative = Self::translate_timeout(deadline)?;
53        let mut waiting_wr = self.state.load(Ordering::Relaxed) & WAITING_WR;
54
55        loop {
56            match self.state.compare_exchange_weak(
57                waiting_wr,
58                EXCLUSIVE,
59                Ordering::Acquire,
60                Ordering::Relaxed,
61            ) {
62                Ok(_) => break,
63                Err(actual) => {
64                    let expected = actual;
65                    let expected = if actual & COUNT_MASK != EXCLUSIVE {
66                        // Set the exclusive bit, but only if we're waiting for readers, to avoid
67                        // reader starvation by overprioritizing write locks.
68                        self.state.fetch_or(WAITING_WR, Ordering::Relaxed);
69
70                        actual | WAITING_WR
71                    } else {
72                        actual
73                    };
74                    waiting_wr = expected & WAITING_WR;
75
76                    if actual & COUNT_MASK > 0 {
77                        if crate::sync::futex_wait(&self.state, expected, relative.as_ref())
78                            == super::FutexWaitResult::TimedOut
79                        {
80                            return Err(Errno(ETIMEDOUT));
81                        }
82                    } else {
83                        // We must avoid blocking indefinitely in our `futex_wait()`, in this case
84                        // where it's possible that `self.state == expected` but our futex might
85                        // never be woken again, because it's possible that all other threads
86                        // already did their `futex_wake()` before we would've done our
87                        // `futex_wait()`.
88                    }
89                }
90            }
91        }
92
93        Ok(())
94    }
95    pub fn acquire_read_lock(&self, deadline: Option<(&timespec, clockid_t)>) -> Result<(), Errno> {
96        let relative = Self::translate_timeout(deadline)?;
97        while let Err(old) = self.try_acquire_read_lock() {
98            if crate::sync::futex_wait(&self.state, old, relative.as_ref())
99                == super::FutexWaitResult::TimedOut
100            {
101                return Err(Errno(ETIMEDOUT));
102            }
103        }
104
105        Ok(())
106    }
107    pub fn try_acquire_read_lock(&self) -> Result<(), u32> {
108        let mut cached = self.state.load(Ordering::Acquire);
109
110        loop {
111            let waiting_wr = cached & WAITING_WR;
112            let old = if cached & COUNT_MASK == EXCLUSIVE {
113                0
114            } else {
115                cached & COUNT_MASK
116            };
117            let new = old + 1;
118
119            // TODO: Return with error code instead?
120            assert_ne!(
121                new & COUNT_MASK,
122                EXCLUSIVE,
123                "maximum number of rwlock readers reached"
124            );
125
126            match self.state.compare_exchange_weak(
127                (old & COUNT_MASK) | waiting_wr,
128                new | waiting_wr,
129                Ordering::Acquire,
130                Ordering::Relaxed,
131            ) {
132                Ok(_) => return Ok(()),
133
134                Err(value) if value & COUNT_MASK == EXCLUSIVE => return Err(value),
135                Err(value) => {
136                    cached = value;
137                    // TODO: SCHED_YIELD?
138                    core::hint::spin_loop();
139                }
140            }
141        }
142    }
143    pub fn try_acquire_write_lock(&self) -> Result<(), u32> {
144        let mut waiting_wr = self.state.load(Ordering::Relaxed) & WAITING_WR;
145
146        loop {
147            match self.state.compare_exchange_weak(
148                waiting_wr,
149                EXCLUSIVE,
150                Ordering::Acquire,
151                Ordering::Relaxed,
152            ) {
153                Ok(_) => return Ok(()),
154                Err(actual) if actual & COUNT_MASK > 0 => return Err(actual),
155                Err(can_retry) => {
156                    waiting_wr = can_retry & WAITING_WR;
157
158                    core::hint::spin_loop();
159                    continue;
160                }
161            }
162        }
163    }
164
165    pub fn unlock(&self) {
166        let state = self.state.load(Ordering::Relaxed);
167
168        if state & COUNT_MASK == EXCLUSIVE {
169            // Unlocking a write lock.
170
171            // This discards the writer-waiting bit, in order to ensure some level of fairness
172            // between read and write locks.
173            self.state.store(0, Ordering::Release);
174
175            let _ = crate::sync::futex_wake(&self.state, i32::MAX);
176        } else {
177            // Unlocking a read lock. Subtract one from the reader count, but preserve the
178            // WAITING_WR bit.
179
180            if self.state.fetch_sub(1, Ordering::Release) & COUNT_MASK == 1 {
181                let _ = crate::sync::futex_wake(&self.state, i32::MAX);
182            }
183        }
184    }
185}
186
187pub struct RwLock<T: ?Sized> {
188    inner: InnerRwLock,
189    data: UnsafeCell<T>,
190}
191
192unsafe impl<T: ?Sized + Send> Send for RwLock<T> {}
193unsafe impl<T: ?Sized + Send + Sync> Sync for RwLock<T> {}
194
195impl<T> RwLock<T> {
196    pub const fn new(val: T) -> Self {
197        Self {
198            inner: InnerRwLock::new(Pshared::Private),
199            data: UnsafeCell::new(val),
200        }
201    }
202}
203
204impl<T: ?Sized> RwLock<T> {
205    pub fn read(&self) -> ReadGuard<'_, T> {
206        let _ = self.inner.acquire_read_lock(None);
207        unsafe { ReadGuard::new(self) }
208    }
209
210    pub fn write(&self) -> WriteGuard<'_, T> {
211        let _ = self.inner.acquire_write_lock(None);
212        unsafe { WriteGuard::new(self) }
213    }
214
215    pub fn try_read(&self) -> Option<ReadGuard<'_, T>> {
216        if self.inner.try_acquire_read_lock().is_ok() {
217            Some(unsafe { ReadGuard::new(self) })
218        } else {
219            None
220        }
221    }
222
223    pub fn try_write(&self) -> Option<WriteGuard<'_, T>> {
224        if self.inner.try_acquire_write_lock().is_ok() {
225            Some(unsafe { WriteGuard::new(self) })
226        } else {
227            None
228        }
229    }
230}
231
232pub struct ReadGuard<'a, T: ?Sized + 'a> {
233    lock: &'a RwLock<T>,
234}
235
236impl<T: ?Sized> !Send for ReadGuard<'_, T> {}
237unsafe impl<T: ?Sized + Sync> Sync for ReadGuard<'_, T> {}
238
239impl<'a, T: ?Sized> ReadGuard<'a, T> {
240    unsafe fn new(lock: &'a RwLock<T>) -> Self {
241        Self { lock }
242    }
243}
244
245impl<'a, T: ?Sized> ops::Deref for ReadGuard<'a, T> {
246    type Target = T;
247
248    fn deref(&self) -> &Self::Target {
249        // SAFETY: We have shared reference to the data.
250        unsafe { &*self.lock.data.get() }
251    }
252}
253
254impl<'a, T: ?Sized> Drop for ReadGuard<'a, T> {
255    fn drop(&mut self) {
256        self.lock.inner.unlock();
257    }
258}
259
260impl<'a, T: ?Sized + fmt::Debug> fmt::Debug for ReadGuard<'a, T> {
261    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
262        fmt::Debug::fmt(&**self, f)
263    }
264}
265
266impl<'a, T: ?Sized + fmt::Display> fmt::Display for ReadGuard<'a, T> {
267    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
268        fmt::Display::fmt(&**self, f)
269    }
270}
271
272pub struct WriteGuard<'a, T: ?Sized + 'a> {
273    lock: &'a RwLock<T>,
274}
275
276impl<T: ?Sized> !Send for WriteGuard<'_, T> {}
277unsafe impl<T: ?Sized + Sync> Sync for WriteGuard<'_, T> {}
278
279impl<'a, T: ?Sized> WriteGuard<'a, T> {
280    unsafe fn new(lock: &'a RwLock<T>) -> Self {
281        Self { lock }
282    }
283}
284
285impl<'a, T: ?Sized> ops::Deref for WriteGuard<'a, T> {
286    type Target = T;
287
288    fn deref(&self) -> &Self::Target {
289        // SAFETY: We have exclusive reference to the data.
290        unsafe { &*self.lock.data.get() }
291    }
292}
293
294impl<'a, T: ?Sized> ops::DerefMut for WriteGuard<'a, T> {
295    fn deref_mut(&mut self) -> &mut Self::Target {
296        // SAFETY: We have exclusive reference to the data.
297        unsafe { &mut *self.lock.data.get() }
298    }
299}
300
301impl<'a, T: ?Sized> Drop for WriteGuard<'a, T> {
302    fn drop(&mut self) {
303        self.lock.inner.unlock();
304    }
305}
306
307impl<'a, T: ?Sized + fmt::Debug> fmt::Debug for WriteGuard<'a, T> {
308    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
309        fmt::Debug::fmt(&**self, f)
310    }
311}
312
313impl<'a, T: ?Sized + fmt::Display> fmt::Display for WriteGuard<'a, T> {
314    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
315        fmt::Display::fmt(&**self, f)
316    }
317}