[ISSUE #1053] Support request code UNLOCK_BATCH_MQ(42) (#1062)

Author: mxsm

rocketmq-broker/src/out_api/broker_outer_api.rs

@@ -37,6 +37,7 @@ use rocketmq_remoting::protocol::header::lock_batch_mq_request_header::LockBatch
 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::header::unlock_batch_mq_request_header::UnlockBatchMqRequestHeader;
 use rocketmq_remoting::protocol::namesrv::RegisterBrokerResult;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 use rocketmq_remoting::protocol::route::route_data_view::QueueData;
@@ -359,6 +360,37 @@ impl BrokerOuterAPI {
             Err(e) => Err(BrokerClientError(e)),
         }
     }
+
+    pub async fn unlock_batch_mq_async(
+        &self,
+        addr: String,
+        request_body: bytes::Bytes,
+        timeout_millis: u64,
+    ) -> Result<()> {
+        let mut request = RemotingCommand::create_request_command(
+            RequestCode::UnlockBatchMq,
+            UnlockBatchMqRequestHeader::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 {
+                    Ok(())
+                } 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> {

rocketmq-broker/src/processor/admin_broker_processor.rs

@@ -209,6 +209,11 @@ impl AdminBrokerProcessor {
                     .await
             }
 
+            RequestCode::UnlockBatchMq => {
+                self.batch_mq_handler
+                    .unlock_batch_mq(channel, ctx, request_code, request)
+                    .await
+            }
             _ => Some(get_unknown_cmd_response(request_code)),
         }
     }

rocketmq-broker/src/processor/admin_broker_processor/batch_mq_handler.rs

@@ -24,6 +24,7 @@ use rocketmq_remoting::code::request_code::RequestCode;
 use rocketmq_remoting::net::channel::Channel;
 use rocketmq_remoting::protocol::body::request::lock_batch_request_body::LockBatchRequestBody;
 use rocketmq_remoting::protocol::body::response::lock_batch_response_body::LockBatchResponseBody;
+use rocketmq_remoting::protocol::body::unlock_batch_request_body::UnlockBatchRequestBody;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 use rocketmq_remoting::protocol::RemotingDeserializable;
 use rocketmq_remoting::protocol::RemotingSerializable;
@@ -122,8 +123,32 @@ impl BatchMqHandler {
         _channel: Channel,
         _ctx: ConnectionHandlerContext,
         _request_code: RequestCode,
-        _request: RemotingCommand,
+        request: RemotingCommand,
     ) -> Option<RemotingCommand> {
-        unimplemented!("unlockBatchMQ")
+        let mut request_body = UnlockBatchRequestBody::decode(request.get_body().unwrap()).unwrap();
+        if request_body.only_this_broker || !self.inner.broker_config.lock_in_strict_mode {
+            self.inner.rebalance_lock_manager.unlock_batch(
+                request_body.consumer_group.as_ref().unwrap(),
+                &request_body.mq_set,
+                request_body.client_id.as_ref().unwrap(),
+            );
+        } else {
+            request_body.only_this_broker = true;
+            let request_body = Bytes::from(request_body.encode());
+            for broker_addr in self.inner.broker_member_group.broker_addrs.values() {
+                match self
+                    .inner
+                    .broker_out_api
+                    .unlock_batch_mq_async(broker_addr.clone(), request_body.clone(), 1000)
+                    .await
+                {
+                    Ok(_) => {}
+                    Err(e) => {
+                        warn!("unlockBatchMQ exception on {}, {}", broker_addr, e);
+                    }
+                }
+            }
+        }
+        Some(RemotingCommand::create_response_command())
     }
 }