rocketmq_client_rust/consumer/consumer_impl/
pop_process_queue.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::fmt::Display;
16use std::hash::Hash;
17use std::sync::atomic::AtomicBool;
18use std::sync::atomic::AtomicI32;
19use std::sync::atomic::AtomicU64;
20use std::sync::atomic::Ordering;
21use std::sync::atomic::Ordering::Acquire;
22use std::sync::atomic::Ordering::Relaxed;
23use std::sync::Arc;
24
25use rocketmq_common::TimeUtils::current_millis;
26use rocketmq_remoting::protocol::body::pop_process_queue_info::PopProcessQueueInfo;
27
28use crate::consumer::consumer_impl::PULL_MAX_IDLE_TIME;
29
30#[derive(Clone)]
31pub struct PopProcessQueue {
32    last_pop_timestamp: Arc<AtomicU64>,
33    wait_ack_counter: Arc<AtomicI32>,
34    dropped: Arc<AtomicBool>,
35}
36
37impl Hash for PopProcessQueue {
38    #[inline]
39    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
40        self.last_pop_timestamp.load(Relaxed).hash(state);
41        self.wait_ack_counter.load(Relaxed).hash(state);
42        self.dropped.load(Relaxed).hash(state);
43    }
44}
45
46impl Default for PopProcessQueue {
47    fn default() -> Self {
48        PopProcessQueue {
49            last_pop_timestamp: Arc::new(AtomicU64::new(current_millis())),
50            wait_ack_counter: Arc::new(AtomicI32::new(0)),
51            dropped: Arc::new(AtomicBool::new(false)),
52        }
53    }
54}
55
56impl PopProcessQueue {
57    #[inline]
58    pub fn new() -> Self {
59        PopProcessQueue::default()
60    }
61
62    #[inline]
63    pub fn get_last_pop_timestamp(&self) -> u64 {
64        self.last_pop_timestamp.load(Ordering::Relaxed)
65    }
66
67    #[inline]
68    pub fn set_last_pop_timestamp(&self, last_pop_timestamp: u64) {
69        self.last_pop_timestamp.store(last_pop_timestamp, Ordering::Release);
70    }
71
72    #[inline]
73    pub fn inc_found_msg(&self, count: usize) {
74        self.wait_ack_counter
75            .fetch_add(count.min(i32::MAX as usize) as i32, Ordering::AcqRel);
76    }
77
78    #[inline]
79    pub fn ack(&self) -> i32 {
80        self.wait_ack_counter.fetch_sub(1, Ordering::AcqRel)
81    }
82
83    #[inline]
84    pub fn dec_found_msg(&self, count: usize) {
85        self.wait_ack_counter
86            .fetch_sub(count.min(i32::MAX as usize) as i32, Ordering::AcqRel);
87    }
88
89    #[inline]
90    pub fn get_wai_ack_msg_count(&self) -> i32 {
91        self.wait_ack_counter.load(Ordering::Acquire)
92    }
93
94    #[inline]
95    pub fn is_dropped(&self) -> bool {
96        self.dropped.load(Ordering::Acquire)
97    }
98
99    #[inline]
100    pub fn set_dropped(&self, dropped: bool) {
101        self.dropped.store(dropped, Ordering::Release);
102    }
103
104    #[inline]
105    pub fn fill_pop_process_queue_info(&self, info: &mut PopProcessQueueInfo) {
106        info.set_wait_ack_count(self.get_wai_ack_msg_count());
107        info.set_droped(self.is_dropped());
108        info.set_last_pop_timestamp(self.get_last_pop_timestamp());
109    }
110
111    #[inline]
112    pub fn is_pull_expired(&self) -> bool {
113        let current_time = current_millis();
114        current_time.saturating_sub(self.last_pop_timestamp.load(Acquire)) > *PULL_MAX_IDLE_TIME
115    }
116}
117
118impl Display for PopProcessQueue {
119    #[inline]
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        write!(
122            f,
123            "PopProcessQueue[waitAckCounter:{}, lastPopTimestamp:{}, drop:{}]",
124            self.wait_ack_counter.load(Ordering::Relaxed),
125            self.last_pop_timestamp.load(Ordering::Relaxed),
126            self.dropped.load(Ordering::Relaxed)
127        )
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn pop_process_queue_initializes_correctly() {
137        let queue = PopProcessQueue::new();
138        assert!(queue.get_last_pop_timestamp() > 0);
139        assert_eq!(queue.get_wai_ack_msg_count(), 0);
140        assert!(!queue.is_dropped());
141    }
142
143    #[test]
144    fn pop_process_queue_updates_last_pop_timestamp() {
145        let queue = PopProcessQueue::new();
146        let new_timestamp = queue.get_last_pop_timestamp() + 1000;
147        queue.set_last_pop_timestamp(new_timestamp);
148        assert_eq!(queue.get_last_pop_timestamp(), new_timestamp);
149    }
150
151    #[test]
152    fn pop_process_queue_increments_found_msg() {
153        let queue = PopProcessQueue::new();
154        queue.inc_found_msg(5);
155        assert_eq!(queue.get_wai_ack_msg_count(), 5);
156    }
157
158    #[test]
159    fn pop_process_queue_decrements_found_msg() {
160        let queue = PopProcessQueue::new();
161        queue.inc_found_msg(5);
162        queue.dec_found_msg(3);
163        assert_eq!(queue.get_wai_ack_msg_count(), 2);
164    }
165
166    #[test]
167    fn pop_process_queue_acknowledges_msg() {
168        let queue = PopProcessQueue::new();
169        queue.inc_found_msg(5);
170        assert_eq!(queue.ack(), 5);
171        assert_eq!(queue.get_wai_ack_msg_count(), 4);
172    }
173
174    #[test]
175    fn pop_process_queue_ack_at_zero_matches_java_signed_counter() {
176        let queue = PopProcessQueue::new();
177        assert_eq!(queue.ack(), 0);
178        assert_eq!(queue.get_wai_ack_msg_count(), -1);
179    }
180
181    #[test]
182    fn pop_process_queue_sets_dropped_flag() {
183        let queue = PopProcessQueue::new();
184        queue.set_dropped(true);
185        assert!(queue.is_dropped());
186    }
187
188    #[test]
189    fn pop_process_queue_detects_pull_expired() {
190        let queue = PopProcessQueue::new();
191        queue.set_last_pop_timestamp(queue.get_last_pop_timestamp() - *PULL_MAX_IDLE_TIME - 1);
192        assert!(queue.is_pull_expired());
193    }
194
195    #[test]
196    fn pop_process_queue_info_preserves_negative_wait_ack_count_like_java() {
197        let queue = PopProcessQueue::new();
198        queue.ack();
199        let mut info = PopProcessQueueInfo::new(0, false, 0);
200        queue.fill_pop_process_queue_info(&mut info);
201        assert_eq!(info.wait_ack_count(), -1);
202    }
203}
rocketmq_client_rust/consumer/consumer_impl/pop_process_queue.rs

rocketmq_client_rust/consumer/consumer_impl/
pop_process_queue.rs
