[ISSUE ##1052] Support request code LOCK_BATCH_MQ(41) (#1061)

Author: mxsm

rocketmq-broker/src/broker_runtime.rs

@@ -30,6 +30,7 @@ use rocketmq_common::common::statistics::state_getter::StateGetter;
 use rocketmq_common::ArcRefCellWrapper;
 use rocketmq_common::TimeUtils::get_current_millis;
 use rocketmq_common::UtilAll::compute_next_morning_time_millis;
+use rocketmq_remoting::protocol::body::broker_body::broker_member_group::BrokerMemberGroup;
 use rocketmq_remoting::protocol::body::topic_info_wrapper::topic_config_wrapper::TopicConfigAndMappingSerializeWrapper;
 use rocketmq_remoting::protocol::body::topic_info_wrapper::topic_config_wrapper::TopicConfigSerializeWrapper;
 use rocketmq_remoting::protocol::namespace_util::NamespaceUtil;
@@ -114,6 +115,7 @@ pub(crate) struct BrokerRuntime {
     #[cfg(feature = "local_file_store")]
     pull_request_hold_service: Option<PullRequestHoldService<DefaultMessageStore>>,
     rebalance_lock_manager: Arc<RebalanceLockManager>,
+    broker_member_group: Arc<BrokerMemberGroup>,
 }
 
 impl Clone for BrokerRuntime {
@@ -147,6 +149,7 @@ impl Clone for BrokerRuntime {
             is_isolated: self.is_isolated.clone(),
             pull_request_hold_service: self.pull_request_hold_service.clone(),
             rebalance_lock_manager: self.rebalance_lock_manager.clone(),
+            broker_member_group: self.broker_member_group.clone(),
         }
     }
 }
@@ -190,6 +193,14 @@ impl BrokerRuntime {
         }));
         let broker_stats_manager = Arc::new(stats_manager);
         consumer_manager.set_broker_stats_manager(Some(Arc::downgrade(&broker_stats_manager)));
+        let mut broker_member_group = BrokerMemberGroup::new(
+            broker_config.broker_identity.broker_cluster_name.clone(),
+            broker_config.broker_identity.broker_name.clone(),
+        );
+        broker_member_group.broker_addrs.insert(
+            broker_config.broker_identity.broker_id,
+            broker_config.get_broker_addr(),
+        );
         Self {
             broker_config: broker_config.clone(),
             message_store_config,
@@ -222,6 +233,7 @@ impl BrokerRuntime {
             is_isolated: Arc::new(AtomicBool::new(false)),
             pull_request_hold_service: None,
             rebalance_lock_manager: Arc::new(Default::default()),
+            broker_member_group: Arc::new(broker_member_group),
         }
     }
 
@@ -471,6 +483,8 @@ impl BrokerRuntime {
             self.consumer_manager.clone(),
             self.broker_out_api.clone(),
             self.broker_stats_manager.clone(),
+            self.rebalance_lock_manager.clone(),
+            self.broker_member_group.clone(),
         );
 
         BrokerRequestProcessor {

rocketmq-broker/src/client/net/broker_to_client.rs

@@ -18,7 +18,7 @@ use rocketmq_remoting::net::channel::Channel;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 
 use crate::error::BrokerError::BrokerClientError;
-use crate::BrokerResult;
+use crate::Result;
 
 #[derive(Default, Clone)]
 pub struct Broker2Client;
@@ -29,7 +29,7 @@ impl Broker2Client {
         channel: &mut Channel,
         request: RemotingCommand,
         timeout_millis: u64,
-    ) -> BrokerResult<RemotingCommand> {
+    ) -> Result<RemotingCommand> {
         match channel.send_wait_response(request, timeout_millis).await {
             Ok(value) => Ok(value),
             Err(e) => Err(BrokerClientError(e)),

rocketmq-broker/src/client/rebalance/rebalance_lock_manager.rs

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::sync::atomic::AtomicI64;
 use std::sync::Arc;
 
@@ -34,7 +35,7 @@ lazy_static! {
     };
 }
 
-type MessageQueueLockTable = HashMap<String, HashMap<Arc<MessageQueue>, LockEntry>>;
+type MessageQueueLockTable = HashMap<String, HashMap<MessageQueue, LockEntry>>;
 
 #[derive(Clone, Default)]
 pub struct RebalanceLockManager {
@@ -60,26 +61,25 @@ impl RebalanceLockManager {
     pub fn try_lock_batch(
         &self,
         group: &str,
-        mqs: Vec<Arc<MessageQueue>>,
+        mqs: &HashSet<MessageQueue>,
         client_id: &str,
-    ) -> Vec<Arc<MessageQueue>> {
-        let mut lock_mqs = Vec::new();
-        let mut not_locked_mqs = Vec::new();
+    ) -> HashSet<MessageQueue> {
+        let mut lock_mqs = HashSet::with_capacity(mqs.len());
+        let mut not_locked_mqs = HashSet::with_capacity(mqs.len());
         for mq in mqs.iter() {
             if self.is_locked(group, mq, client_id) {
-                lock_mqs.push(mq.clone());
+                lock_mqs.insert(mq.clone());
             } else {
-                not_locked_mqs.push(mq.clone());
+                not_locked_mqs.insert(mq.clone());
             }
         }
         if !not_locked_mqs.is_empty() {
             let mut write_guard = self.mq_lock_table.write();
-            let mut group_value = write_guard.get_mut(group);
-            if group_value.is_none() {
-                group_value = Some(write_guard.entry(group.to_string()).or_default());
-            }
-            let group_value = group_value.unwrap();
-            for mq in not_locked_mqs.iter() {
+            let group_value = write_guard
+                .entry(group.to_string())
+                .or_insert(HashMap::with_capacity(32));
+
+            for mq in not_locked_mqs {
                 let lock_entry = group_value.entry(mq.clone()).or_insert_with(|| {
                     info!(
                         "RebalanceLockManager#tryLockBatch: lock a message which has not been \
@@ -96,7 +96,7 @@ impl RebalanceLockManager {
                         get_current_millis() as i64,
                         std::sync::atomic::Ordering::Relaxed,
                     );
-                    lock_mqs.push(mq.clone());
+                    lock_mqs.insert(mq);
                     continue;
                 }
                 let old_client_id = lock_entry.client_id.as_str().to_string();
@@ -106,12 +106,12 @@ impl RebalanceLockManager {
                         get_current_millis() as i64,
                         std::sync::atomic::Ordering::Relaxed,
                     );
-                    lock_mqs.push(mq.clone());
                     warn!(
                         "RebalanceLockManager#tryLockBatch: try to lock a expired message queue, \
                          group={} mq={:?}, old client id={}, new client id={}",
                         group, mq, old_client_id, client_id
                     );
+                    lock_mqs.insert(mq);
                     continue;
                 }
                 warn!(
@@ -124,7 +124,7 @@ impl RebalanceLockManager {
         lock_mqs
     }
 
-    pub fn unlock_batch(&self, group: &str, mqs: Vec<Arc<MessageQueue>>, client_id: &str) {
+    pub fn unlock_batch(&self, group: &str, mqs: &HashSet<MessageQueue>, client_id: &str) {
         let mut write_guard = self.mq_lock_table.write();
         let group_value = write_guard.get_mut(group);
         if group_value.is_none() {
@@ -163,7 +163,7 @@ impl RebalanceLockManager {
         }
     }
 
-    fn is_locked(&self, group: &str, mq: &Arc<MessageQueue>, client_id: &str) -> bool {
+    fn is_locked(&self, group: &str, mq: &MessageQueue, client_id: &str) -> bool {
         let lock_table = self.mq_lock_table.read();
         let group_value = lock_table.get(group);
         if group_value.is_none() {
@@ -231,59 +231,71 @@ mod rebalance_lock_manager_tests {
     #[test]
     fn lock_all_expired_returns_false_when_active_locks_exist() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        manager.try_lock_batch("test_group", &set, "client_1");
         assert!(!manager.is_lock_all_expired("test_group"));
     }
 
     #[test]
     fn try_lock_batch_locks_message_queues_for_new_group() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        let locked_mqs = manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        let locked_mqs = manager.try_lock_batch("test_group", &set, "client_1");
         assert_eq!(locked_mqs.len(), 1);
     }
 
     #[test]
     fn try_lock_batch_does_not_lock_already_locked_message_queues() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
-        let locked_mqs = manager.try_lock_batch("test_group", vec![mq.clone()], "client_2");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        manager.try_lock_batch("test_group", &set, "client_1");
+        let locked_mqs = manager.try_lock_batch("test_group", &set, "client_2");
         assert!(locked_mqs.is_empty());
     }
 
     #[test]
     fn unlock_batch_unlocks_message_queues_locked_by_client() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
-        manager.unlock_batch("test_group", vec![mq.clone()], "client_1");
-        let locked_mqs = manager.try_lock_batch("test_group", vec![mq.clone()], "client_2");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        manager.try_lock_batch("test_group", &set, "client_1");
+        manager.unlock_batch("test_group", &set, "client_1");
+        let locked_mqs = manager.try_lock_batch("test_group", &set, "client_2");
         assert_eq!(locked_mqs.len(), 1);
     }
 
     #[test]
     fn unlock_batch_does_not_unlock_message_queues_locked_by_other_clients() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
-        manager.unlock_batch("test_group", vec![mq.clone()], "client_2");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        manager.try_lock_batch("test_group", &set, "client_1");
+        manager.unlock_batch("test_group", &set, "client_2");
         assert!(!manager.is_lock_all_expired("test_group"));
     }
 
     #[test]
     fn is_locked_returns_true_for_locked_message_queue() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
-        manager.try_lock_batch("test_group", vec![mq.clone()], "client_1");
+        let mq = MessageQueue::default();
+        let mut set = HashSet::new();
+        set.insert(mq.clone());
+        manager.try_lock_batch("test_group", &set, "client_1");
         assert!(manager.is_locked("test_group", &mq, "client_1"));
     }
 
     #[test]
     fn is_locked_returns_false_for_unlocked_message_queue() {
         let manager = RebalanceLockManager::default();
-        let mq = Arc::new(MessageQueue::default());
+        let mq = MessageQueue::default();
         assert!(!manager.is_locked("test_group", &mq, "client_1"));
     }
 }

rocketmq-broker/src/error.rs

@@ -20,4 +20,7 @@ use thiserror::Error;
 pub enum BrokerError {
     #[error("broker client error: {0}")]
     BrokerClientError(#[from] rocketmq_remoting::error::Error),
+
+    #[error("Client exception occurred: CODE:{0}, broker address:{2}, Message:{1}")]
+    MQBrokerError(i32, String, String),
 }

rocketmq-broker/src/lib.rs

@@ -47,4 +47,4 @@ pub(crate) mod topic;
 pub(crate) mod util;
 
 type RemotingError = rocketmq_remoting::error::Error;
-type BrokerResult<T> = Result<T, BrokerError>;
+type Result<T> = std::result::Result<T, BrokerError>;

rocketmq-broker/src/out_api/broker_outer_api.rs

@@ -14,12 +14,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Weak;
 
 use dns_lookup::lookup_host;
 use rocketmq_common::common::broker::broker_config::BrokerIdentity;
 use rocketmq_common::common::config::TopicConfig;
+use rocketmq_common::common::message::message_queue::MessageQueue;
 use rocketmq_common::utils::crc32_utils;
 use rocketmq_common::utils::serde_json_utils::SerdeJsonUtils;
 use rocketmq_common::ArcRefCellWrapper;
@@ -29,14 +31,17 @@ use rocketmq_remoting::code::request_code::RequestCode;
 use rocketmq_remoting::code::response_code::ResponseCode;
 use rocketmq_remoting::protocol::body::broker_body::register_broker_body::RegisterBrokerBody;
 use rocketmq_remoting::protocol::body::kv_table::KVTable;
+use rocketmq_remoting::protocol::body::response::lock_batch_response_body::LockBatchResponseBody;
 use rocketmq_remoting::protocol::body::topic_info_wrapper::topic_config_wrapper::TopicConfigAndMappingSerializeWrapper;
+use rocketmq_remoting::protocol::header::lock_batch_mq_request_header::LockBatchMqRequestHeader;
 use rocketmq_remoting::protocol::header::namesrv::register_broker_header::RegisterBrokerRequestHeader;
 use rocketmq_remoting::protocol::header::namesrv::register_broker_header::RegisterBrokerResponseHeader;
 use rocketmq_remoting::protocol::header::namesrv::topic_operation_header::RegisterTopicRequestHeader;
 use rocketmq_remoting::protocol::namesrv::RegisterBrokerResult;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 use rocketmq_remoting::protocol::route::route_data_view::QueueData;
 use rocketmq_remoting::protocol::route::topic_route_data::TopicRouteData;
+use rocketmq_remoting::protocol::RemotingDeserializable;
 use rocketmq_remoting::protocol::RemotingSerializable;
 use rocketmq_remoting::remoting::RemotingService;
 use rocketmq_remoting::request_processor::default_request_processor::DefaultRemotingRequestProcessor;
@@ -48,6 +53,10 @@ use tracing::debug;
 use tracing::error;
 use tracing::info;
 
+use crate::error::BrokerError;
+use crate::error::BrokerError::BrokerClientError;
+use crate::Result;
+
 pub struct BrokerOuterAPI {
     remoting_client: ArcRefCellWrapper<RocketmqDefaultClient<DefaultRemotingRequestProcessor>>,
     name_server_address: Option<String>,
@@ -317,6 +326,39 @@ impl BrokerOuterAPI {
     pub fn rpc_client(&self) -> &RpcClientImpl {
         &self.rpc_client
     }
+
+    pub async fn lock_batch_mq_async(
+        &self,
+        addr: String,
+        request_body: bytes::Bytes,
+        timeout_millis: u64,
+    ) -> Result<HashSet<MessageQueue>> {
+        let mut request = RemotingCommand::create_request_command(
+            RequestCode::LockBatchMq,
+            LockBatchMqRequestHeader::default(),
+        );
+        request.set_body_mut_ref(Some(request_body));
+        let result = self
+            .remoting_client
+            .invoke_async(Some(addr), request, timeout_millis)
+            .await;
+        match result {
+            Ok(response) => {
+                if ResponseCode::from(response.code()) == ResponseCode::Success {
+                    let lock_batch_response_body =
+                        LockBatchResponseBody::decode(response.get_body().unwrap()).unwrap();
+                    Ok(lock_batch_response_body.lock_ok_mq_set)
+                } else {
+                    Err(BrokerError::MQBrokerError(
+                        response.code(),
+                        response.remark().cloned().unwrap_or("".to_string()),
+                        "".to_string(),
+                    ))
+                }
+            }
+            Err(e) => Err(BrokerClientError(e)),
+        }
+    }
 }
 
 fn dns_lookup_address_by_domain(domain: &str) -> Vec<String> {