|
| 1 | +use crate::sync::atomic::{ |
| 2 | + AtomicI32, |
| 3 | + Ordering::{Acquire, Relaxed, Release}, |
| 4 | +}; |
| 5 | +use crate::sys::futex::{futex_wait, futex_wake, futex_wake_all}; |
| 6 | + |
| 7 | +pub type MovableRwLock = RwLock; |
| 8 | + |
| 9 | +pub struct RwLock { |
| 10 | + // The state consists of a 30-bit reader counter, a 'readers waiting' flag, and a 'writers waiting' flag. |
| 11 | + // Bits 0..30: |
| 12 | + // 0: Unlocked |
| 13 | + // 1..=0x3FFF_FFFE: Locked by N readers |
| 14 | + // 0x3FFF_FFFF: Write locked |
| 15 | + // Bit 30: Readers are waiting on this futex. |
| 16 | + // Bit 31: Writers are waiting on the writer_notify futex. |
| 17 | + state: AtomicI32, |
| 18 | + // The 'condition variable' to notify writers through. |
| 19 | + // Incremented on every signal. |
| 20 | + writer_notify: AtomicI32, |
| 21 | +} |
| 22 | + |
| 23 | +const READ_LOCKED: i32 = 1; |
| 24 | +const MASK: i32 = (1 << 30) - 1; |
| 25 | +const WRITE_LOCKED: i32 = MASK; |
| 26 | +const MAX_READERS: i32 = MASK - 1; |
| 27 | +const READERS_WAITING: i32 = 1 << 30; |
| 28 | +const WRITERS_WAITING: i32 = 1 << 31; |
| 29 | + |
| 30 | +fn is_unlocked(state: i32) -> bool { |
| 31 | + state & MASK == 0 |
| 32 | +} |
| 33 | + |
| 34 | +fn is_write_locked(state: i32) -> bool { |
| 35 | + state & MASK == WRITE_LOCKED |
| 36 | +} |
| 37 | + |
| 38 | +fn has_readers_waiting(state: i32) -> bool { |
| 39 | + state & READERS_WAITING != 0 |
| 40 | +} |
| 41 | + |
| 42 | +fn has_writers_waiting(state: i32) -> bool { |
| 43 | + state & WRITERS_WAITING != 0 |
| 44 | +} |
| 45 | + |
| 46 | +fn is_read_lockable(state: i32) -> bool { |
| 47 | + // This also returns false if the counter could overflow if we tried to read lock it. |
| 48 | + // |
| 49 | + // We don't allow read-locking if there's readers waiting, even if the lock is unlocked |
| 50 | + // and there's no writers waiting. The only situation when this happens is after unlocking, |
| 51 | + // at which point the unlocking thread might be waking up writers, which have priority over readers. |
| 52 | + // The unlocking thread will clear the readers waiting bit and wake up readers, if necssary. |
| 53 | + state & MASK < MAX_READERS && !has_readers_waiting(state) && !has_writers_waiting(state) |
| 54 | +} |
| 55 | + |
| 56 | +fn has_reached_max_readers(state: i32) -> bool { |
| 57 | + state & MASK == MAX_READERS |
| 58 | +} |
| 59 | + |
| 60 | +impl RwLock { |
| 61 | + #[inline] |
| 62 | + pub const fn new() -> Self { |
| 63 | + Self { state: AtomicI32::new(0), writer_notify: AtomicI32::new(0) } |
| 64 | + } |
| 65 | + |
| 66 | + #[inline] |
| 67 | + pub unsafe fn destroy(&self) {} |
| 68 | + |
| 69 | + #[inline] |
| 70 | + pub unsafe fn try_read(&self) -> bool { |
| 71 | + self.state |
| 72 | + .fetch_update(Acquire, Relaxed, |s| is_read_lockable(s).then(|| s + READ_LOCKED)) |
| 73 | + .is_ok() |
| 74 | + } |
| 75 | + |
| 76 | + #[inline] |
| 77 | + pub unsafe fn read(&self) { |
| 78 | + let state = self.state.load(Relaxed); |
| 79 | + if !is_read_lockable(state) |
| 80 | + || self |
| 81 | + .state |
| 82 | + .compare_exchange_weak(state, state + READ_LOCKED, Acquire, Relaxed) |
| 83 | + .is_err() |
| 84 | + { |
| 85 | + self.read_contended(); |
| 86 | + } |
| 87 | + } |
| 88 | + |
| 89 | + #[inline] |
| 90 | + pub unsafe fn read_unlock(&self) { |
| 91 | + let state = self.state.fetch_sub(READ_LOCKED, Release) - READ_LOCKED; |
| 92 | + |
| 93 | + // It's impossible for a reader to be waiting on a read-locked RwLock, |
| 94 | + // except if there is also a writer waiting. |
| 95 | + debug_assert!(!has_readers_waiting(state) || has_writers_waiting(state)); |
| 96 | + |
| 97 | + // Wake up a writer if we were the last reader and there's a writer waiting. |
| 98 | + if is_unlocked(state) && has_writers_waiting(state) { |
| 99 | + self.wake_writer_or_readers(state); |
| 100 | + } |
| 101 | + } |
| 102 | + |
| 103 | + #[cold] |
| 104 | + fn read_contended(&self) { |
| 105 | + let mut state = self.spin_read(); |
| 106 | + |
| 107 | + loop { |
| 108 | + // If we can lock it, lock it. |
| 109 | + if is_read_lockable(state) { |
| 110 | + match self.state.compare_exchange_weak(state, state + READ_LOCKED, Acquire, Relaxed) |
| 111 | + { |
| 112 | + Ok(_) => return, // Locked! |
| 113 | + Err(s) => { |
| 114 | + state = s; |
| 115 | + continue; |
| 116 | + } |
| 117 | + } |
| 118 | + } |
| 119 | + |
| 120 | + // Check for overflow. |
| 121 | + if has_reached_max_readers(state) { |
| 122 | + panic!("too many active read locks on RwLock"); |
| 123 | + } |
| 124 | + |
| 125 | + // Make sure the readers waiting bit is set before we go to sleep. |
| 126 | + if !has_readers_waiting(state) { |
| 127 | + if let Err(s) = |
| 128 | + self.state.compare_exchange(state, state | READERS_WAITING, Relaxed, Relaxed) |
| 129 | + { |
| 130 | + state = s; |
| 131 | + continue; |
| 132 | + } |
| 133 | + } |
| 134 | + |
| 135 | + // Wait for the state to change. |
| 136 | + futex_wait(&self.state, state | READERS_WAITING, None); |
| 137 | + |
| 138 | + // Spin again after waking up. |
| 139 | + state = self.spin_read(); |
| 140 | + } |
| 141 | + } |
| 142 | + |
| 143 | + #[inline] |
| 144 | + pub unsafe fn try_write(&self) -> bool { |
| 145 | + self.state |
| 146 | + .fetch_update(Acquire, Relaxed, |s| is_unlocked(s).then(|| s + WRITE_LOCKED)) |
| 147 | + .is_ok() |
| 148 | + } |
| 149 | + |
| 150 | + #[inline] |
| 151 | + pub unsafe fn write(&self) { |
| 152 | + if self.state.compare_exchange_weak(0, WRITE_LOCKED, Acquire, Relaxed).is_err() { |
| 153 | + self.write_contended(); |
| 154 | + } |
| 155 | + } |
| 156 | + |
| 157 | + #[inline] |
| 158 | + pub unsafe fn write_unlock(&self) { |
| 159 | + let state = self.state.fetch_sub(WRITE_LOCKED, Release) - WRITE_LOCKED; |
| 160 | + |
| 161 | + debug_assert!(is_unlocked(state)); |
| 162 | + |
| 163 | + if has_writers_waiting(state) || has_readers_waiting(state) { |
| 164 | + self.wake_writer_or_readers(state); |
| 165 | + } |
| 166 | + } |
| 167 | + |
| 168 | + #[cold] |
| 169 | + fn write_contended(&self) { |
| 170 | + let mut state = self.spin_write(); |
| 171 | + |
| 172 | + let mut other_writers_waiting = 0; |
| 173 | + |
| 174 | + loop { |
| 175 | + // If it's unlocked, we try to lock it. |
| 176 | + if is_unlocked(state) { |
| 177 | + match self.state.compare_exchange_weak( |
| 178 | + state, |
| 179 | + state | WRITE_LOCKED | other_writers_waiting, |
| 180 | + Acquire, |
| 181 | + Relaxed, |
| 182 | + ) { |
| 183 | + Ok(_) => return, // Locked! |
| 184 | + Err(s) => { |
| 185 | + state = s; |
| 186 | + continue; |
| 187 | + } |
| 188 | + } |
| 189 | + } |
| 190 | + |
| 191 | + // Set the waiting bit indicating that we're waiting on it. |
| 192 | + if !has_writers_waiting(state) { |
| 193 | + if let Err(s) = |
| 194 | + self.state.compare_exchange(state, state | WRITERS_WAITING, Relaxed, Relaxed) |
| 195 | + { |
| 196 | + state = s; |
| 197 | + continue; |
| 198 | + } |
| 199 | + } |
| 200 | + |
| 201 | + // Other writers might be waiting now too, so we should make sure |
| 202 | + // we keep that bit on once we manage lock it. |
| 203 | + other_writers_waiting = WRITERS_WAITING; |
| 204 | + |
| 205 | + // Examine the notification counter before we check if `state` has changed, |
| 206 | + // to make sure we don't miss any notifications. |
| 207 | + let seq = self.writer_notify.load(Acquire); |
| 208 | + |
| 209 | + // Don't go to sleep if the lock has become available, |
| 210 | + // or if the writers waiting bit is no longer set. |
| 211 | + let s = self.state.load(Relaxed); |
| 212 | + if is_unlocked(state) || !has_writers_waiting(s) { |
| 213 | + state = s; |
| 214 | + continue; |
| 215 | + } |
| 216 | + |
| 217 | + // Wait for the state to change. |
| 218 | + futex_wait(&self.writer_notify, seq, None); |
| 219 | + |
| 220 | + // Spin again after waking up. |
| 221 | + state = self.spin_write(); |
| 222 | + } |
| 223 | + } |
| 224 | + |
| 225 | + /// Wake up waiting threads after unlocking. |
| 226 | + /// |
| 227 | + /// If both are waiting, this will wake up only one writer, but will fall |
| 228 | + /// back to waking up readers if there was no writer to wake up. |
| 229 | + #[cold] |
| 230 | + fn wake_writer_or_readers(&self, mut state: i32) { |
| 231 | + assert!(is_unlocked(state)); |
| 232 | + |
| 233 | + // The readers waiting bit might be turned on at any point now, |
| 234 | + // since readers will block when there's anything waiting. |
| 235 | + // Writers will just lock the lock though, regardless of the waiting bits, |
| 236 | + // so we don't have to worry about the writer waiting bit. |
| 237 | + // |
| 238 | + // If the lock gets locked in the meantime, we don't have to do |
| 239 | + // anything, because then the thread that locked the lock will take |
| 240 | + // care of waking up waiters when it unlocks. |
| 241 | + |
| 242 | + // If only writers are waiting, wake one of them up. |
| 243 | + if state == WRITERS_WAITING { |
| 244 | + match self.state.compare_exchange(state, 0, Relaxed, Relaxed) { |
| 245 | + Ok(_) => { |
| 246 | + self.wake_writer(); |
| 247 | + return; |
| 248 | + } |
| 249 | + Err(s) => { |
| 250 | + // Maybe some readers are now waiting too. So, continue to the next `if`. |
| 251 | + state = s; |
| 252 | + } |
| 253 | + } |
| 254 | + } |
| 255 | + |
| 256 | + // If both writers and readers are waiting, leave the readers waiting |
| 257 | + // and only wake up one writer. |
| 258 | + if state == READERS_WAITING + WRITERS_WAITING { |
| 259 | + if self.state.compare_exchange(state, READERS_WAITING, Relaxed, Relaxed).is_err() { |
| 260 | + // The lock got locked. Not our problem anymore. |
| 261 | + return; |
| 262 | + } |
| 263 | + if self.wake_writer() { |
| 264 | + return; |
| 265 | + } |
| 266 | + // No writers were actually blocked on futex_wait, so we continue |
| 267 | + // to wake up readers instead, since we can't be sure if we notified a writer. |
| 268 | + state = READERS_WAITING; |
| 269 | + } |
| 270 | + |
| 271 | + // If readers are waiting, wake them all up. |
| 272 | + if state == READERS_WAITING { |
| 273 | + if self.state.compare_exchange(state, 0, Relaxed, Relaxed).is_ok() { |
| 274 | + futex_wake_all(&self.state); |
| 275 | + } |
| 276 | + } |
| 277 | + } |
| 278 | + |
| 279 | + /// This wakes one writer and returns true if we woke up a writer that was |
| 280 | + /// blocked on futex_wait. |
| 281 | + /// |
| 282 | + /// If this returns false, it might still be the case that we notified a |
| 283 | + /// writer that was about to go to sleep. |
| 284 | + fn wake_writer(&self) -> bool { |
| 285 | + self.writer_notify.fetch_add(1, Release); |
| 286 | + futex_wake(&self.writer_notify) |
| 287 | + } |
| 288 | + |
| 289 | + /// Spin for a while, but stop directly at the given condition. |
| 290 | + fn spin_until(&self, f: impl Fn(i32) -> bool) -> i32 { |
| 291 | + let mut spin = 100; // Chosen by fair dice roll. |
| 292 | + loop { |
| 293 | + let state = self.state.load(Relaxed); |
| 294 | + if f(state) || spin == 0 { |
| 295 | + return state; |
| 296 | + } |
| 297 | + crate::hint::spin_loop(); |
| 298 | + spin -= 1; |
| 299 | + } |
| 300 | + } |
| 301 | + |
| 302 | + fn spin_write(&self) -> i32 { |
| 303 | + // Stop spinning when it's unlocked or when there's waiting writers, to keep things somewhat fair. |
| 304 | + self.spin_until(|state| is_unlocked(state) || has_writers_waiting(state)) |
| 305 | + } |
| 306 | + |
| 307 | + fn spin_read(&self) -> i32 { |
| 308 | + // Stop spinning when it's unlocked or read locked, or when there's waiting threads. |
| 309 | + self.spin_until(|state| { |
| 310 | + !is_write_locked(state) || has_readers_waiting(state) || has_writers_waiting(state) |
| 311 | + }) |
| 312 | + } |
| 313 | +} |
0 commit comments