rocketmq_protocol/protocol/admin/
consume_stats_list.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 cheetah_string::CheetahString;
18use serde::Deserialize;
19use serde::Serialize;
20
21use super::consume_stats::normalize_nonstandard_offset_table_keys;
22use super::consume_stats::ConsumeStats;
23use crate::protocol::RemotingDeserializable;
24
25#[derive(Serialize, Deserialize)]
26pub struct ConsumeStatsList {
27    #[serde(rename = "consumeStatsList")]
28    pub consume_stats_list: Vec<HashMap<CheetahString, Vec<ConsumeStats>>>,
29
30    #[serde(rename = "brokerAddr")]
31    pub broker_addr: Option<CheetahString>,
32
33    #[serde(rename = "totalDiff")]
34    pub total_diff: i64,
35
36    #[serde(rename = "totalInflightDiff")]
37    pub total_inflight_diff: i64,
38}
39
40impl ConsumeStatsList {
41    pub fn encode_java_compatible(&self) -> rocketmq_error::Result<Vec<u8>> {
42        Ok(self.to_java_compatible_json()?.into_bytes())
43    }
44
45    pub fn to_java_compatible_json(&self) -> rocketmq_error::Result<String> {
46        let mut body = String::new();
47        body.push_str("{\"consumeStatsList\":[");
48
49        for (list_index, group_map) in self.consume_stats_list.iter().enumerate() {
50            if list_index > 0 {
51                body.push(',');
52            }
53            body.push('{');
54            for (group_index, (group, stats_list)) in group_map.iter().enumerate() {
55                if group_index > 0 {
56                    body.push(',');
57                }
58                body.push_str(
59                    &serde_json::to_string(group)
60                        .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
61                );
62                body.push(':');
63                body.push('[');
64                for (stats_index, stats) in stats_list.iter().enumerate() {
65                    if stats_index > 0 {
66                        body.push(',');
67                    }
68                    body.push_str(&stats.to_java_compatible_json()?);
69                }
70                body.push(']');
71            }
72            body.push('}');
73        }
74
75        body.push_str("],\"brokerAddr\":");
76        match &self.broker_addr {
77            Some(broker_addr) => body.push_str(
78                &serde_json::to_string(broker_addr)
79                    .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
80            ),
81            None => body.push_str("null"),
82        }
83        body.push_str(",\"totalDiff\":");
84        body.push_str(&self.total_diff.to_string());
85        body.push_str(",\"totalInflightDiff\":");
86        body.push_str(&self.total_inflight_diff.to_string());
87        body.push('}');
88        Ok(body)
89    }
90
91    pub fn decode(body: &[u8]) -> rocketmq_error::Result<Self> {
92        match <Self as RemotingDeserializable>::decode(body) {
93            Ok(stats) => Ok(stats),
94            Err(error) => {
95                let Ok(raw_body) = std::str::from_utf8(body) else {
96                    return Err(error);
97                };
98                let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
99                if normalized_body == raw_body {
100                    return Err(error);
101                }
102                <Self as RemotingDeserializable>::decode_str(&normalized_body)
103            }
104        }
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use serde_json;
111
112    use super::*;
113
114    #[test]
115    fn consume_status_list_serialization() {
116        let mut map = HashMap::new();
117        let consume_stats_list = vec![ConsumeStats {
118            offset_table: HashMap::new(),
119            consume_tps: 1.2,
120        }];
121        map.insert(CheetahString::from("group1"), consume_stats_list);
122        let consume_status_list = ConsumeStatsList {
123            consume_stats_list: vec![map],
124            broker_addr: Some(CheetahString::from("addr")),
125            total_diff: 2,
126            total_inflight_diff: 1,
127        };
128        let serialized = serde_json::to_string(&consume_status_list).unwrap();
129        assert!(serialized.contains("\"consumeStatsList\""));
130        let deserialized: ConsumeStatsList = serde_json::from_str(&serialized).unwrap();
131        assert_eq!(deserialized.broker_addr.unwrap(), CheetahString::from("addr"));
132        assert_eq!(deserialized.total_diff, 2);
133        assert_eq!(deserialized.total_inflight_diff, 1);
134        let a = deserialized.consume_stats_list[0]
135            .get(&CheetahString::from("group1"))
136            .unwrap();
137        assert_eq!(a[0].consume_tps, 1.2);
138    }
139
140    #[test]
141    fn consume_status_list_java_compatible_serialization() {
142        use rocketmq_model::message::MessageQueue;
143
144        let mut stats = ConsumeStats::new();
145        stats
146            .get_offset_table_mut()
147            .insert(MessageQueue::from_parts("TopicTest", "broker-a", 0), Default::default());
148        let mut map = HashMap::new();
149        map.insert(CheetahString::from("group1"), vec![stats]);
150        let list = ConsumeStatsList {
151            consume_stats_list: vec![map],
152            broker_addr: Some(CheetahString::from("127.0.0.1:10911")),
153            total_diff: 0,
154            total_inflight_diff: 0,
155        };
156
157        let encoded = String::from_utf8(list.encode_java_compatible().expect("encode consume stats list"))
158            .expect("utf8 consume stats list");
159
160        assert!(encoded.contains(r#""consumeStatsList":[{"group1":[{"offsetTable":{{"topic":"TopicTest""#));
161        assert!(!encoded.contains(r#""{\"topic\""#));
162
163        let decoded = ConsumeStatsList::decode(encoded.as_bytes()).expect("decode consume stats list");
164        let group = decoded.consume_stats_list[0]
165            .get(&CheetahString::from("group1"))
166            .expect("group1 stats");
167        assert_eq!(group[0].offset_table.len(), 1);
168    }
169}
rocketmq_protocol/protocol/admin/consume_stats_list.rs

rocketmq_protocol/protocol/admin/
consume_stats_list.rs
