rocketmq_protocol/protocol/body/
consumer_connection.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::collections::HashSet;
17
18use cheetah_string::CheetahString;
19use rocketmq_model::common::consumer::consume_from_where::ConsumeFromWhere;
20use serde::Deserialize;
21use serde::Serialize;
22
23use crate::protocol::body::connection::Connection;
24use crate::protocol::heartbeat::consume_type::ConsumeType;
25use crate::protocol::heartbeat::message_model::MessageModel;
26use crate::protocol::heartbeat::subscription_data::SubscriptionData;
27
28/// Represents a consumer connection with subscription information.
29///
30/// This struct is serialization-compatible with Java's ConsumerConnection class.
31/// It contains consumer connection details, subscription data, and consumption configuration.
32#[derive(Debug, Clone, Default, Serialize, Deserialize)]
33#[serde(rename_all = "camelCase")]
34pub struct ConsumerConnection {
35    /// Set of active connections from consumers
36    connection_set: HashSet<Connection>,
37    /// Topic subscription table (Topic -> SubscriptionData)
38    #[serde(default)]
39    subscription_table: HashMap<CheetahString, SubscriptionData>,
40    /// Type of consumption (push/pull)
41    #[serde(skip_serializing_if = "Option::is_none")]
42    consume_type: Option<ConsumeType>,
43    /// Message consumption model (clustering/broadcasting)
44    #[serde(skip_serializing_if = "Option::is_none")]
45    message_model: Option<MessageModel>,
46    /// Position to start consuming from
47    #[serde(skip_serializing_if = "Option::is_none")]
48    consume_from_where: Option<ConsumeFromWhere>,
49}
50
51impl ConsumerConnection {
52    /// Creates a new ConsumerConnection instance.
53    pub fn new() -> Self {
54        ConsumerConnection {
55            connection_set: HashSet::new(),
56            subscription_table: HashMap::new(),
57            consume_type: None,
58            message_model: None,
59            consume_from_where: None,
60        }
61    }
62
63    /// Computes the minimum version across all connections.
64    ///
65    /// This method iterates through all connections and finds the minimum version number.
66    /// Returns `i32::MAX` if there are no connections.
67    ///
68    /// # Returns
69    /// The minimum version number, or `i32::MAX` if connection set is empty
70    pub fn compute_min_version(&self) -> i32 {
71        let mut min_version = i32::MAX;
72        for connection in &self.connection_set {
73            let version = connection.get_version();
74            if version < min_version {
75                min_version = version;
76            }
77        }
78        min_version
79    }
80
81    /// Gets a clone of the connection set.
82    ///
83    /// # Returns
84    /// A cloned HashSet of all connections
85    pub fn get_connection_set(&self) -> &HashSet<Connection> {
86        &self.connection_set
87    }
88
89    /// Sets the connection set.
90    ///
91    /// # Arguments
92    /// * `connection_set` - The new set of connections
93    pub fn set_connection_set(&mut self, connection_set: HashSet<Connection>) {
94        self.connection_set = connection_set;
95    }
96
97    /// Inserts a connection into the connection set.
98    ///
99    /// # Arguments
100    /// * `connection` - The connection to insert
101    ///
102    /// # Returns
103    /// `true` if the connection was newly inserted, `false` if it already existed
104    pub fn insert_connection(&mut self, connection: Connection) -> bool {
105        self.connection_set.insert(connection)
106    }
107
108    /// Gets a reference to the subscription table.
109    ///
110    /// # Returns
111    /// Reference to the subscription table (Topic -> SubscriptionData mapping)
112    pub fn get_subscription_table(&self) -> &HashMap<CheetahString, SubscriptionData> {
113        &self.subscription_table
114    }
115
116    /// Gets a mutable reference to the subscription table.
117    ///
118    /// # Returns
119    /// Mutable reference to the subscription table
120    pub fn get_subscription_table_mut(&mut self) -> &mut HashMap<CheetahString, SubscriptionData> {
121        &mut self.subscription_table
122    }
123
124    /// Sets the subscription table.
125    ///
126    /// # Arguments
127    /// * `subscription_table` - The new subscription table
128    pub fn set_subscription_table(&mut self, subscription_table: HashMap<CheetahString, SubscriptionData>) {
129        self.subscription_table = subscription_table;
130    }
131
132    /// Gets the consume type.
133    ///
134    /// # Returns
135    /// The consume type, or None if not set
136    pub fn get_consume_type(&self) -> Option<ConsumeType> {
137        self.consume_type
138    }
139
140    /// Sets the consume type.
141    ///
142    /// # Arguments
143    /// * `consume_type` - The consume type (push/pull)
144    pub fn set_consume_type(&mut self, consume_type: ConsumeType) {
145        self.consume_type = Some(consume_type);
146    }
147
148    /// Gets the message model.
149    ///
150    /// # Returns
151    /// The message model, or None if not set
152    pub fn get_message_model(&self) -> Option<MessageModel> {
153        self.message_model
154    }
155
156    /// Sets the message model.
157    ///
158    /// # Arguments
159    /// * `message_model` - The message model (clustering/broadcasting)
160    pub fn set_message_model(&mut self, message_model: MessageModel) {
161        self.message_model = Some(message_model);
162    }
163
164    /// Gets the consume from where setting.
165    ///
166    /// # Returns
167    /// The consume from where setting, or None if not set
168    pub fn get_consume_from_where(&self) -> Option<ConsumeFromWhere> {
169        self.consume_from_where
170    }
171
172    /// Sets the consume from where setting.
173    ///
174    /// # Arguments
175    /// * `consume_from_where` - Where to start consuming from
176    pub fn set_consume_from_where(&mut self, consume_from_where: ConsumeFromWhere) {
177        self.consume_from_where = Some(consume_from_where);
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn create_test_connection() -> Connection {
186        let mut conn = Connection::new();
187        conn.set_client_id(CheetahString::from_static_str("test_client_1"));
188        conn.set_client_addr(CheetahString::from_static_str("127.0.0.1:9876"));
189        conn.set_version(1);
190        conn
191    }
192
193    fn create_test_subscription() -> SubscriptionData {
194        SubscriptionData {
195            topic: CheetahString::from_static_str("test_topic"),
196            sub_string: CheetahString::from_static_str("*"),
197            ..Default::default()
198        }
199    }
200
201    #[test]
202    fn consumer_connection_serializes_correctly() {
203        let mut consumer_conn = ConsumerConnection::new();
204        let conn = create_test_connection();
205        consumer_conn.insert_connection(conn);
206        consumer_conn.set_consume_type(ConsumeType::ConsumeActively);
207        consumer_conn.set_message_model(MessageModel::Clustering);
208
209        let json_str = serde_json::to_string(&consumer_conn).unwrap();
210        assert!(json_str.contains("connectionSet"));
211        assert!(json_str.contains("consumeType"));
212        assert!(json_str.contains("messageModel"));
213    }
214
215    #[test]
216    fn consumer_connection_deserializes_correctly() {
217        let json_str = r#"{
218            "connectionSet": [
219                {
220                    "clientId": "test_client",
221                    "clientAddr": "127.0.0.1:9876",
222                    "language": "RUST",
223                    "version": 1
224                }
225            ],
226            "subscriptionTable": {
227                "test_topic": {
228                    "classFilterMode": false,
229                    "topic": "test_topic",
230                    "subString": "*",
231                    "tagsSet": [],
232                    "codeSet": [],
233                    "subVersion": 123456,
234                    "expressionType": "TAG"
235                }
236            },
237            "consumeType": "CONSUME_ACTIVELY",
238            "messageModel": "CLUSTERING",
239            "consumeFromWhere": "CONSUME_FROM_LAST_OFFSET"
240        }"#;
241
242        let consumer_conn: ConsumerConnection = serde_json::from_str(json_str).unwrap();
243        assert_eq!(consumer_conn.get_connection_set().len(), 1);
244        assert_eq!(consumer_conn.get_subscription_table().len(), 1);
245        assert!(consumer_conn.get_consume_type().is_some());
246        assert!(consumer_conn.get_message_model().is_some());
247        assert!(consumer_conn.get_consume_from_where().is_some());
248    }
249
250    #[test]
251    fn consumer_connection_compute_min_version() {
252        let consumer_conn = ConsumerConnection::new();
253        assert_eq!(consumer_conn.compute_min_version(), i32::MAX);
254
255        let mut consumer_conn = ConsumerConnection::new();
256        let mut conn = Connection::new();
257        conn.set_version(42);
258        consumer_conn.insert_connection(conn);
259        assert_eq!(consumer_conn.compute_min_version(), 42);
260
261        let mut consumer_conn = ConsumerConnection::new();
262        let mut conn1 = Connection::new();
263        conn1.set_client_id(CheetahString::from_static_str("client1"));
264        conn1.set_version(100);
265        let mut conn2 = Connection::new();
266        conn2.set_client_id(CheetahString::from_static_str("client2"));
267        conn2.set_version(50);
268        let mut conn3 = Connection::new();
269        conn3.set_client_id(CheetahString::from_static_str("client3"));
270        conn3.set_version(75);
271        consumer_conn.insert_connection(conn1);
272        consumer_conn.insert_connection(conn2);
273        consumer_conn.insert_connection(conn3);
274        assert_eq!(consumer_conn.compute_min_version(), 50);
275    }
276
277    #[test]
278    fn consumer_connection_set_connection_set() {
279        let mut consumer_conn = ConsumerConnection::new();
280        let mut conn_set = HashSet::new();
281        conn_set.insert(create_test_connection());
282
283        consumer_conn.set_connection_set(conn_set.clone());
284        assert_eq!(consumer_conn.get_connection_set().len(), 1);
285    }
286
287    #[test]
288    fn consumer_connection_insert_connection() {
289        let mut consumer_conn = ConsumerConnection::new();
290        let conn = create_test_connection();
291
292        let result = consumer_conn.insert_connection(conn.clone());
293        assert!(result);
294
295        let result = consumer_conn.insert_connection(conn.clone());
296        assert!(!result);
297
298        assert_eq!(consumer_conn.get_connection_set().len(), 1);
299        assert!(consumer_conn.get_connection_set().contains(&conn));
300    }
301
302    #[test]
303    fn consumer_connection_subscription_table_operations() {
304        let mut consumer_conn = ConsumerConnection::new();
305        let mut sub_table = HashMap::new();
306
307        let sub_data = create_test_subscription();
308        sub_table.insert(CheetahString::from_static_str("test_topic"), sub_data.clone());
309
310        consumer_conn.set_subscription_table(sub_table);
311        assert_eq!(consumer_conn.get_subscription_table().len(), 1);
312        assert!(consumer_conn
313            .get_subscription_table()
314            .contains_key(&CheetahString::from_static_str("test_topic")));
315
316        let mut consumer_conn = ConsumerConnection::new();
317        let sub_data = create_test_subscription();
318
319        {
320            let sub_table = consumer_conn.get_subscription_table_mut();
321            sub_table.insert(CheetahString::from_static_str("test_topic"), sub_data.clone());
322        }
323
324        assert_eq!(consumer_conn.get_subscription_table().len(), 1);
325        let retrieved_sub = consumer_conn
326            .get_subscription_table()
327            .get(&CheetahString::from_static_str("test_topic"));
328        assert!(retrieved_sub.is_some());
329    }
330
331    #[test]
332    fn consumer_connection_consumption_options() {
333        let mut consumer_conn = ConsumerConnection::new();
334        assert!(consumer_conn.get_consume_type().is_none());
335        assert!(consumer_conn.get_message_model().is_none());
336        assert!(consumer_conn.get_consume_from_where().is_none());
337
338        consumer_conn.set_consume_type(ConsumeType::ConsumeActively);
339        consumer_conn.set_message_model(MessageModel::Clustering);
340        consumer_conn.set_consume_from_where(ConsumeFromWhere::ConsumeFromLastOffset);
341        assert_eq!(consumer_conn.get_consume_type(), Some(ConsumeType::ConsumeActively));
342        assert_eq!(consumer_conn.get_message_model(), Some(MessageModel::Clustering));
343        assert_eq!(
344            consumer_conn.get_consume_from_where(),
345            Some(ConsumeFromWhere::ConsumeFromLastOffset)
346        );
347    }
348}
rocketmq_protocol/protocol/body/consumer_connection.rs

rocketmq_protocol/protocol/body/
consumer_connection.rs
