rocketmq_remoting/protocol/admin/
topic_stats_table.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;
16
17use rocketmq_common::common::message::message_queue::MessageQueue;
18use serde::Deserialize;
19use serde::Serialize;
20use serde_json_any_key::*;
21
22use crate::protocol::admin::consume_stats::normalize_nonstandard_offset_table_keys;
23use crate::protocol::admin::topic_offset::TopicOffset;
24use crate::protocol::RemotingDeserializable;
25
26#[derive(Debug, Clone, Serialize, Deserialize, Default)]
27#[serde(rename_all = "camelCase")]
28pub struct TopicStatsTable {
29    topic_put_tps: f64,
30    #[serde(with = "any_key_map")]
31    offset_table: HashMap<MessageQueue, TopicOffset>,
32}
33
34impl TopicStatsTable {
35    pub fn new() -> Self {
36        Self {
37            topic_put_tps: 0.0,
38            offset_table: HashMap::new(),
39        }
40    }
41
42    pub fn get_topic_put_tps(&self) -> f64 {
43        self.topic_put_tps
44    }
45
46    pub fn set_topic_put_tps(&mut self, topic_put_tps: f64) {
47        self.topic_put_tps = topic_put_tps;
48    }
49
50    pub fn get_offset_table(&self) -> &HashMap<MessageQueue, TopicOffset> {
51        &self.offset_table
52    }
53
54    pub fn get_offset_table_mut(&mut self) -> &mut HashMap<MessageQueue, TopicOffset> {
55        &mut self.offset_table
56    }
57
58    pub fn into_offset_table(self) -> HashMap<MessageQueue, TopicOffset> {
59        self.offset_table
60    }
61
62    pub fn set_offset_table(&mut self, offset_table: HashMap<MessageQueue, TopicOffset>) {
63        self.offset_table = offset_table;
64    }
65
66    pub fn decode(body: &[u8]) -> rocketmq_error::RocketMQResult<Self> {
67        match <Self as RemotingDeserializable>::decode(body) {
68            Ok(stats) => Ok(stats),
69            Err(error) => {
70                let Ok(raw_body) = std::str::from_utf8(body) else {
71                    return Err(error);
72                };
73                let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
74                if normalized_body == raw_body {
75                    return Err(error);
76                }
77                <Self as RemotingDeserializable>::decode_str(&normalized_body)
78            }
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    fn create_custom_mq(topic: &str, queue_id: i32) -> MessageQueue {
88        let json = format!(
89            r#"{{"topic": "{}", "brokerName": "default", "queueId": {}}}"#,
90            topic, queue_id
91        );
92        serde_json::from_str(&json).unwrap()
93    }
94
95    #[test]
96    fn test_topic_stats_table_initialization() {
97        let table = TopicStatsTable::new();
98        assert!(table.get_offset_table().is_empty());
99    }
100
101    #[test]
102    fn test_message_queue_default_new() {
103        let mq = MessageQueue::new();
104        let json = serde_json::to_value(&mq).unwrap();
105        assert_eq!(json["queueId"], 0);
106        assert_eq!(json["topic"], "");
107    }
108
109    #[test]
110    fn test_set_and_get_offsets() {
111        let mut table = TopicStatsTable::new();
112        let mut map = HashMap::new();
113
114        let mq = create_custom_mq("test_topic", 5);
115        let mut offset = TopicOffset::new();
116        offset.set_min_offset(100);
117        offset.set_max_offset(200);
118
119        map.insert(mq.clone(), offset);
120        table.set_offset_table(map);
121        let result_table = table.get_offset_table();
122        assert_eq!(result_table.len(), 1);
123
124        let retrieved_offset = result_table.get(&mq).expect("MQ should exist in map");
125        assert_eq!(retrieved_offset.get_min_offset(), 100);
126        assert_eq!(retrieved_offset.get_max_offset(), 200);
127    }
128
129    #[test]
130    fn test_serialization_cycle_with_any_key() {
131        let mut table = TopicStatsTable::new();
132        let mut map = HashMap::new();
133
134        let mq = create_custom_mq("order_topic", 1);
135        let mut offset = TopicOffset::new();
136        offset.set_last_update_timestamp(11111111);
137
138        map.insert(mq, offset);
139        table.set_offset_table(map);
140
141        let serialized = serde_json::to_string(&table).expect("Serialization failed");
142
143        assert!(serialized.contains("offsetTable"));
144
145        let deserialized: TopicStatsTable = serde_json::from_str(&serialized).expect("Deserialization failed");
146        assert_eq!(deserialized.get_offset_table().len(), 1);
147
148        let offset_val = deserialized.get_offset_table().values().next().unwrap().clone();
149        assert_eq!(offset_val.get_last_update_timestamp(), 11111111);
150    }
151
152    #[test]
153    fn test_decode_java_fastjson_message_queue_keys() {
154        let body = br#"{"offsetTable":{{"topic":"TBW102","brokerName":"broker-a","queueId":0}:{"minOffset":0,"maxOffset":0,"lastUpdateTimestamp":0}},"topicPutTps":1.5}"#;
155
156        let table = TopicStatsTable::decode(body).expect("decode Java fastjson topic stats");
157
158        let queue = MessageQueue::from_parts("TBW102", "broker-a", 0);
159        let offset = table
160            .get_offset_table()
161            .get(&queue)
162            .expect("queue offset should decode");
163        assert_eq!(offset.get_min_offset(), 0);
164        assert_eq!(offset.get_max_offset(), 0);
165        assert_eq!(table.get_topic_put_tps(), 1.5);
166    }
167
168    #[test]
169    fn test_topic_offset_display_format() {
170        let mut offset = TopicOffset::new();
171        offset.set_min_offset(5);
172        offset.set_max_offset(15);
173
174        let output = format!("{}", offset);
175        let expected = "TopicOffset{min_offset=5, max_offset=15, last_update_timestamp=0}";
176        assert_eq!(output, expected);
177    }
178}
rocketmq_remoting/protocol/admin/topic_stats_table.rs

rocketmq_remoting/protocol/admin/
topic_stats_table.rs
