ax_task/sync/
rwsem.rs

1//! Reader/writer sleeping locks sharing Linux RT's single-writer PI gate.
2
3use alloc::sync::Arc;
4use core::sync::atomic::{AtomicUsize, Ordering};
5
6use super::{Mutex, RawMutex, RawSpinLock};
7use crate::{
8    runtime::{context::runtime_task_system, sync::rt_lock::RtLockWaitGuard},
9    thread::{
10        ThreadCore, ThreadWakeHandle,
11        current::{self, CurrentParkStart},
12    },
13};
14
15/// Raw implementation of a sleeping reader/writer semaphore.
16///
17/// A writer prevents new readers through its PI mutex, then waits for existing
18/// readers to drain. Existing readers have no single owner to receive PI;
19/// they must finish their critical sections, as with Linux `rwbase_rt`.
20pub struct RawRwSemaphore {
21    gate: Mutex<()>,
22    readers: AtomicUsize,
23    drain: RawSpinLock<Option<DrainWake>>,
24    rt_lock: bool,
25}
26
27enum DrainWake {
28    Ordinary(ThreadWakeHandle),
29    RtLock {
30        core: Arc<ThreadCore>,
31        generation: u64,
32    },
33}
34
35impl DrainWake {
36    fn wake(self) {
37        match self {
38            Self::Ordinary(wake) => {
39                wake.wake();
40            }
41            Self::RtLock { core, generation } => {
42                runtime_task_system()
43                    .expect("reader drain retains its task system")
44                    .wake_rt_lock_park(&core, generation);
45            }
46        }
47    }
48}
49
50impl RawRwSemaphore {
51    /// Creates a sleeping reader/writer semaphore with ordinary task waits.
52    pub const fn new() -> Self {
53        Self::with_wait_state(false)
54    }
55
56    pub(super) const fn with_wait_state(rt_lock: bool) -> Self {
57        Self {
58            gate: Mutex::const_new(
59                if rt_lock {
60                    RawMutex::new_rt_lock()
61                } else {
62                    RawMutex::new()
63                },
64                (),
65            ),
66            readers: AtomicUsize::new(0),
67            drain: RawSpinLock::new(None),
68            rt_lock,
69        }
70    }
71
72    fn add_reader(&self) {
73        self.readers
74            .try_update(Ordering::AcqRel, Ordering::Acquire, |readers| {
75                readers.checked_add(1)
76            })
77            .expect("reader reference count exhausted");
78    }
79
80    fn drain_readers(&self) {
81        if self.readers.load(Ordering::Acquire) == 0 {
82            return;
83        }
84        let _saved_state = self
85            .rt_lock
86            .then(|| RtLockWaitGuard::enter().expect("save RT writer wait state"));
87        loop {
88            let CurrentParkStart::Prepared(park) =
89                current::begin_current_park().expect("prepare reader drain park")
90            else {
91                continue;
92            };
93            let wake = if self.rt_lock {
94                DrainWake::RtLock {
95                    core: current::current_thread_core_arc().expect("current writer"),
96                    generation: park.generation(),
97                }
98            } else {
99                DrainWake::Ordinary(park.wake_handle())
100            };
101            let mut drain = self.drain.lock();
102            if self.readers.load(Ordering::Acquire) == 0 {
103                drop(drain);
104                park.cancel().expect("cancel completed reader drain");
105                return;
106            }
107            assert!(drain.is_none(), "the writer gate owns one drain waiter");
108            *drain = Some(wake);
109            drop(drain);
110            park.commit().expect("commit reader drain park");
111            let stale = self.drain.lock().take();
112            drop(stale);
113            if self.readers.load(Ordering::Acquire) == 0 {
114                return;
115            }
116        }
117    }
118}
119
120impl Default for RawRwSemaphore {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126// SAFETY: every shared acquisition increments readers while holding gate.
127// An exclusive acquisition retains gate and waits for all counted readers to
128// release. Reader release publishes protected reads before the writer's
129// Acquire observation of zero. Gate's PI ownership serializes all writers.
130// GuardNoSend prevents transfer of task-owned unlock authority.
131unsafe impl lock_api::RawRwLock for RawRwSemaphore {
132    const INIT: Self = Self::new();
133    type GuardMarker = lock_api::GuardNoSend;
134
135    fn lock_shared(&self) {
136        let _gate = self.gate.lock();
137        self.add_reader();
138    }
139
140    fn try_lock_shared(&self) -> bool {
141        let Some(_gate) = self.gate.try_lock() else {
142            return false;
143        };
144        self.add_reader();
145        true
146    }
147
148    unsafe fn unlock_shared(&self) {
149        // The caller owns one counted read guard; it cannot underflow.
150        if self.readers.fetch_sub(1, Ordering::AcqRel) == 1 {
151            let wake = self.drain.lock().take();
152            if let Some(wake) = wake {
153                wake.wake();
154            }
155        }
156    }
157
158    fn lock_exclusive(&self) {
159        let gate = self.gate.lock();
160        self.drain_readers();
161        // RawRwLock transfers this established ownership to its caller.
162        core::mem::forget(gate);
163    }
164
165    fn try_lock_exclusive(&self) -> bool {
166        let Some(gate) = self.gate.try_lock() else {
167            return false;
168        };
169        if self.readers.load(Ordering::Acquire) != 0 {
170            return false;
171        }
172        core::mem::forget(gate);
173        true
174    }
175
176    unsafe fn unlock_exclusive(&self) {
177        // SAFETY: the exclusive guard retained gate throughout its lifetime.
178        unsafe {
179            self.gate.force_unlock();
180        }
181    }
182}
183
184/// A sleeping reader/writer semaphore; holding it does not pin the CPU.
185pub type RwSemaphore<T> = lock_api::RwLock<RawRwSemaphore, T>;
186/// A shared, task-bound semaphore guard.
187pub type RwSemaphoreReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, RawRwSemaphore, T>;
188/// An exclusive, task-bound semaphore guard.
189pub type RwSemaphoreWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, RawRwSemaphore, T>;
ax_task/sync/rwsem.rs

ax_task/sync/
rwsem.rs
