rocketmq_protocol/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_model::message::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::Result<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 methods_and_serde_preserve_topic_stats() {
97        let mut table = TopicStatsTable::new();
98        assert_eq!(table.get_topic_put_tps(), 0.0);
99        assert!(table.get_offset_table().is_empty());
100
101        table.set_topic_put_tps(2.5);
102        let mut map = HashMap::new();
103        let mq = create_custom_mq("order_topic", 5);
104        let mut offset = TopicOffset::new();
105        offset.set_min_offset(100);
106        offset.set_max_offset(200);
107        map.insert(mq.clone(), offset);
108        table.set_offset_table(map);
109        table
110            .get_offset_table_mut()
111            .get_mut(&mq)
112            .expect("queue offset")
113            .set_last_update_timestamp(11_111_111);
114
115        let serialized = serde_json::to_string(&table).expect("Serialization failed");
116        assert!(serialized.contains("offsetTable"));
117
118        let deserialized: TopicStatsTable = serde_json::from_str(&serialized).expect("Deserialization failed");
119        assert_eq!(deserialized.get_topic_put_tps(), 2.5);
120        let offsets = deserialized.into_offset_table();
121        let offset = offsets.get(&mq).expect("queue offset should round-trip");
122        assert_eq!(offset.get_min_offset(), 100);
123        assert_eq!(offset.get_max_offset(), 200);
124        assert_eq!(offset.get_last_update_timestamp(), 11_111_111);
125    }
126
127    #[test]
128    fn test_decode_java_fastjson_message_queue_keys() {
129        let body = br#"{"offsetTable":{{"topic":"TBW102","brokerName":"broker-a","queueId":0}:{"minOffset":0,"maxOffset":0,"lastUpdateTimestamp":0}},"topicPutTps":1.5}"#;
130
131        let table = TopicStatsTable::decode(body).expect("decode Java fastjson topic stats");
132
133        let queue = MessageQueue::from_parts("TBW102", "broker-a", 0);
134        let offset = table
135            .get_offset_table()
136            .get(&queue)
137            .expect("queue offset should decode");
138        assert_eq!(offset.get_min_offset(), 0);
139        assert_eq!(offset.get_max_offset(), 0);
140        assert_eq!(table.get_topic_put_tps(), 1.5);
141    }
142}
rocketmq_protocol/protocol/admin/topic_stats_table.rs

rocketmq_protocol/protocol/admin/
topic_stats_table.rs
