[ISSUE #1303]🚀Add QueryAssignmentRequestBody and QueryAssignmentResponseBody (#1304)

Author: mxsm

rocketmq-common/src/common/message.rs

@@ -37,6 +37,7 @@ pub mod message_ext;
 pub mod message_ext_broker_inner;
 pub mod message_id;
 pub mod message_queue;
+pub mod message_queue_assignment;
 pub mod message_single;
 
 /// This module defines the `MessageTrait` trait, which provides a flexible interface for working

rocketmq-common/src/common/message/message_enum.rs

@@ -14,6 +14,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+use std::fmt;
+
+use serde::de;
+use serde::de::Visitor;
+use serde::Deserialize;
+use serde::Deserializer;
+use serde::Serialize;
+use serde::Serializer;
+
 #[derive(Debug, PartialEq, Copy, Clone, Default)]
 pub enum MessageType {
     #[default]
@@ -47,7 +56,7 @@ impl MessageType {
     }
 }
 
-#[derive(Debug, PartialEq, Copy, Clone)]
+#[derive(Debug, PartialEq, Copy, Clone, Hash, Eq)]
 pub enum MessageRequestMode {
     Pull,
     Pop,
@@ -62,6 +71,48 @@ impl MessageRequestMode {
     }
 }
 
+impl Serialize for MessageRequestMode {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str(match *self {
+            MessageRequestMode::Pull => "PULL",
+            MessageRequestMode::Pop => "POP",
+        })
+    }
+}
+
+impl<'de> Deserialize<'de> for MessageRequestMode {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct MessageRequestModeVisitor;
+
+        impl Visitor<'_> for MessageRequestModeVisitor {
+            type Value = MessageRequestMode;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                formatter.write_str("a string representing a MessageRequestMode")
+            }
+
+            fn visit_str<E>(self, value: &str) -> Result<MessageRequestMode, E>
+            where
+                E: de::Error,
+            {
+                match value {
+                    "PULL" => Ok(MessageRequestMode::Pull),
+                    "POP" => Ok(MessageRequestMode::Pop),
+                    _ => Err(de::Error::unknown_variant(value, &["PULL", "POP"])),
+                }
+            }
+        }
+
+        deserializer.deserialize_str(MessageRequestModeVisitor)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -108,4 +159,39 @@ mod tests {
         assert_eq!(MessageRequestMode::Pull.get_name(), "PULL");
         assert_eq!(MessageRequestMode::Pop.get_name(), "POP");
     }
+
+    #[test]
+    fn serialize_message_request_mode_pull() {
+        let mode = MessageRequestMode::Pull;
+        let serialized = serde_json::to_string(&mode).unwrap();
+        assert_eq!(serialized, "\"PULL\"");
+    }
+
+    #[test]
+    fn serialize_message_request_mode_pop() {
+        let mode = MessageRequestMode::Pop;
+        let serialized = serde_json::to_string(&mode).unwrap();
+        assert_eq!(serialized, "\"POP\"");
+    }
+
+    #[test]
+    fn deserialize_message_request_mode_pull() {
+        let json = "\"PULL\"";
+        let deserialized: MessageRequestMode = serde_json::from_str(json).unwrap();
+        assert_eq!(deserialized, MessageRequestMode::Pull);
+    }
+
+    #[test]
+    fn deserialize_message_request_mode_pop() {
+        let json = "\"POP\"";
+        let deserialized: MessageRequestMode = serde_json::from_str(json).unwrap();
+        assert_eq!(deserialized, MessageRequestMode::Pop);
+    }
+
+    #[test]
+    fn deserialize_message_request_mode_invalid() {
+        let json = "\"INVALID\"";
+        let deserialized: Result<MessageRequestMode, _> = serde_json::from_str(json);
+        assert!(deserialized.is_err());
+    }
 }

rocketmq-common/src/common/message/message_queue_assignment.rs

@@ -0,0 +1,57 @@
+/*
+ * 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 std::collections::HashMap;
+use std::hash::Hash;
+use std::hash::Hasher;
+
+use cheetah_string::CheetahString;
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::common::message::message_enum::MessageRequestMode;
+use crate::common::message::message_queue::MessageQueue;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "camelCase")]
+pub struct MessageQueueAssignment {
+    pub message_queue: Option<MessageQueue>,
+    pub mode: MessageRequestMode,
+    pub attachments: Option<HashMap<CheetahString, CheetahString>>,
+}
+
+impl Hash for MessageQueueAssignment {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.message_queue.hash(state);
+        self.mode.hash(state);
+        if let Some(ref attachments) = self.attachments {
+            for (key, value) in attachments {
+                key.hash(state);
+                value.hash(state);
+            }
+        }
+    }
+}
+
+impl Default for MessageQueueAssignment {
+    fn default() -> Self {
+        MessageQueueAssignment {
+            message_queue: None,
+            mode: MessageRequestMode::Pull,
+            attachments: None,
+        }
+    }
+}

rocketmq-remoting/src/protocol/body.rs

@@ -30,6 +30,8 @@ pub mod group_list;
 pub mod kv_table;
 pub mod pop_process_queue_info;
 pub mod process_queue_info;
+pub mod query_assignment_request_body;
+pub mod query_assignment_response_body;
 pub mod request;
 pub mod response;
 pub mod topic;

rocketmq-remoting/src/protocol/body/query_assignment_request_body.rs

@@ -0,0 +1,31 @@
+/*
+ * 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 cheetah_string::CheetahString;
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::protocol::heartbeat::message_model::MessageModel;
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct QueryAssignmentRequestBody {
+    pub topic: CheetahString,
+    pub consumer_group: CheetahString,
+    pub client_id: CheetahString,
+    pub strategy_name: CheetahString,
+    pub message_model: MessageModel,
+}

rocketmq-remoting/src/protocol/body/query_assignment_response_body.rs

@@ -0,0 +1,27 @@
+/*
+ * 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 std::collections::HashSet;
+
+use rocketmq_common::common::message::message_queue_assignment::MessageQueueAssignment;
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct QueryAssignmentResponseBody {
+    pub message_queue_assignments: HashSet<MessageQueueAssignment>,
+}