[ISSUE #1301]🔥Rocketmq-rust broker supports request code QUERY_ASSIGNMENT(400)🚀

mxsm · mxsm · commit 870b3d203678 · 2024-11-28T12:14:57.000Z
diff --git a/rocketmq-broker/src/broker_runtime.rs b/rocketmq-broker/src/broker_runtime.rs
@@ -81,6 +81,7 @@ use crate::schedule::schedule_message_service::ScheduleMessageService;
 use crate::subscription::manager::subscription_group_manager::SubscriptionGroupManager;
 use crate::topic::manager::topic_config_manager::TopicConfigManager;
 use crate::topic::manager::topic_queue_mapping_manager::TopicQueueMappingManager;
+use crate::topic::manager::topic_route_info_manager::TopicRouteInfoManager;
 use crate::topic::topic_queue_mapping_clean_service::TopicQueueMappingCleanService;
 use crate::transaction::queue::default_transactional_message_check_listener::DefaultTransactionalMessageCheckListener;
 use crate::transaction::queue::default_transactional_message_service::DefaultTransactionalMessageService;
@@ -133,6 +134,7 @@ pub(crate) struct BrokerRuntime {
         Option<Arc<DefaultTransactionalMessageCheckListener<DefaultMessageStore>>>,
     transactional_message_check_service: Option<Arc<TransactionalMessageCheckService>>,
     transaction_metrics_flush_service: Option<Arc<TransactionMetricsFlushService>>,
+    topic_route_info_manager: Arc<TopicRouteInfoManager>,
 }
 
 impl Clone for BrokerRuntime {
@@ -171,6 +173,7 @@ impl Clone for BrokerRuntime {
             transactional_message_check_listener: self.transactional_message_check_listener.clone(),
             transactional_message_check_service: None,
             transaction_metrics_flush_service: None,
+            topic_route_info_manager: self.topic_route_info_manager.clone(),
         }
     }
 }
@@ -239,7 +242,7 @@ impl BrokerRuntime {
             broker_stats: None,
             schedule_message_service: Default::default(),
             timer_message_store: None,
-            broker_out_api: broker_outer_api,
+            broker_out_api: broker_outer_api.clone(),
             broker_runtime: Some(runtime),
             producer_manager,
             consumer_manager,
@@ -259,6 +262,10 @@ impl BrokerRuntime {
             transactional_message_check_listener: None,
             transactional_message_check_service: None,
             transaction_metrics_flush_service: None,
+            topic_route_info_manager: Arc::new(TopicRouteInfoManager::new(
+                broker_outer_api,
+                broker_config,
+            )),
         }
     }
 
@@ -540,6 +547,9 @@ impl BrokerRuntime {
             consumer_manage_processor: ArcMut::new(consumer_manage_processor),
             query_assignment_processor: ArcMut::new(QueryAssignmentProcessor::new(
                 self.message_store_config.clone(),
+                self.broker_config.clone(),
+                self.topic_route_info_manager.clone(),
+                self.consumer_manager.clone(),
             )),
             query_message_processor: ArcMut::new(query_message_processor),
             end_transaction_processor: ArcMut::new(EndTransactionProcessor::new(
@@ -736,6 +746,8 @@ impl BrokerRuntime {
             let this = pull_request_hold_service.clone();
             pull_request_hold_service.start(this);
         }
+
+        self.topic_route_info_manager.start();
     }
 
     async fn update_namesrv_addr(&mut self) {
diff --git a/rocketmq-broker/src/error.rs b/rocketmq-broker/src/error.rs
@@ -27,4 +27,10 @@ pub enum BrokerError {
 
     #[error("Client exception occurred: CODE:{0}, broker address:{2}, Message:{1}")]
     MQBrokerError(i32, String, String),
+
+    #[error("{0}")]
+    IllegalArgumentError(String),
+
+    #[error("Client error: {0}")]
+    ClientError(#[from] rocketmq_client_rust::error::MQClientError),
 }
diff --git a/rocketmq-broker/src/processor/query_assignment_processor.rs b/rocketmq-broker/src/processor/query_assignment_processor.rs
@@ -25,21 +25,43 @@ use crate::load_balance::message_request_mode_manager::MessageRequestModeManager
  use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
  use rocketmq_remoting::runtime::connection_handler_context::ConnectionHandlerContext;
  use rocketmq_store::config::message_store_config::MessageStoreConfig;
- use std::collections::HashMap;
+ use std::collections::{HashMap, HashSet};
  use std::sync::Arc;
+use tracing::warn;
+use rocketmq_common::common::broker::broker_config::BrokerConfig;
+use rocketmq_common::common::message::message_enum::MessageRequestMode;
+use rocketmq_common::common::message::message_queue::MessageQueue;
+use rocketmq_common::common::message::message_queue_assignment::MessageQueueAssignment;
+use rocketmq_common::common::mix_all;
 use rocketmq_common::common::mix_all::RETRY_GROUP_TOPIC_PREFIX;
 use rocketmq_remoting::code::response_code::ResponseCode;
+use rocketmq_remoting::protocol::body::query_assignment_request_body::QueryAssignmentRequestBody;
+use rocketmq_remoting::protocol::body::query_assignment_response_body::QueryAssignmentResponseBody;
 use rocketmq_remoting::protocol::body::set_message_request_mode_request_body::SetMessageRequestModeRequestBody;
-use rocketmq_remoting::protocol::RemotingDeserializable;
+use rocketmq_remoting::protocol::heartbeat::message_model::MessageModel;
+use rocketmq_remoting::protocol::{RemotingDeserializable, RemotingSerializable};
+use crate::client::manager::consumer_manager::ConsumerManager;
+use crate::error::BrokerError;
+use crate::error::BrokerError::IllegalArgumentError;
+use crate::topic::manager::topic_route_info_manager::TopicRouteInfoManager;
+use crate::Result;
 
 pub struct QueryAssignmentProcessor {
     message_request_mode_manager: MessageRequestModeManager,
     load_strategy: HashMap<CheetahString, Arc<dyn AllocateMessageQueueStrategy>>,
     message_store_config: Arc<MessageStoreConfig>,
+    broker_config: Arc<BrokerConfig>,
+    topic_route_info_manager: Arc<TopicRouteInfoManager>,
+    consumer_manager: Arc<ConsumerManager>,
 }
 
 impl QueryAssignmentProcessor {
-    pub fn new(message_store_config: Arc<MessageStoreConfig>) -> Self {
+    pub fn new(
+        message_store_config: Arc<MessageStoreConfig>,
+        broker_config: Arc<BrokerConfig>,
+        topic_route_info_manager: Arc<TopicRouteInfoManager>,
+        consumer_manager: Arc<ConsumerManager>,
+    ) -> Self {
         let allocate_message_queue_averagely: Arc<dyn AllocateMessageQueueStrategy> =
             Arc::new(AllocateMessageQueueAveragely);
         let allocate_message_queue_averagely_by_circle: Arc<dyn AllocateMessageQueueStrategy> =
@@ -59,6 +81,9 @@ impl QueryAssignmentProcessor {
             message_request_mode_manager: manager,
             load_strategy,
             message_store_config,
+            broker_config,
+            topic_route_info_manager,
+            consumer_manager,
         }
     }
 }
@@ -82,11 +107,220 @@ impl QueryAssignmentProcessor {
 
     async fn query_assignment(
         &mut self,
-        _channel: Channel,
+        channel: Channel,
         _ctx: ConnectionHandlerContext,
-        _request: RemotingCommand,
+        request: RemotingCommand,
     ) -> Option<RemotingCommand> {
-        unimplemented!()
+        let request_body =
+            QueryAssignmentRequestBody::decode(request.get_body().expect("empty body"))
+                .expect("decode QueryAssignmentRequestBody failed");
+        let set_message_request_mode_request_body = self
+            .message_request_mode_manager
+            .get_message_request_mode(&request_body.topic, &request_body.consumer_group);
+        let set_message_request_mode_request_body =
+            if let Some(set_message_request_mode_request_body) =
+                set_message_request_mode_request_body
+            {
+                set_message_request_mode_request_body
+            } else {
+                let mut body = SetMessageRequestModeRequestBody {
+                    topic: request_body.topic.clone(),
+                    consumer_group: request_body.consumer_group.clone(),
+                    ..Default::default()
+                };
+                if request_body.topic.starts_with(RETRY_GROUP_TOPIC_PREFIX) {
+                    body.mode = MessageRequestMode::Pull;
+                } else {
+                    body.mode = self.broker_config.default_message_request_mode;
+                }
+                if body.mode == MessageRequestMode::Pop {
+                    body.pop_share_queue_num = self.broker_config.default_pop_share_queue_num;
+                }
+                body
+            };
+        let mode = set_message_request_mode_request_body.mode;
+        let message_queues = self
+            .do_load_balance(
+                &request_body.topic,
+                &request_body.consumer_group,
+                &request_body.client_id,
+                request_body.message_model,
+                &request_body.strategy_name,
+                set_message_request_mode_request_body,
+                channel,
+            )
+            .await;
+        let assignments = if let Some(message_queues) = message_queues {
+            message_queues
+                .into_iter()
+                .map(|mq| MessageQueueAssignment {
+                    message_queue: Some(mq),
+                    mode,
+                    attachments: None,
+                })
+                .collect()
+        } else {
+            HashSet::new()
+        };
+        let body = QueryAssignmentResponseBody {
+            message_queue_assignments: assignments,
+        };
+        Some(RemotingCommand::create_response_command().set_body(body.encode()))
+    }
+
+    async fn do_load_balance(
+        &mut self,
+        topic: &CheetahString,
+        consumer_group: &CheetahString,
+        client_id: &CheetahString,
+        message_model: MessageModel,
+        strategy_name: &CheetahString,
+        set_message_request_mode_request_body: SetMessageRequestModeRequestBody,
+        channel: Channel,
+    ) -> Option<HashSet<MessageQueue>> {
+        match message_model {
+            MessageModel::Broadcasting => {
+                let assigned_queue_set = self
+                    .topic_route_info_manager
+                    .get_topic_subscribe_info(topic)
+                    .await;
+                if assigned_queue_set.is_none() {
+                    warn!(
+                        "QueryLoad: no assignment for group[{}], the topic[{}] does not exist.",
+                        consumer_group, topic
+                    );
+                }
+                assigned_queue_set
+            }
+            MessageModel::Clustering => {
+                let mq_set = if mix_all::is_lmq(Some(topic.as_str())) {
+                    let mut set = HashSet::new();
+                    let queue = MessageQueue::from_parts(
+                        topic.clone(),
+                        self.broker_config.broker_name.clone(),
+                        mix_all::LMQ_QUEUE_ID as i32,
+                    );
+                    set.insert(queue);
+                    Some(set)
+                } else {
+                    self.topic_route_info_manager
+                        .get_topic_subscribe_info(topic)
+                        .await
+                };
+
+                if mq_set.is_none() || mq_set.as_ref().unwrap().is_empty() {
+                    if topic.starts_with(RETRY_GROUP_TOPIC_PREFIX) {
+                        warn!(
+                            "QueryLoad: no assignment for group[{}], the topic[{}] does not exist.",
+                            consumer_group, topic
+                        );
+                    }
+                    return None;
+                }
+
+                if !self.broker_config.server_load_balancer_enable {
+                    return mq_set;
+                }
+                let consumer_group_info = self
+                    .consumer_manager
+                    .get_consumer_group_info(consumer_group);
+                let mut cid_all = if let Some(consumer_group_info) = consumer_group_info {
+                    consumer_group_info.get_all_client_ids()
+                } else {
+                    vec![]
+                };
+                if cid_all.is_empty() {
+                    warn!(
+                        "QueryLoad: no assignment for group[{}] topic[{}], get consumer id list \
+                         failed",
+                        consumer_group, topic
+                    );
+                    return None;
+                }
+                let mut mq_all = mq_set.unwrap().into_iter().collect::<Vec<MessageQueue>>();
+                mq_all.sort();
+                cid_all.sort();
+
+                let strategy = self.load_strategy.get(strategy_name);
+                if strategy.is_none() {
+                    warn!(
+                        "QueryLoad: unsupported strategy [{}],  {}",
+                        strategy_name,
+                        channel.remote_address()
+                    );
+                    return None;
+                }
+                let strategy = strategy.unwrap();
+                let result =
+                    if set_message_request_mode_request_body.mode == MessageRequestMode::Pop {
+                        self.allocate_for_pop(
+                            strategy,
+                            consumer_group,
+                            client_id,
+                            mq_all.as_slice(),
+                            cid_all.as_slice(),
+                            set_message_request_mode_request_body.pop_share_queue_num,
+                        )
+                    } else {
+                        match strategy.allocate(
+                            consumer_group,
+                            client_id,
+                            mq_all.as_slice(),
+                            cid_all.as_slice(),
+                        ) {
+                            Ok(value) => Ok(value.into_iter().collect::<HashSet<MessageQueue>>()),
+                            Err(e) => Err(BrokerError::ClientError(e)),
+                        }
+                    };
+                match result {
+                    Ok(value) => Some(value),
+                    Err(_) => None,
+                }
+            }
+        }
+    }
+
+    pub fn allocate_for_pop(
+        &self,
+        strategy: &Arc<dyn AllocateMessageQueueStrategy>,
+        consumer_group: &CheetahString,
+        current_cid: &CheetahString,
+        mq_all: &[MessageQueue],
+        cid_all: &[CheetahString],
+        pop_share_queue_num: i32,
+    ) -> Result<HashSet<MessageQueue>> {
+        if pop_share_queue_num <= 0 || pop_share_queue_num > cid_all.len() as i32 - 1 {
+            Ok(mq_all
+                .iter()
+                .map(|mq| {
+                    MessageQueue::from_parts(
+                        mq.get_topic_cs().clone(),
+                        mq.get_broker_name().clone(),
+                        -1,
+                    )
+                })
+                .collect::<HashSet<MessageQueue>>())
+        } else if cid_all.len() <= mq_all.len() {
+            let mut allocate_result = strategy
+                .allocate(consumer_group, current_cid, mq_all, cid_all)
+                .unwrap();
+            let index = cid_all.iter().position(|cid| cid == current_cid);
+            if let Some(mut index) = index {
+                for _i in 1..pop_share_queue_num {
+                    index += index;
+                    index %= cid_all.len();
+                    let result = strategy
+                        .allocate(consumer_group, &cid_all[index], mq_all, cid_all)
+                        .unwrap();
+                    allocate_result.extend(result);
+                }
+            }
+            Ok(allocate_result
+                .into_iter()
+                .collect::<HashSet<MessageQueue>>())
+        } else {
+            allocate(consumer_group, current_cid, mq_all, cid_all)
+        }
     }
 
     async fn set_message_request_mode(
@@ -117,3 +351,39 @@ impl QueryAssignmentProcessor {
         ))
     }
 }
+
+fn allocate(
+    consumer_group: &CheetahString,
+    current_cid: &CheetahString,
+    mq_all: &[MessageQueue],
+    cid_all: &[CheetahString],
+) -> Result<HashSet<MessageQueue>> {
+    if current_cid.is_empty() {
+        return Err(IllegalArgumentError("currentCID is empty".to_string()));
+    }
+    if mq_all.is_empty() {
+        return Err(IllegalArgumentError(
+            "mqAll is null or mqAll empty".to_string(),
+        ));
+    }
+    if cid_all.is_empty() {
+        return Err(IllegalArgumentError(
+            "cidAll is null or cidAll empty".to_string(),
+        ));
+    }
+
+    let mut result = HashSet::new();
+    if !cid_all.contains(current_cid) {
+        log::info!(
+            "[BUG] ConsumerGroup: {} The consumerId: {} not in cidAll: {:?}",
+            consumer_group,
+            current_cid,
+            cid_all
+        );
+        return Ok(result);
+    }
+
+    let index = cid_all.iter().position(|cid| cid == current_cid).unwrap();
+    result.insert(mq_all[index % mq_all.len()].clone());
+    Ok(result)
+}
diff --git a/rocketmq-common/src/common/broker/broker_config.rs b/rocketmq-common/src/common/broker/broker_config.rs
