rocketmq_client_rust/consumer/
message_queue_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::sync::Arc;
16
17use dashmap::DashMap;
18use rocketmq_common::common::message::message_queue::MessageQueue;
19use tokio::sync::Mutex;
20
21type LockObject = Arc<Mutex<()>>;
22
23#[derive(Default)]
24pub struct MessageQueueLock {
25    mq_lock_table: Arc<DashMap<MessageQueue, Arc<DashMap<i32, LockObject>>>>,
26}
27
28impl MessageQueueLock {
29    pub fn new() -> Self {
30        MessageQueueLock {
31            mq_lock_table: Arc::new(DashMap::new()),
32        }
33    }
34
35    pub async fn fetch_lock_object(&self, mq: &MessageQueue) -> Arc<Mutex<()>> {
36        self.fetch_lock_object_with_sharding_key(mq, -1).await
37    }
38
39    pub async fn fetch_lock_object_with_sharding_key(
40        &self,
41        mq: &MessageQueue,
42        sharding_key_index: i32,
43    ) -> Arc<Mutex<()>> {
44        // Get or create the inner DashMap for this message queue
45        let inner_map = self
46            .mq_lock_table
47            .entry(mq.clone())
48            .or_insert_with(|| Arc::new(DashMap::new()))
49            .value()
50            .clone();
51
52        // Get or create the lock for this sharding key
53        let lock = inner_map
54            .entry(sharding_key_index)
55            .or_insert_with(|| Arc::new(Mutex::new(())))
56            .value()
57            .clone();
58
59        lock
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    fn message_queue() -> MessageQueue {
68        MessageQueue::from_parts("TopicA", "BrokerA", 0)
69    }
70
71    #[tokio::test]
72    async fn fetch_lock_object_reuses_lock_for_same_queue_like_java() {
73        let lock = MessageQueueLock::new();
74        let mq = message_queue();
75
76        let first = lock.fetch_lock_object(&mq).await;
77        let second = lock.fetch_lock_object(&mq).await;
78
79        assert!(Arc::ptr_eq(&first, &second));
80    }
81
82    #[tokio::test]
83    async fn fetch_lock_object_uses_independent_locks_per_sharding_key_like_java() {
84        let lock = MessageQueueLock::new();
85        let mq = message_queue();
86
87        let first = lock.fetch_lock_object_with_sharding_key(&mq, 0).await;
88        let second = lock.fetch_lock_object_with_sharding_key(&mq, 1).await;
89
90        assert!(!Arc::ptr_eq(&first, &second));
91    }
92}
rocketmq_client_rust/consumer/message_queue_lock.rs

rocketmq_client_rust/consumer/
message_queue_lock.rs
