[ISSUE #790]🚀Support Query message(request code :12,33)🔥

rocketmq-broker/src/broker_runtime.rs

@@ -66,6 +66,7 @@
 use crate::processor::default_pull_message_result_handler::DefaultPullMessageResultHandler;
 use crate::processor::pull_message_processor::PullMessageProcessor;
 use crate::processor::pull_message_result_handler::PullMessageResultHandler;
+use crate::processor::query_message_processor::QueryMessageProcessor;
 use crate::processor::send_message_processor::SendMessageProcessor;
 use crate::processor::BrokerRequestProcessor;
 use crate::schedule::schedule_message_service::ScheduleMessageService;
@@ -410,7 +411,7 @@
             self.message_store.clone().unwrap(),
         );
         self.pull_request_hold_service = Some(PullRequestHoldService::new(
-            message_store,
+            message_store.clone(),
             Arc::new(pull_message_processor.clone()),
             self.broker_config.clone(),
         ));
@@ -430,6 +431,8 @@
             .set_message_arriving_listener(Some(Arc::new(Box::new(
                 NotifyMessageArrivingListener::new(self.pull_request_hold_service.clone().unwrap()),
             ))));
+        let query_message_processor =
+            QueryMessageProcessor::new(self.message_store_config.clone(), message_store.clone());
 
         let admin_broker_processor = AdminBrokerProcessor::new(
             self.broker_config.clone(),
@@ -459,7 +462,7 @@
             ),
             consumer_manage_processor,
             query_assignment_processor: Default::default(),
-            query_message_processor: Default::default(),
+            query_message_processor,
             end_transaction_processor: Default::default(),
         }
     }

rocketmq-broker/src/processor.rs

@@ -95,7 +95,7 @@
     pub(crate) notification_processor: NotificationProcessor,
     pub(crate) polling_info_processor: PollingInfoProcessor,
     pub(crate) reply_message_processor: ReplyMessageProcessor,
-    pub(crate) query_message_processor: QueryMessageProcessor,
+    pub(crate) query_message_processor: QueryMessageProcessor<MS>,
     pub(crate) client_manage_processor: ClientManageProcessor<MS>,
     pub(crate) consumer_manage_processor: ConsumerManageProcessor<MS>,
     pub(crate) query_assignment_processor: QueryAssignmentProcessor,
@@ -161,6 +161,13 @@
                     .process_request(channel, ctx, request_code, request)
                     .await
             }
+
+            RequestCode::QueryMessage | RequestCode::ViewMessageById => {
+                self.query_message_processor
+                    .process_request(channel, ctx, request_code, request)
+                    .await
+            }
+
             _ => {
                 self.admin_broker_processor
                     .process_request(channel, ctx, request_code, request)

rocketmq-broker/src/processor/query_message_processor.rs

@@ -14,19 +14,127 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+use std::sync::Arc;
 
+use rocketmq_common::common::mix_all::UNIQUE_MSG_QUERY_FLAG;
+use rocketmq_remoting::code::request_code::RequestCode;
+use rocketmq_remoting::code::response_code::ResponseCode;
+use rocketmq_remoting::net::channel::Channel;
+use rocketmq_remoting::protocol::header::query_message_request_header::QueryMessageRequestHeader;
+use rocketmq_remoting::protocol::header::query_message_response_header::QueryMessageResponseHeader;
+use rocketmq_remoting::protocol::header::view_message_request_header::ViewMessageRequestHeader;
 use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
 use rocketmq_remoting::runtime::server::ConnectionHandlerContext;
+use rocketmq_store::config::message_store_config::MessageStoreConfig;
+use rocketmq_store::log_file::MessageStore;
 
 #[derive(Default, Clone)]
-pub struct QueryMessageProcessor {}
+pub struct QueryMessageProcessor<MS> {
+    message_store_config: Arc<MessageStoreConfig>,
+    message_store: Arc<MS>,
+}
+
+impl<MS> QueryMessageProcessor<MS> {
+    pub fn new(message_store_config: Arc<MessageStoreConfig>, message_store: Arc<MS>) -> Self {
+        Self {
+            message_store_config,
+            message_store,
+        }
+    }
+}
+
+impl<MS> QueryMessageProcessor<MS>
+where
+    MS: MessageStore + Send + Sync + 'static,
+{
+    pub async fn process_request(
+        &mut self,
+        channel: Channel,
+        ctx: ConnectionHandlerContext,
+        request_code: RequestCode,
+        request: RemotingCommand,
+    ) -> Option<RemotingCommand> {
+        match request_code {
+            RequestCode::QueryMessage => self.query_message(channel, ctx, request).await,
+            RequestCode::ViewMessageById => self.view_message_by_id(channel, ctx, request).await,
+            _ => None,
+        }
+    }
+
+    async fn query_message(
+        &mut self,
+        _channel: Channel,
+        _ctx: ConnectionHandlerContext,
+        request: RemotingCommand,
+    ) -> Option<RemotingCommand> {
+        let mut response = RemotingCommand::create_response_command_with_header(
+            QueryMessageResponseHeader::default(),
+        );
+        let mut request_header = request
+            .decode_command_custom_header::<QueryMessageRequestHeader>()
+            .unwrap();
+        response.set_opaque_mut(request.opaque());
+        let is_unique_key = request.ext_fields().unwrap().get(UNIQUE_MSG_QUERY_FLAG);
+        if is_unique_key.is_some() && is_unique_key.unwrap() == "true" {
+            request_header.max_num = self.message_store_config.default_query_max_num as i32;
+        }
+        let query_message_result = self
+            .message_store
+            .query_message(
+                request_header.topic.as_str(),
+                request_header.key.as_str(),
+                request_header.max_num,
+                request_header.begin_timestamp,
+                request_header.end_timestamp,
+            )
+            .await?;
+
+        let response_header = response
+            .read_custom_header_mut::<QueryMessageResponseHeader>()
+            .unwrap();
+        response_header.index_last_update_phyoffset =
+            query_message_result.index_last_update_phyoffset;
+        response_header.index_last_update_timestamp =
+            query_message_result.index_last_update_timestamp;
+
+        if query_message_result.buffer_total_size > 0 {
+            let message_data = query_message_result.get_message_data();
+            return Some(response.set_body(message_data));
+        }
+        Some(
+            response
+                .set_code(ResponseCode::QueryNotFound)
+                .set_remark(Some(
+                    "can not find message, maybe time range not correct".to_string(),
+                )),
+        )
+    }
 
-impl QueryMessageProcessor {
-    fn process_request(
-        &self,
+    async fn view_message_by_id(
+        &mut self,
+        _channel: Channel,
         _ctx: ConnectionHandlerContext,
-        _request: RemotingCommand,
-    ) -> RemotingCommand {
-        todo!()
+        request: RemotingCommand,
+    ) -> Option<RemotingCommand> {
+        let response = RemotingCommand::create_response_command();
+        let request_header = request
+            .decode_command_custom_header::<ViewMessageRequestHeader>()
+            .unwrap();
+        let select_mapped_buffer_result = self
+            .message_store
+            .select_one_message_by_offset(request_header.offset)
+            .await;
+        if let Some(result) = select_mapped_buffer_result {
+            let message_data = result.get_bytes();
+            return Some(response.set_body(message_data));
+        }
+        Some(
+            response
+                .set_code(ResponseCode::SystemError)
+                .set_remark(Some(format!(
+                    "can not find message by offset: {}",
+                    request_header.offset
+                ))),
+        )
     }
 }

rocketmq-remoting/src/protocol/header.rs

@@ -30,6 +30,10 @@ pub mod pull_message_request_header;
 pub mod pull_message_response_header;
 pub mod query_consumer_offset_request_header;
 pub mod query_consumer_offset_response_header;
+pub mod query_message_request_header;
+pub mod query_message_response_header;
 pub mod search_offset_response_header;
 pub mod unregister_client_request_header;
 pub mod update_consumer_offset_header;
+pub mod view_message_request_header;
+pub mod view_message_response_header;

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

@@ -0,0 +1,184 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::protocol::command_custom_header::CommandCustomHeader;
+use crate::protocol::command_custom_header::FromMap;
+use crate::rpc::topic_request_header::TopicRequestHeader;
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct QueryMessageRequestHeader {
+    pub topic: String,
+
+    pub key: String,
+
+    pub max_num: i32,
+
+    pub begin_timestamp: i64,
+
+    pub end_timestamp: i64,
+
+    #[serde(flatten)]
+    pub topic_request_header: Option<TopicRequestHeader>,
+}
+
+impl QueryMessageRequestHeader {
+    pub const TOPIC: &'static str = "topic";
+    pub const KEY: &'static str = "key";
+    pub const MAX_NUM: &'static str = "maxNum";
+    pub const BEGIN_TIMESTAMP: &'static str = "beginTimestamp";
+    pub const END_TIMESTAMP: &'static str = "endTimestamp";
+}
+
+/*impl From<&QueryMessageRequestHeader> for String {
+    fn from(header: &QueryMessageRequestHeader) -> Self {
+        SerdeJsonUtils::to_json(header).unwrap()
+    }
+}*/
+
+impl CommandCustomHeader for QueryMessageRequestHeader {
+    fn to_map(&self) -> Option<std::collections::HashMap<String, String>> {
+        let mut map = std::collections::HashMap::new();
+        map.insert(Self::TOPIC.to_string(), self.topic.clone());
+        map.insert(Self::KEY.to_string(), self.key.clone());
+        map.insert(Self::MAX_NUM.to_string(), self.max_num.to_string());
+        map.insert(
+            Self::BEGIN_TIMESTAMP.to_string(),
+            self.begin_timestamp.to_string(),
+        );
+        map.insert(
+            Self::END_TIMESTAMP.to_string(),
+            self.end_timestamp.to_string(),
+        );
+        if let Some(value) = self.topic_request_header.as_ref() {
+            if let Some(val) = value.to_map() {
+                map.extend(val);
+            }
+        }
+        Some(map)
+    }
+}
+
+impl FromMap for QueryMessageRequestHeader {
+    type Target = Self;
+
+    fn from(map: &std::collections::HashMap<String, String>) -> Option<Self::Target> {
+        Some(QueryMessageRequestHeader {
+            topic: map.get(Self::TOPIC)?.clone(),
+            key: map.get(Self::KEY)?.clone(),
+            max_num: map.get(Self::MAX_NUM)?.parse().ok()?,
+            begin_timestamp: map.get(Self::BEGIN_TIMESTAMP)?.parse().ok()?,
+            end_timestamp: map.get(Self::END_TIMESTAMP)?.parse().ok()?,
+            topic_request_header: <TopicRequestHeader as FromMap>::from(map),
+        })
+    }
+}
+
+#[cfg(test)]
+mod query_message_request_header_tests {
+    use std::collections::HashMap;
+
+    use super::*;
+
+    #[test]
+    fn creating_from_map_with_all_fields_populates_struct_correctly() {
+        let mut map = HashMap::new();
+        map.insert("topic".to_string(), "test_topic".to_string());
+        map.insert("key".to_string(), "test_key".to_string());
+        map.insert("maxNum".to_string(), "10".to_string());
+        map.insert("beginTimestamp".to_string(), "1000".to_string());
+        map.insert("endTimestamp".to_string(), "2000".to_string());
+
+        let header = <QueryMessageRequestHeader as FromMap>::from(&map).unwrap();
+
+        assert_eq!(header.topic, "test_topic");
+        assert_eq!(header.key, "test_key");
+        assert_eq!(header.max_num, 10);
+        assert_eq!(header.begin_timestamp, 1000);
+        assert_eq!(header.end_timestamp, 2000);
+    }
+
+    #[test]
+    fn creating_from_map_missing_optional_fields_still_succeeds() {
+        let mut map = HashMap::new();
+        map.insert("topic".to_string(), "test_topic".to_string());
+        map.insert("key".to_string(), "test_key".to_string());
+        map.insert("maxNum".to_string(), "10".to_string());
+        map.insert("beginTimestamp".to_string(), "1000".to_string());
+        map.insert("endTimestamp".to_string(), "2000".to_string());
+
+        let header = <QueryMessageRequestHeader as FromMap>::from(&map).unwrap();
+
+        assert_eq!(header.topic, "test_topic");
+        assert_eq!(header.key, "test_key");
+        assert_eq!(header.max_num, 10);
+    }
+
+    #[test]
+    fn creating_from_map_with_invalid_number_fields_returns_none() {
+        let mut map = HashMap::new();
+        map.insert("topic".to_string(), "test_topic".to_string());
+        map.insert("key".to_string(), "test_key".to_string());
+        map.insert("maxNum".to_string(), "invalid".to_string());
+
+        let header = <QueryMessageRequestHeader as FromMap>::from(&map);
+
+        assert!(header.is_none());
+    }
+
+    #[test]
+    fn to_map_includes_all_fields() {
+        let header = QueryMessageRequestHeader {
+            topic: "test_topic".to_string(),
+            key: "test_key".to_string(),
+            max_num: 10,
+            begin_timestamp: 1000,
+            end_timestamp: 2000,
+            topic_request_header: None,
+        };
+
+        let map = header.to_map().unwrap();
+
+        assert_eq!(map.get("topic").unwrap(), "test_topic");
+        assert_eq!(map.get("key").unwrap(), "test_key");
+        assert_eq!(map.get("maxNum").unwrap(), "10");
+        assert_eq!(map.get("beginTimestamp").unwrap(), "1000");
+        assert_eq!(map.get("endTimestamp").unwrap(), "2000");
+    }
+
+    #[test]
+    fn to_map_with_topic_request_header_includes_nested_fields() {
+        let topic_request_header = TopicRequestHeader::default();
+        let header = QueryMessageRequestHeader {
+            topic: "test_topic".to_string(),
+            key: "test_key".to_string(),
+            max_num: 10,
+            begin_timestamp: 1000,
+            end_timestamp: 2000,
+            topic_request_header: Some(topic_request_header),
+        };
+
+        let map = header.to_map().unwrap();
+
+        // Verify that nested fields are included, assuming specific fields in TopicRequestHeader
+        // This is a placeholder assertion; actual fields and checks depend on TopicRequestHeader's
+        // structure
+        assert!(!map.contains_key("nestedField"));
+    }
+}