[ISSUE #433]👷Improving Methods of TopicQueueMappingManager Struct (#434)

* Improving Methods of TopicQueueMappingManager Struct

* fix ci error

* fix ci error

* fix ci error

Author: mxsm

rocketmq-broker/src/mqtrace/send_message_context.rs

@@ -74,8 +74,8 @@ impl SendMessageContext {
     pub fn broker_addr(&mut self, broker_addr: String) {
         self.broker_addr = broker_addr;
     }
-    pub fn queue_id(&mut self, queue_id: i32) {
-        self.queue_id = Some(queue_id);
+    pub fn queue_id(&mut self, queue_id: Option<i32>) {
+        self.queue_id = queue_id;
     }
     pub fn queue_offset(&mut self, queue_offset: i64) {
         self.queue_offset = Some(queue_offset);

rocketmq-broker/src/processor/send_message_processor.rs

@@ -101,9 +101,10 @@ impl<MS: MessageStore + Send> SendMessageProcessor<MS> {
                 let mapping_context = self
                     .topic_queue_mapping_manager
                     .build_topic_queue_mapping_context(&request_header, true);
-                let rewrite_result = self
-                    .topic_queue_mapping_manager
-                    .rewrite_request_for_static_topic(&request_header, &mapping_context);
+                let rewrite_result = TopicQueueMappingManager::rewrite_request_for_static_topic(
+                    &mut request_header,
+                    &mapping_context,
+                );
                 if let Some(rewrite_result) = rewrite_result {
                     return Some(rewrite_result);
                 }
@@ -222,13 +223,13 @@ impl<MS: MessageStore + Send + Clone> SendMessageProcessor<MS> {
             .select_topic_config(request_header.topic().as_str())
             .unwrap();
         let mut queue_id = request_header.queue_id;
-        if queue_id < 0 {
-            queue_id = self.inner.random_queue_id(topic_config.write_queue_nums) as i32;
+        if queue_id.is_none() || queue_id.unwrap() < 0 {
+            queue_id = Some(self.inner.random_queue_id(topic_config.write_queue_nums) as i32);
         }
 
         let mut message_ext = MessageExtBrokerInner::default();
         message_ext.message_ext_inner.message.topic = request_header.topic().to_string();
-        message_ext.message_ext_inner.queue_id = queue_id;
+        message_ext.message_ext_inner.queue_id = *queue_id.as_ref().unwrap();
         let mut ori_props =
             MessageDecoder::string_to_message_properties(request_header.properties.as_ref());
         if self.handle_retry_and_dlq(
@@ -321,7 +322,7 @@ impl<MS: MessageStore + Send + Clone> SendMessageProcessor<MS> {
             response,
             &request,
             topic.as_str(),
-            queue_id,
+            *queue_id.as_ref().unwrap(),
         )
     }
 

rocketmq-broker/src/topic/manager/topic_queue_mapping_manager.rs

@@ -18,16 +18,20 @@
 use std::{collections::HashMap, sync::Arc};
 
 use rocketmq_common::common::{broker::broker_config::BrokerConfig, config_manager::ConfigManager};
-use rocketmq_remoting::protocol::{
-    body::topic_info_wrapper::topic_queue_wrapper::TopicQueueMappingSerializeWrapper,
-    header::message_operation_header::TopicRequestHeaderTrait,
-    remoting_command::RemotingCommand,
-    static_topic::{
-        topic_queue_mapping_context::TopicQueueMappingContext,
-        topic_queue_mapping_detail::TopicQueueMappingDetail,
+use rocketmq_remoting::{
+    code::response_code::ResponseCode,
+    protocol::{
+        body::topic_info_wrapper::topic_queue_wrapper::TopicQueueMappingSerializeWrapper,
+        header::message_operation_header::TopicRequestHeaderTrait,
+        remoting_command::RemotingCommand,
+        static_topic::{
+            topic_queue_mapping_context::TopicQueueMappingContext,
+            topic_queue_mapping_detail::TopicQueueMappingDetail,
+        },
+        DataVersion, RemotingSerializable,
     },
-    DataVersion,
 };
+use tracing::{info, warn};
 
 use crate::broker_path_config_helper::get_topic_queue_mapping_path;
 
@@ -47,6 +51,35 @@ impl TopicQueueMappingManager {
         }
     }
 
+    pub(crate) fn rewrite_request_for_static_topic(
+        request_header: &mut impl TopicRequestHeaderTrait,
+        mapping_context: &TopicQueueMappingContext,
+    ) -> Option<RemotingCommand> {
+        mapping_context.mapping_detail.as_ref()?;
+
+        if !mapping_context.is_leader() {
+            let mapping_detail = mapping_context.mapping_detail.as_ref().unwrap();
+            return Some(RemotingCommand::create_response_command_with_code_remark(
+                ResponseCode::NotLeaderForQueue,
+                format!(
+                    "{}-{:?} does not exit in request process of current broker {}",
+                    request_header.topic(),
+                    request_header.queue_id(),
+                    mapping_detail
+                        .topic_queue_mapping_info
+                        .bname
+                        .clone()
+                        .unwrap_or_default()
+                ),
+            ));
+        }
+        let mapping_item = mapping_context.leader_item.as_ref().unwrap();
+        request_header.set_queue_id(Some(mapping_item.queue_id));
+        None
+    }
+}
+
+impl TopicQueueMappingManager {
     pub(crate) fn build_topic_queue_mapping_context(
         &self,
         request_header: &impl TopicRequestHeaderTrait,
@@ -66,13 +99,8 @@ impl TopicQueueMappingManager {
         }
         let topic = request_header.topic();
 
-        let mut global_id: Option<i32> = None;
-        /*if let Some(header) = request_header
-            .as_any()
-            .downcast_ref::<TopicQueueRequestHeaderT>()
-        {
-            global_id = header.queue_id;
-        }*/
+        let mut global_id: Option<i32> = request_header.queue_id();
+
         if let Some(mapping_detail) = self.topic_queue_mapping_table.lock().get(&topic) {
             // it is not static topic
             if mapping_detail
@@ -164,41 +192,45 @@ impl TopicQueueMappingManager {
         }
     }
 
-    pub(crate) fn rewrite_request_for_static_topic(
-        &self,
-        _request_header: &impl TopicRequestHeaderTrait,
-        _mapping_context: &TopicQueueMappingContext,
-    ) -> Option<RemotingCommand> {
-        //TODO
-        None
-    }
-
     pub fn get_topic_queue_mapping(&self, topic: &str) -> Option<TopicQueueMappingDetail> {
         self.topic_queue_mapping_table.lock().get(topic).cloned()
     }
+
+    pub fn delete(&self, topic: &str) {
+        let old = self.topic_queue_mapping_table.lock().remove(topic);
+        match old {
+            None => {
+                warn!(
+                    "delete topic queue mapping failed, static topic: {} not exists",
+                    topic
+                )
+            }
+            Some(value) => {
+                info!(
+                    "delete topic queue mapping OK, static topic queue mapping: {:?}",
+                    value
+                );
+                self.data_version.lock().next_version();
+                self.persist();
+            }
+        }
+    }
 }
 
 //Fully implemented will be removed
-#[allow(unused_variables)]
 impl ConfigManager for TopicQueueMappingManager {
-    fn decode0(&mut self, key: &[u8], body: &[u8]) {
-        todo!()
-    }
-
-    fn stop(&mut self) -> bool {
-        todo!()
-    }
-
     fn config_file_path(&self) -> String {
         get_topic_queue_mapping_path(self.broker_config.store_path_root_dir.as_str())
     }
-
-    fn encode(&mut self) -> String {
-        todo!()
-    }
-
     fn encode_pretty(&self, pretty_format: bool) -> String {
-        todo!()
+        let wrapper = TopicQueueMappingSerializeWrapper::new(
+            Some(self.topic_queue_mapping_table.lock().clone()),
+            Some(self.data_version.lock().clone()),
+        );
+        match pretty_format {
+            true => wrapper.to_json_pretty(),
+            false => wrapper.to_json(),
+        }
     }
 
     fn decode(&self, json_string: &str) {
@@ -219,3 +251,60 @@ impl ConfigManager for TopicQueueMappingManager {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use rocketmq_common::common::broker::broker_config::BrokerConfig;
+
+    use super::*;
+
+    #[test]
+    fn new_creates_default_manager() {
+        let broker_config = Arc::new(BrokerConfig::default());
+        let manager = TopicQueueMappingManager::new(broker_config.clone());
+
+        assert_eq!(Arc::ptr_eq(&manager.broker_config, &broker_config), true);
+        assert_eq!(manager.data_version.lock().get_state_version(), 0);
+        assert_eq!(manager.topic_queue_mapping_table.lock().len(), 0);
+    }
+
+    #[test]
+    fn get_topic_queue_mapping_returns_none_for_non_existent_topic() {
+        let broker_config = Arc::new(BrokerConfig::default());
+        let manager = TopicQueueMappingManager::new(broker_config);
+
+        assert!(manager
+            .get_topic_queue_mapping("non_existent_topic")
+            .is_none());
+    }
+
+    #[test]
+    fn get_topic_queue_mapping_returns_mapping_for_existing_topic() {
+        let broker_config = Arc::new(BrokerConfig::default());
+        let manager = TopicQueueMappingManager::new(broker_config);
+        let detail = TopicQueueMappingDetail::default();
+        manager
+            .topic_queue_mapping_table
+            .lock()
+            .insert("existing_topic".to_string(), detail.clone());
+
+        assert!(manager.get_topic_queue_mapping("existing_topic").is_some());
+    }
+
+    #[test]
+    fn delete_removes_existing_topic() {
+        let broker_config = Arc::new(BrokerConfig::default());
+        let manager = TopicQueueMappingManager::new(broker_config);
+        let detail = TopicQueueMappingDetail::default();
+        manager
+            .topic_queue_mapping_table
+            .lock()
+            .insert("existing_topic".to_string(), detail.clone());
+
+        manager.delete("existing_topic");
+
+        assert!(manager.get_topic_queue_mapping("existing_topic").is_none());
+    }
+}

rocketmq-remoting/src/protocol/body/topic_info_wrapper/topic_queue_wrapper.rs

@@ -31,6 +31,18 @@ pub struct TopicQueueMappingSerializeWrapper {
     data_version: Option<DataVersion>,
 }
 
+impl TopicQueueMappingSerializeWrapper {
+    pub fn new(
+        topic_queue_mapping_info_map: Option<HashMap<String, TopicQueueMappingDetail>>,
+        data_version: Option<DataVersion>,
+    ) -> Self {
+        Self {
+            topic_queue_mapping_info_map,
+            data_version,
+        }
+    }
+}
+
 impl TopicQueueMappingSerializeWrapper {
     pub fn topic_queue_mapping_info_map(
         &self,

rocketmq-remoting/src/protocol/header/message_operation_header.rs

@@ -41,4 +41,8 @@ pub trait TopicRequestHeaderTrait {
     fn oneway(&self) -> Option<bool>;
 
     fn with_oneway(&mut self, oneway: bool);
+
+    fn queue_id(&self) -> Option<i32>;
+
+    fn set_queue_id(&mut self, queue_id: Option<i32>);
 }