rocketmq_client_rust/consumer/consumer_impl/
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::collections::BTreeMap;
16use std::sync::atomic::AtomicBool;
17use std::sync::atomic::AtomicI64;
18use std::sync::atomic::AtomicU64;
19use std::sync::atomic::Ordering;
20use std::sync::Arc;
21use std::sync::LazyLock;
22
23use cheetah_string::CheetahString;
24use rocketmq_common::common::message::message_ext::MessageExt;
25use rocketmq_common::common::message::MessageConst;
26use rocketmq_common::common::message::MessageTrait;
27use rocketmq_common::MessageAccessor::MessageAccessor;
28use rocketmq_common::TimeUtils::current_millis;
29use rocketmq_remoting::protocol::body::process_queue_info::ProcessQueueInfo;
30use rocketmq_rust::ArcMut;
31use tokio::sync::RwLock;
32use tracing::info;
33
34use crate::consumer::consumer_impl::default_mq_push_consumer_impl::DefaultMQPushConsumerImpl;
35use crate::consumer::consumer_impl::PULL_MAX_IDLE_TIME;
36
37pub static REBALANCE_LOCK_MAX_LIVE_TIME: LazyLock<u64> = LazyLock::new(|| {
38    std::env::var("rocketmq.client.rebalance.lockMaxLiveTime")
39        .unwrap_or_else(|_| "30000".into())
40        .parse()
41        .unwrap_or(30000)
42});
43
44pub static REBALANCE_LOCK_INTERVAL: LazyLock<u64> = LazyLock::new(|| {
45    std::env::var("rocketmq.client.rebalance.lockInterval")
46        .unwrap_or_else(|_| "20000".into())
47        .parse()
48        .unwrap_or(20000)
49});
50
51struct ProcessQueueStore {
52    msg_tree_map: BTreeMap<i64, ArcMut<MessageExt>>,
53    consuming_msg_orderly_tree_map: BTreeMap<i64, ArcMut<MessageExt>>,
54    queue_offset_max: i64,
55}
56
57impl ProcessQueueStore {
58    fn new() -> Self {
59        ProcessQueueStore {
60            msg_tree_map: BTreeMap::new(),
61            consuming_msg_orderly_tree_map: BTreeMap::new(),
62            queue_offset_max: 0,
63        }
64    }
65}
66
67pub struct ProcessQueue {
68    store: RwLock<ProcessQueueStore>,
69    pub(crate) consume_lock: Arc<RwLock<()>>,
70
71    msg_count: AtomicI64,
72    msg_size: AtomicU64,
73    msg_acc_cnt: AtomicI64,
74    try_unlock_times: AtomicI64,
75
76    dropped: AtomicBool,
77    locked: AtomicBool,
78    consuming: AtomicBool,
79
80    last_pull_timestamp: AtomicU64,
81    last_consume_timestamp: AtomicU64,
82    last_lock_timestamp: AtomicU64,
83}
84
85impl Default for ProcessQueue {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91impl ProcessQueue {
92    pub fn new() -> Self {
93        let now = current_millis();
94        ProcessQueue {
95            store: RwLock::new(ProcessQueueStore::new()),
96            consume_lock: Arc::new(RwLock::new(())),
97
98            msg_count: AtomicI64::new(0),
99            msg_size: AtomicU64::new(0),
100            msg_acc_cnt: AtomicI64::new(0),
101            try_unlock_times: AtomicI64::new(0),
102
103            dropped: AtomicBool::new(false),
104            locked: AtomicBool::new(false),
105            consuming: AtomicBool::new(false),
106
107            last_pull_timestamp: AtomicU64::new(now),
108            last_consume_timestamp: AtomicU64::new(now),
109            last_lock_timestamp: AtomicU64::new(now),
110        }
111    }
112}
113
114impl ProcessQueue {
115    pub(crate) fn set_dropped(&self, dropped: bool) {
116        self.dropped.store(dropped, Ordering::Release);
117    }
118
119    pub(crate) fn is_dropped(&self) -> bool {
120        self.dropped.load(Ordering::Acquire)
121    }
122
123    pub(crate) fn get_last_lock_timestamp(&self) -> u64 {
124        self.last_lock_timestamp.load(Ordering::Acquire)
125    }
126
127    pub(crate) fn set_locked(&self, locked: bool) {
128        self.locked.store(locked, Ordering::Release);
129    }
130
131    pub(crate) fn is_pull_expired(&self) -> bool {
132        (current_millis() - self.last_pull_timestamp.load(Ordering::Acquire)) > *PULL_MAX_IDLE_TIME
133    }
134
135    pub(crate) fn is_lock_expired(&self) -> bool {
136        (current_millis() - self.last_lock_timestamp.load(Ordering::Acquire)) > *REBALANCE_LOCK_MAX_LIVE_TIME
137    }
138
139    pub(crate) fn inc_try_unlock_times(&self) {
140        self.try_unlock_times.fetch_add(1, Ordering::AcqRel);
141    }
142
143    pub(crate) async fn clean_expired_msg(&self, push_consumer: Option<ArcMut<DefaultMQPushConsumerImpl>>) {
144        let mut push_consumer = match push_consumer {
145            Some(pc) => pc,
146            None => return,
147        };
148
149        if push_consumer.is_consume_orderly() {
150            return;
151        }
152
153        let loop_ = 16.min(self.store.read().await.msg_tree_map.len());
154        for _ in 0..loop_ {
155            let msg = {
156                let store = self.store.read().await;
157                if let Some((_, value)) = store.msg_tree_map.first_key_value() {
158                    let consume_start_time_stamp = MessageAccessor::get_consume_start_time_stamp(value.as_ref());
159                    if let Some(ts_str) = consume_start_time_stamp {
160                        if let Ok(ts) = ts_str.parse::<u64>() {
161                            if current_millis() - ts > push_consumer.consumer_config.consume_timeout * 1000 * 60 {
162                                Some(value.clone())
163                            } else {
164                                None
165                            }
166                        } else {
167                            None
168                        }
169                    } else {
170                        None
171                    }
172                } else {
173                    None
174                }
175            };
176
177            let mut msg = match msg {
178                Some(m) => m,
179                None => break,
180            };
181
182            let msg_inner = msg.as_mut();
183            let topic = push_consumer.client_config.with_namespace(msg_inner.topic());
184            let queue_id = msg_inner.queue_id();
185            let msg_id = msg_inner.msg_id().to_string();
186            let store_host = format!("{:?}", msg_inner.store_host());
187            let queue_offset = msg_inner.queue_offset;
188            msg_inner.set_topic(topic);
189            let topic_name = msg_inner.topic().to_string();
190            let _ = push_consumer
191                .send_message_back_with_broker_name(msg_inner, 3, None, None)
192                .await;
193            info!(
194                "send expire msg back. topic={}, msgId={}, storeHost={}, queueId={}, queueOffset={}",
195                topic_name, msg_id, store_host, queue_id, queue_offset
196            );
197
198            let should_remove = {
199                let store = self.store.read().await;
200                store
201                    .msg_tree_map
202                    .first_key_value()
203                    .is_some_and(|(offset, _)| queue_offset == *offset)
204            };
205            if should_remove {
206                self.remove_message(&[msg]).await;
207            }
208        }
209    }
210
211    pub(crate) async fn put_message(&self, messages: &[ArcMut<MessageExt>]) -> bool {
212        let mut dispatch_to_consume = false;
213        let mut store = self.store.write().await;
214        let mut valid_msg_cnt = 0i64;
215
216        let acc_total = if let Some(last_msg) = messages.last() {
217            if let Some(property) =
218                last_msg.property(&CheetahString::from_static_str(MessageConst::PROPERTY_MAX_OFFSET))
219            {
220                property.parse::<i64>().unwrap_or(0) - last_msg.queue_offset
221            } else {
222                0
223            }
224        } else {
225            0
226        };
227
228        for message in messages {
229            if store
230                .msg_tree_map
231                .insert(message.queue_offset, message.clone())
232                .is_none()
233            {
234                valid_msg_cnt += 1;
235                if message.queue_offset > store.queue_offset_max {
236                    store.queue_offset_max = message.queue_offset;
237                }
238                let body_len = message.body().map_or(0, |b| b.len()) as u64;
239                self.msg_size.fetch_add(body_len, Ordering::AcqRel);
240            }
241        }
242        self.msg_count.fetch_add(valid_msg_cnt, Ordering::AcqRel);
243        if !store.msg_tree_map.is_empty() && !self.consuming.load(Ordering::Acquire) {
244            dispatch_to_consume = true;
245            self.consuming.store(true, Ordering::Release);
246        }
247        if acc_total > 0 {
248            self.msg_acc_cnt.store(acc_total, Ordering::Release);
249        }
250        dispatch_to_consume
251    }
252
253    pub(crate) async fn get_max_span(&self) -> u64 {
254        let store = self.store.read().await;
255        match (
256            store.msg_tree_map.first_key_value(),
257            store.msg_tree_map.last_key_value(),
258        ) {
259            (Some((first, _)), Some((last, _))) => (last - first) as u64,
260            _ => 0,
261        }
262    }
263
264    pub(crate) async fn remove_message(&self, messages: &[ArcMut<MessageExt>]) -> i64 {
265        let now = current_millis();
266        let mut store = self.store.write().await;
267
268        self.last_consume_timestamp.store(now, Ordering::Release);
269
270        if store.msg_tree_map.is_empty() {
271            return -1;
272        }
273
274        let mut result = store.queue_offset_max + 1;
275        let mut removed_cnt = 0i64;
276
277        for message in messages {
278            if store.msg_tree_map.remove(&message.queue_offset).is_some() {
279                removed_cnt += 1;
280                let body_len = message.body().map_or(0, |b| b.len()) as u64;
281                if body_len > 0 {
282                    self.msg_size.fetch_sub(body_len, Ordering::AcqRel);
283                }
284            }
285        }
286
287        if removed_cnt > 0 {
288            let prev = self.msg_count.fetch_sub(removed_cnt, Ordering::AcqRel);
289            if prev == removed_cnt {
290                self.msg_size.store(0, Ordering::Release);
291            }
292        }
293
294        if let Some((first_offset, _)) = store.msg_tree_map.first_key_value() {
295            result = *first_offset;
296        }
297
298        result
299    }
300
301    pub(crate) async fn rollback(&self) {
302        let mut store = self.store.write().await;
303        let mut consuming = std::mem::take(&mut store.consuming_msg_orderly_tree_map);
304        store.msg_tree_map.append(&mut consuming);
305    }
306
307    pub(crate) async fn commit(&self) -> i64 {
308        let mut store = self.store.write().await;
309        let offset = store
310            .consuming_msg_orderly_tree_map
311            .last_key_value()
312            .map_or(-1, |(k, _)| *k + 1);
313
314        let consumed_count = store.consuming_msg_orderly_tree_map.len() as i64;
315        let prev_count = self.msg_count.fetch_sub(consumed_count, Ordering::AcqRel);
316
317        if prev_count == consumed_count {
318            self.msg_size.store(0, Ordering::Release);
319        } else {
320            for message in store.consuming_msg_orderly_tree_map.values() {
321                let body_len = message.body().map_or(0, |b| b.len()) as u64;
322                if body_len > 0 {
323                    self.msg_size.fetch_sub(body_len, Ordering::AcqRel);
324                }
325            }
326        }
327        store.consuming_msg_orderly_tree_map.clear();
328        offset
329    }
330
331    pub(crate) async fn make_message_to_consume_again(&self, messages: &[ArcMut<MessageExt>]) {
332        let mut store = self.store.write().await;
333        for message in messages {
334            store.consuming_msg_orderly_tree_map.remove(&message.queue_offset);
335            store.msg_tree_map.insert(message.queue_offset, message.clone());
336        }
337    }
338
339    pub(crate) async fn take_messages(&self, batch_size: u32) -> Vec<ArcMut<MessageExt>> {
340        let mut messages = Vec::with_capacity(batch_size as usize);
341        let now = current_millis();
342        let mut store = self.store.write().await;
343
344        self.last_consume_timestamp.store(now, Ordering::Release);
345
346        for _ in 0..batch_size {
347            if let Some((offset, message)) = store.msg_tree_map.pop_first() {
348                store.consuming_msg_orderly_tree_map.insert(offset, message.clone());
349                messages.push(message);
350            } else {
351                break;
352            }
353        }
354
355        if messages.is_empty() {
356            self.consuming.store(false, Ordering::Release);
357        }
358
359        messages
360    }
361
362    pub(crate) async fn contains_message(&self, message_ext: &MessageExt) -> bool {
363        self.store
364            .read()
365            .await
366            .msg_tree_map
367            .contains_key(&message_ext.queue_offset)
368    }
369
370    pub(crate) async fn clear(&self) {
371        let mut store = self.store.write().await;
372        store.msg_tree_map.clear();
373        store.consuming_msg_orderly_tree_map.clear();
374        store.queue_offset_max = 0;
375        self.msg_count.store(0, Ordering::Release);
376        self.msg_size.store(0, Ordering::Release);
377        self.msg_acc_cnt.store(0, Ordering::Release);
378        self.consuming.store(false, Ordering::Release);
379    }
380
381    pub(crate) async fn fill_process_queue_info(&self, info: &mut ProcessQueueInfo) {
382        let store = self.store.read().await;
383
384        if !store.msg_tree_map.is_empty() {
385            if let Some((min_offset, _)) = store.msg_tree_map.first_key_value() {
386                info.cached_msg_min_offset = *min_offset as u64;
387            }
388            if let Some((max_offset, _)) = store.msg_tree_map.last_key_value() {
389                info.cached_msg_max_offset = *max_offset as u64;
390            }
391            info.cached_msg_count = store.msg_tree_map.len() as u32;
392        }
393
394        info.cached_msg_size_in_mib = (self.msg_size.load(Ordering::Acquire) / (1024 * 1024)) as u32;
395
396        if !store.consuming_msg_orderly_tree_map.is_empty() {
397            if let Some((min_offset, _)) = store.consuming_msg_orderly_tree_map.first_key_value() {
398                info.transaction_msg_min_offset = *min_offset as u64;
399            }
400            if let Some((max_offset, _)) = store.consuming_msg_orderly_tree_map.last_key_value() {
401                info.transaction_msg_max_offset = *max_offset as u64;
402            }
403            info.transaction_msg_count = store.consuming_msg_orderly_tree_map.len() as u32;
404        }
405
406        info.locked = self.locked.load(Ordering::Acquire);
407        info.try_unlock_times = self.try_unlock_times.load(Ordering::Acquire) as u64;
408        info.last_lock_timestamp = self.last_lock_timestamp.load(Ordering::Acquire);
409        info.droped = self.dropped.load(Ordering::Acquire);
410        info.last_pull_timestamp = self.last_pull_timestamp.load(Ordering::Acquire);
411        info.last_consume_timestamp = self.last_consume_timestamp.load(Ordering::Acquire);
412    }
413
414    /// Returns `Some((min_offset, max_offset))` of the pending message tree, or `None` if empty.
415    pub(crate) async fn get_offset_span(&self) -> Option<(i64, i64)> {
416        let store = self.store.read().await;
417        match (
418            store.msg_tree_map.first_key_value(),
419            store.msg_tree_map.last_key_value(),
420        ) {
421            (Some((min, _)), Some((max, _))) => Some((*min, *max)),
422            _ => None,
423        }
424    }
425
426    pub(crate) fn set_last_pull_timestamp(&self, last_pull_timestamp: u64) {
427        self.last_pull_timestamp.store(last_pull_timestamp, Ordering::Release);
428    }
429
430    pub(crate) fn set_last_lock_timestamp(&self, last_lock_timestamp: u64) {
431        self.last_lock_timestamp.store(last_lock_timestamp, Ordering::Release);
432    }
433
434    pub fn msg_count(&self) -> u64 {
435        self.msg_count.load(Ordering::Acquire).max(0) as u64
436    }
437
438    pub(crate) fn msg_size(&self) -> u64 {
439        self.msg_size.load(Ordering::Acquire)
440    }
441
442    pub(crate) fn is_locked(&self) -> bool {
443        self.locked.load(Ordering::Acquire)
444    }
445
446    pub(crate) async fn has_temp_message(&self) -> bool {
447        !self.store.read().await.msg_tree_map.is_empty()
448    }
449
450    pub(crate) fn get_last_pull_timestamp(&self) -> u64 {
451        self.last_pull_timestamp.load(Ordering::Acquire)
452    }
453
454    pub(crate) fn get_last_consume_timestamp(&self) -> u64 {
455        self.last_consume_timestamp.load(Ordering::Acquire)
456    }
457
458    pub(crate) fn update_last_consume_timestamp(&self) {
459        self.last_consume_timestamp.store(current_millis(), Ordering::Release);
460    }
461
462    pub(crate) fn get_msg_acc_cnt(&self) -> i64 {
463        self.msg_acc_cnt.load(Ordering::Acquire)
464    }
465
466    pub(crate) fn get_try_unlock_times(&self) -> i64 {
467        self.try_unlock_times.load(Ordering::Acquire)
468    }
469
470    pub(crate) fn is_consuming(&self) -> bool {
471        self.consuming.load(Ordering::Acquire)
472    }
473
474    pub(crate) fn set_consuming(&self, consuming: bool) {
475        self.consuming.store(consuming, Ordering::Release);
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use super::*;
482    use bytes::Bytes;
483    use cheetah_string::CheetahString;
484    use rocketmq_common::common::message::MessageTrait;
485
486    fn create_test_messages(count: usize) -> Vec<ArcMut<MessageExt>> {
487        let mut messages = Vec::with_capacity(count);
488        for i in 0..count {
489            let mut msg = MessageExt {
490                queue_offset: i as i64,
491                ..Default::default()
492            };
493            msg.set_body(Bytes::from(vec![0u8; 100]));
494            msg.set_topic(CheetahString::from_static_str("test_topic"));
495            messages.push(ArcMut::new(msg));
496        }
497        messages
498    }
499
500    #[tokio::test]
501    async fn test_remove_message_count_correct() {
502        let pq = ProcessQueue::new();
503
504        let msgs = create_test_messages(10);
505        pq.put_message(&msgs).await;
506
507        assert_eq!(pq.msg_count(), 10);
508
509        pq.remove_message(&msgs[0..5]).await;
510
511        assert_eq!(pq.msg_count(), 5);
512        assert_eq!(pq.msg_size(), 500);
513    }
514
515    #[tokio::test]
516    async fn test_remove_message_all() {
517        let pq = ProcessQueue::new();
518
519        let msgs = create_test_messages(5);
520        pq.put_message(&msgs).await;
521
522        assert_eq!(pq.msg_count(), 5);
523
524        pq.remove_message(&msgs).await;
525
526        assert_eq!(pq.msg_count(), 0);
527        assert_eq!(pq.msg_size(), 0);
528    }
529
530    #[tokio::test]
531    async fn test_clear_resets_consuming_state_for_new_dispatch() {
532        let pq = ProcessQueue::new();
533
534        let msgs = create_test_messages(5);
535        assert!(pq.put_message(&msgs).await);
536        assert!(pq.is_consuming());
537
538        pq.clear().await;
539
540        assert_eq!(pq.msg_count(), 0);
541        assert_eq!(pq.msg_size(), 0);
542        assert!(!pq.is_consuming());
543        assert!(pq.put_message(&msgs).await);
544    }
545
546    #[tokio::test]
547    async fn test_take_messages_updates_consuming_map() {
548        let pq = ProcessQueue::new();
549
550        let msgs = create_test_messages(10);
551        pq.put_message(&msgs).await;
552
553        let taken = pq.take_messages(5).await;
554
555        assert_eq!(taken.len(), 5);
556
557        let store = pq.store.read().await;
558        assert_eq!(store.consuming_msg_orderly_tree_map.len(), 5);
559
560        for msg in &taken {
561            assert!(store.consuming_msg_orderly_tree_map.contains_key(&msg.queue_offset));
562        }
563    }
564
565    #[tokio::test]
566    async fn test_take_messages_updates_timestamp() {
567        let pq = ProcessQueue::new();
568
569        let msgs = create_test_messages(5);
570        pq.put_message(&msgs).await;
571
572        let timestamp_before = pq.last_consume_timestamp.load(Ordering::Acquire);
573
574        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
575
576        pq.take_messages(3).await;
577
578        let timestamp_after = pq.last_consume_timestamp.load(Ordering::Acquire);
579
580        assert!(timestamp_after > timestamp_before);
581    }
582
583    #[tokio::test]
584    async fn test_take_messages_empty_sets_consuming_false() {
585        let pq = ProcessQueue::new();
586
587        pq.consuming.store(true, Ordering::Release);
588
589        let taken = pq.take_messages(5).await;
590
591        assert_eq!(taken.len(), 0);
592        assert!(!pq.consuming.load(Ordering::Acquire));
593    }
594
595    #[tokio::test]
596    async fn test_fill_process_queue_info() {
597        let pq = ProcessQueue::new();
598
599        let msgs = create_test_messages(10);
600        pq.put_message(&msgs).await;
601
602        pq.take_messages(3).await;
603
604        let mut info = ProcessQueueInfo {
605            commit_offset: 0,
606            cached_msg_min_offset: 0,
607            cached_msg_max_offset: 0,
608            cached_msg_count: 0,
609            cached_msg_size_in_mib: 0,
610            transaction_msg_min_offset: 0,
611            transaction_msg_max_offset: 0,
612            transaction_msg_count: 0,
613            locked: false,
614            try_unlock_times: 0,
615            last_lock_timestamp: 0,
616            droped: false,
617            last_pull_timestamp: 0,
618            last_consume_timestamp: 0,
619        };
620
621        pq.fill_process_queue_info(&mut info).await;
622
623        assert_eq!(info.cached_msg_count, 7);
624        assert_eq!(info.cached_msg_min_offset, 3);
625        assert_eq!(info.cached_msg_max_offset, 9);
626
627        assert_eq!(info.transaction_msg_count, 3);
628        assert_eq!(info.transaction_msg_min_offset, 0);
629        assert_eq!(info.transaction_msg_max_offset, 2);
630    }
631
632    #[tokio::test]
633    async fn test_has_temp_message() {
634        let pq = ProcessQueue::new();
635
636        assert!(!pq.has_temp_message().await);
637
638        let msgs = create_test_messages(5);
639        pq.put_message(&msgs).await;
640
641        assert!(pq.has_temp_message().await);
642
643        pq.remove_message(&msgs).await;
644
645        assert!(!pq.has_temp_message().await);
646    }
647
648    #[tokio::test]
649    async fn test_rollback_after_take_messages() {
650        let pq = ProcessQueue::new();
651
652        let msgs = create_test_messages(10);
653        pq.put_message(&msgs).await;
654
655        assert_eq!(pq.msg_count(), 10);
656
657        let taken = pq.take_messages(5).await;
658        assert_eq!(taken.len(), 5);
659
660        let msg_tree_map_len = pq.store.read().await.msg_tree_map.len();
661        assert_eq!(msg_tree_map_len, 5);
662
663        pq.rollback().await;
664
665        let msg_tree_map_len_after = pq.store.read().await.msg_tree_map.len();
666        assert_eq!(msg_tree_map_len_after, 10);
667
668        let consuming_map_len = pq.store.read().await.consuming_msg_orderly_tree_map.len();
669        assert_eq!(consuming_map_len, 0);
670    }
671
672    #[tokio::test]
673    async fn test_commit_after_take_messages() {
674        let pq = ProcessQueue::new();
675
676        let msgs = create_test_messages(10);
677        pq.put_message(&msgs).await;
678
679        pq.take_messages(5).await;
680
681        assert_eq!(pq.msg_count(), 10);
682
683        let offset = pq.commit().await;
684
685        assert_eq!(offset, 5);
686        assert_eq!(pq.msg_count(), 5);
687
688        let consuming_map_len = pq.store.read().await.consuming_msg_orderly_tree_map.len();
689        assert_eq!(consuming_map_len, 0);
690    }
691
692    #[tokio::test]
693    async fn test_getter_methods() {
694        let pq = ProcessQueue::new();
695
696        assert_eq!(pq.get_try_unlock_times(), 0);
697        assert!(!pq.is_consuming());
698        assert!(!pq.is_dropped());
699        assert!(!pq.is_locked());
700
701        pq.inc_try_unlock_times();
702        assert_eq!(pq.get_try_unlock_times(), 1);
703
704        pq.set_consuming(true);
705        assert!(pq.is_consuming());
706
707        pq.set_dropped(true);
708        assert!(pq.is_dropped());
709
710        pq.set_locked(true);
711        assert!(pq.is_locked());
712    }
713
714    #[tokio::test]
715    async fn test_timestamp_getters() {
716        let pq = ProcessQueue::new();
717
718        let pull_ts = pq.get_last_pull_timestamp();
719        let consume_ts = pq.get_last_consume_timestamp();
720        let lock_ts = pq.get_last_lock_timestamp();
721
722        assert!(pull_ts > 0);
723        assert!(consume_ts > 0);
724        assert!(lock_ts > 0);
725
726        let new_ts = current_millis() + 1000;
727        pq.set_last_pull_timestamp(new_ts);
728        assert_eq!(pq.get_last_pull_timestamp(), new_ts);
729
730        pq.set_last_lock_timestamp(new_ts);
731        assert_eq!(pq.get_last_lock_timestamp(), new_ts);
732    }
733
734    #[tokio::test]
735    async fn test_msg_acc_cnt() {
736        let pq = ProcessQueue::new();
737
738        assert_eq!(pq.get_msg_acc_cnt(), 0);
739
740        let mut msgs = Vec::new();
741        for i in 0..5 {
742            let mut msg = MessageExt {
743                queue_offset: i as i64,
744                ..Default::default()
745            };
746            msg.set_body(Bytes::from(vec![0u8; 100]));
747            msg.set_topic(CheetahString::from_static_str("test_topic"));
748            msg.put_property(
749                CheetahString::from_static_str(MessageConst::PROPERTY_MAX_OFFSET),
750                CheetahString::from(format!("{}", i + 100)),
751            );
752            msgs.push(ArcMut::new(msg));
753        }
754
755        pq.put_message(&msgs).await;
756
757        assert!(pq.get_msg_acc_cnt() > 0);
758    }
759
760    #[tokio::test]
761    async fn test_clear_resets_all_state() {
762        let pq = ProcessQueue::new();
763
764        let msgs = create_test_messages(10);
765        pq.put_message(&msgs).await;
766        pq.take_messages(5).await;
767
768        assert_eq!(pq.msg_count(), 10);
769        assert!(pq.msg_size() > 0);
770
771        pq.clear().await;
772
773        assert_eq!(pq.msg_count(), 0);
774        assert_eq!(pq.msg_size(), 0);
775
776        let msg_tree_map_len = pq.store.read().await.msg_tree_map.len();
777        assert_eq!(msg_tree_map_len, 0);
778
779        let consuming_map_len = pq.store.read().await.consuming_msg_orderly_tree_map.len();
780        assert_eq!(consuming_map_len, 0);
781    }
782
783    #[tokio::test]
784    async fn test_get_max_span() {
785        let pq = ProcessQueue::new();
786
787        assert_eq!(pq.get_max_span().await, 0);
788
789        let msgs = create_test_messages(10);
790        pq.put_message(&msgs).await;
791
792        let span = pq.get_max_span().await;
793        assert_eq!(span, 9);
794    }
795
796    #[tokio::test]
797    async fn test_consuming_flag_behavior() {
798        let pq = ProcessQueue::new();
799
800        assert!(!pq.is_consuming());
801
802        let msgs = create_test_messages(5);
803        let should_dispatch = pq.put_message(&msgs).await;
804
805        assert!(should_dispatch);
806        assert!(pq.is_consuming());
807
808        pq.set_consuming(false);
809        assert!(!pq.is_consuming());
810    }
811
812    #[tokio::test]
813    async fn test_queue_offset_max_with_unordered_messages() {
814        let pq = ProcessQueue::new();
815
816        let mut msgs = Vec::new();
817        for offset in [5, 2, 8, 1, 9, 3].iter() {
818            let mut msg = MessageExt {
819                queue_offset: *offset,
820                ..Default::default()
821            };
822            msg.set_body(Bytes::from(vec![0u8; 100]));
823            msg.set_topic(CheetahString::from_static_str("test_topic"));
824            msgs.push(ArcMut::new(msg));
825        }
826
827        pq.put_message(&msgs).await;
828
829        let max_offset = pq.store.read().await.queue_offset_max;
830        assert_eq!(max_offset, 9);
831    }
832
833    #[tokio::test]
834    async fn test_commit_preserves_msg_size_correctly() {
835        let pq = ProcessQueue::new();
836
837        let msgs = create_test_messages(10);
838        pq.put_message(&msgs).await;
839
840        assert_eq!(pq.msg_size(), 1000);
841
842        pq.take_messages(3).await;
843
844        assert_eq!(pq.msg_count(), 10);
845        assert_eq!(pq.msg_size(), 1000);
846
847        pq.commit().await;
848
849        assert_eq!(pq.msg_count(), 7);
850        assert_eq!(pq.msg_size(), 700);
851    }
852
853    #[tokio::test]
854    async fn test_commit_when_all_messages_consumed() {
855        let pq = ProcessQueue::new();
856
857        let msgs = create_test_messages(5);
858        pq.put_message(&msgs).await;
859
860        assert_eq!(pq.msg_size(), 500);
861
862        pq.take_messages(5).await;
863
864        pq.commit().await;
865
866        assert_eq!(pq.msg_count(), 0);
867        assert_eq!(pq.msg_size(), 0);
868    }
869}
rocketmq_client_rust/consumer/consumer_impl/process_queue.rs

rocketmq_client_rust/consumer/consumer_impl/
process_queue.rs
