rocketmq_client_rust/
lock.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::hint::spin_loop;
16use std::sync::atomic::AtomicBool;
17use std::sync::atomic::AtomicI32;
18use std::sync::atomic::Ordering;
19
20/// Java-compatible CAS read/write lock used by producer backpressure paths.
21///
22/// Java's `ReadWriteCASLock` is a small spin lock backed by one `AtomicBoolean`
23/// for writer admission and one `AtomicInteger` for active readers. This Rust
24/// implementation intentionally mirrors those semantics instead of wrapping a
25/// fair `RwLock`, because Java exposes `getWriteLock` and `getReadLock` as
26/// direct state probes.
27#[derive(Debug)]
28pub struct ReadWriteCASLock {
29    write_lock: AtomicBool,
30    read_lock: AtomicI32,
31}
32
33impl Default for ReadWriteCASLock {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl ReadWriteCASLock {
40    pub fn new() -> Self {
41        Self {
42            write_lock: AtomicBool::new(true),
43            read_lock: AtomicI32::new(0),
44        }
45    }
46
47    pub fn acquire_write_lock(&self) {
48        while self
49            .write_lock
50            .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed)
51            .is_err()
52        {
53            spin_loop();
54        }
55
56        while self.read_lock.load(Ordering::Acquire) != 0 {
57            spin_loop();
58        }
59    }
60
61    pub fn release_write_lock(&self) {
62        let _ = self
63            .write_lock
64            .compare_exchange(false, true, Ordering::Release, Ordering::Relaxed);
65    }
66
67    pub fn acquire_read_lock(&self) {
68        while !self.write_lock.load(Ordering::Acquire) {
69            spin_loop();
70        }
71        self.read_lock.fetch_add(1, Ordering::AcqRel);
72    }
73
74    pub fn release_read_lock(&self) {
75        self.read_lock.fetch_sub(1, Ordering::AcqRel);
76    }
77
78    pub fn get_write_lock(&self) -> bool {
79        self.write_lock.load(Ordering::Acquire) && self.read_lock.load(Ordering::Acquire) == 0
80    }
81
82    pub fn get_read_lock(&self) -> bool {
83        self.write_lock.load(Ordering::Acquire)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use std::sync::atomic::AtomicBool;
90    use std::sync::Arc;
91    use std::sync::Barrier;
92    use std::thread;
93    use std::time::Duration;
94
95    use super::*;
96
97    #[test]
98    fn read_write_cas_lock_state_matches_java_accessors() {
99        let lock = ReadWriteCASLock::new();
100        assert!(lock.get_write_lock());
101        assert!(lock.get_read_lock());
102
103        lock.acquire_read_lock();
104        assert!(!lock.get_write_lock());
105        assert!(lock.get_read_lock());
106
107        lock.release_read_lock();
108        assert!(lock.get_write_lock());
109        assert!(lock.get_read_lock());
110
111        lock.acquire_write_lock();
112        assert!(!lock.get_write_lock());
113        assert!(!lock.get_read_lock());
114
115        lock.release_write_lock();
116        assert!(lock.get_write_lock());
117        assert!(lock.get_read_lock());
118    }
119
120    #[test]
121    fn write_lock_waits_for_active_readers_like_java() {
122        let lock = Arc::new(ReadWriteCASLock::new());
123        lock.acquire_read_lock();
124
125        let writer_started = Arc::new(Barrier::new(2));
126        let writer_acquired = Arc::new(AtomicBool::new(false));
127        let writer_lock = Arc::clone(&lock);
128        let writer_started_inner = Arc::clone(&writer_started);
129        let writer_acquired_inner = Arc::clone(&writer_acquired);
130
131        let handle = thread::spawn(move || {
132            writer_started_inner.wait();
133            writer_lock.acquire_write_lock();
134            writer_acquired_inner.store(true, Ordering::Release);
135            writer_lock.release_write_lock();
136        });
137
138        writer_started.wait();
139        thread::sleep(Duration::from_millis(20));
140        assert!(!writer_acquired.load(Ordering::Acquire));
141
142        lock.release_read_lock();
143        handle.join().expect("writer thread panicked");
144        assert!(writer_acquired.load(Ordering::Acquire));
145        assert!(lock.get_write_lock());
146    }
147}
rocketmq_client_rust/lock.rs

rocketmq_client_rust/
lock.rs
