rocketmq_protocol/protocol/admin/
consume_stats.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::de::Error as DeError;
19use serde::de::MapAccess;
20use serde::de::Visitor;
21use serde::Deserialize;
22use serde::Deserializer;
23use serde::Serialize;
24use serde_json_any_key::*;
25
26use crate::protocol::admin::offset_wrapper::OffsetWrapper;
27use crate::protocol::RemotingDeserializable;
28
29#[derive(Debug, Default, Serialize, Deserialize)]
30#[serde(rename_all = "camelCase")]
31pub struct ConsumeStats {
32    #[serde(with = "any_key_map")]
33    pub offset_table: HashMap<MessageQueue, OffsetWrapper>,
34    pub consume_tps: f64,
35}
36
37#[derive(Deserialize)]
38#[serde(rename_all = "camelCase")]
39struct StrictConsumeStats {
40    #[serde(deserialize_with = "deserialize_strict_offset_table")]
41    offset_table: HashMap<MessageQueue, OffsetWrapper>,
42    consume_tps: f64,
43}
44
45fn deserialize_strict_offset_table<'de, D>(deserializer: D) -> Result<HashMap<MessageQueue, OffsetWrapper>, D::Error>
46where
47    D: Deserializer<'de>,
48{
49    struct StrictOffsetTableVisitor;
50
51    impl<'de> Visitor<'de> for StrictOffsetTableVisitor {
52        type Value = HashMap<MessageQueue, OffsetWrapper>;
53
54        fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55            formatter.write_str("an offset table with unique MessageQueue keys")
56        }
57
58        fn visit_map<A>(self, mut access: A) -> Result<Self::Value, A::Error>
59        where
60            A: MapAccess<'de>,
61        {
62            let mut offsets = HashMap::with_capacity(access.size_hint().unwrap_or_default());
63            while let Some((encoded_queue, offset)) = access.next_entry::<String, OffsetWrapper>()? {
64                let queue = serde_json::from_str::<MessageQueue>(&encoded_queue).map_err(A::Error::custom)?;
65                if offsets.insert(queue, offset).is_some() {
66                    return Err(A::Error::custom("duplicate MessageQueue key in offset table"));
67                }
68            }
69            Ok(offsets)
70        }
71    }
72
73    deserializer.deserialize_map(StrictOffsetTableVisitor)
74}
75
76impl ConsumeStats {
77    pub fn new() -> Self {
78        Self {
79            offset_table: HashMap::new(),
80            consume_tps: 0.0,
81        }
82    }
83
84    pub fn compute_total_diff(&self) -> i64 {
85        self.offset_table
86            .values()
87            .map(|value| value.get_broker_offset() - value.get_consumer_offset())
88            .sum()
89    }
90
91    pub fn compute_inflight_total_diff(&self) -> i64 {
92        self.offset_table
93            .values()
94            .map(|value| value.get_pull_offset() - value.get_consumer_offset())
95            .sum()
96    }
97
98    pub fn get_offset_table(&self) -> &HashMap<MessageQueue, OffsetWrapper> {
99        &self.offset_table
100    }
101
102    pub fn get_offset_table_mut(&mut self) -> &mut HashMap<MessageQueue, OffsetWrapper> {
103        &mut self.offset_table
104    }
105
106    pub fn set_offset_table(&mut self, offset_table: HashMap<MessageQueue, OffsetWrapper>) {
107        self.offset_table = offset_table;
108    }
109
110    pub fn get_consume_tps(&self) -> f64 {
111        self.consume_tps
112    }
113
114    pub fn set_consume_tps(&mut self, consume_tps: f64) {
115        self.consume_tps = consume_tps;
116    }
117
118    pub fn decode(body: &[u8]) -> rocketmq_error::Result<Self> {
119        match <Self as RemotingDeserializable>::decode(body) {
120            Ok(stats) => Ok(stats),
121            Err(error) => {
122                let Ok(raw_body) = std::str::from_utf8(body) else {
123                    return Err(error);
124                };
125                let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
126                if normalized_body == raw_body {
127                    return Err(error);
128                }
129                <Self as RemotingDeserializable>::decode_str(&normalized_body)
130            }
131        }
132    }
133
134    /// Decodes a consume-stats body while rejecting duplicate logical queue keys.
135    pub fn decode_strict(body: &[u8]) -> rocketmq_error::Result<Self> {
136        let raw_body = std::str::from_utf8(body)
137            .map_err(|source| crate::error::serialization_source("deserialize", "JSON", source))?;
138        let normalized_body = normalize_nonstandard_offset_table_keys(raw_body);
139        let strict = serde_json::from_str::<StrictConsumeStats>(&normalized_body)
140            .map_err(|error| crate::error::serialization_source("deserialize", "JSON", error))?;
141        Ok(Self {
142            offset_table: strict.offset_table,
143            consume_tps: strict.consume_tps,
144        })
145    }
146
147    pub fn encode_java_compatible(&self) -> rocketmq_error::Result<Vec<u8>> {
148        Ok(self.to_java_compatible_json()?.into_bytes())
149    }
150
151    pub fn to_java_compatible_json(&self) -> rocketmq_error::Result<String> {
152        let mut body = String::new();
153        body.push_str("{\"offsetTable\":{");
154
155        for (index, (queue, offset)) in self.offset_table.iter().enumerate() {
156            if index > 0 {
157                body.push(',');
158            }
159            append_message_queue_object_key(&mut body, queue)?;
160            body.push(':');
161            body.push_str(
162                &serde_json::to_string(offset)
163                    .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
164            );
165        }
166
167        body.push_str("},\"consumeTps\":");
168        body.push_str(
169            &serde_json::to_string(&self.consume_tps)
170                .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
171        );
172        body.push('}');
173        Ok(body)
174    }
175}
176
177pub(crate) fn normalize_nonstandard_offset_table_keys(input: &str) -> String {
178    let chars = input.chars().collect::<Vec<_>>();
179    let mut output = String::with_capacity(input.len());
180    let mut index = 0;
181
182    while index < chars.len() {
183        if chars[index] != '"' {
184            output.push(chars[index]);
185            index += 1;
186            continue;
187        }
188
189        let Some(string_end) = consume_string_end(&chars, index) else {
190            return input.to_owned();
191        };
192        output.extend(chars[index..=string_end].iter());
193        let key = chars[index + 1..string_end].iter().collect::<String>();
194        index = string_end + 1;
195
196        if key != "offsetTable" {
197            continue;
198        }
199
200        while index < chars.len() && chars[index].is_whitespace() {
201            output.push(chars[index]);
202            index += 1;
203        }
204        if index >= chars.len() || chars[index] != ':' {
205            continue;
206        }
207        output.push(':');
208        index += 1;
209
210        while index < chars.len() && chars[index].is_whitespace() {
211            output.push(chars[index]);
212            index += 1;
213        }
214        if index >= chars.len() || chars[index] != '{' {
215            continue;
216        }
217
218        output.push('{');
219        index += 1;
220        let Some((normalized_map, next_index)) = normalize_object_key_map_body(&chars, index) else {
221            return input.to_owned();
222        };
223        output.push_str(&normalized_map);
224        index = next_index;
225    }
226    output
227}
228
229pub(crate) fn append_message_queue_object_key(output: &mut String, queue: &MessageQueue) -> rocketmq_error::Result<()> {
230    output.push('{');
231    output.push_str("\"topic\":");
232    output.push_str(
233        &serde_json::to_string(queue.topic_str())
234            .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
235    );
236    output.push_str(",\"brokerName\":");
237    output.push_str(
238        &serde_json::to_string(queue.broker_name().as_str())
239            .map_err(|error| crate::error::serialization_source("serialize", "JSON", error))?,
240    );
241    output.push_str(",\"queueId\":");
242    output.push_str(&queue.queue_id().to_string());
243    output.push('}');
244    Ok(())
245}
246
247fn normalize_object_key_map_body(chars: &[char], mut index: usize) -> Option<(String, usize)> {
248    let mut output = String::new();
249    let mut expecting_key = true;
250    let mut nested_value_depth = 0usize;
251    let length = chars.len();
252
253    while index < length {
254        let current = chars[index];
255        if expecting_key {
256            match current {
257                '}' => {
258                    output.push('}');
259                    return Some((output, index + 1));
260                }
261                '{' => {
262                    let object_end = consume_balanced_object_end(chars, index)?;
263                    let raw_key = chars[index..=object_end].iter().collect::<String>();
264                    append_json_string(&mut output, &raw_key);
265                    index = object_end + 1;
266                    expecting_key = false;
267                }
268                '"' => {
269                    let string_end = consume_string_end(chars, index)?;
270                    output.extend(chars[index..=string_end].iter());
271                    index = string_end + 1;
272                    expecting_key = false;
273                }
274                _ => {
275                    output.push(current);
276                    index += 1;
277                }
278            }
279            continue;
280        }
281
282        match current {
283            '"' => {
284                let string_end = consume_string_end(chars, index)?;
285                output.extend(chars[index..=string_end].iter());
286                index = string_end + 1;
287            }
288            '{' | '[' => {
289                nested_value_depth += 1;
290                output.push(current);
291                index += 1;
292            }
293            '}' => {
294                if nested_value_depth == 0 {
295                    output.push('}');
296                    return Some((output, index + 1));
297                }
298                nested_value_depth -= 1;
299                output.push('}');
300                index += 1;
301            }
302            ']' => {
303                nested_value_depth = nested_value_depth.saturating_sub(1);
304                output.push(']');
305                index += 1;
306            }
307            ',' if nested_value_depth == 0 => {
308                output.push(',');
309                index += 1;
310                expecting_key = true;
311            }
312            _ => {
313                output.push(current);
314                index += 1;
315            }
316        }
317    }
318
319    Some((output, index))
320}
321
322fn append_json_string(output: &mut String, value: &str) {
323    output.push('"');
324    for ch in value.chars() {
325        match ch {
326            '"' => output.push_str("\\\""),
327            '\\' => output.push_str("\\\\"),
328            '\u{08}' => output.push_str("\\b"),
329            '\u{0C}' => output.push_str("\\f"),
330            '\n' => output.push_str("\\n"),
331            '\r' => output.push_str("\\r"),
332            '\t' => output.push_str("\\t"),
333            ch if ch <= '\u{1F}' => {
334                output.push_str("\\u00");
335                const HEX: &[u8; 16] = b"0123456789abcdef";
336                let byte = ch as u8;
337                output.push(HEX[(byte >> 4) as usize] as char);
338                output.push(HEX[(byte & 0x0f) as usize] as char);
339            }
340            ch => output.push(ch),
341        }
342    }
343    output.push('"');
344}
345
346fn consume_string_end(chars: &[char], start: usize) -> Option<usize> {
347    let mut index = start + 1;
348    let mut escaped = false;
349    while index < chars.len() {
350        if escaped {
351            escaped = false;
352        } else if chars[index] == '\\' {
353            escaped = true;
354        } else if chars[index] == '"' {
355            return Some(index);
356        }
357        index += 1;
358    }
359    None
360}
361
362fn consume_balanced_object_end(chars: &[char], start: usize) -> Option<usize> {
363    let mut depth = 0usize;
364    let mut index = start;
365    let mut in_string = false;
366    let mut escaped = false;
367
368    while index < chars.len() {
369        let current = chars[index];
370        if in_string {
371            if escaped {
372                escaped = false;
373            } else if current == '\\' {
374                escaped = true;
375            } else if current == '"' {
376                in_string = false;
377            }
378        } else {
379            match current {
380                '"' => in_string = true,
381                '{' => depth += 1,
382                '}' => {
383                    depth -= 1;
384                    if depth == 0 {
385                        return Some(index);
386                    }
387                }
388                _ => {}
389            }
390        }
391        index += 1;
392    }
393
394    None
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400
401    #[test]
402    fn decode_rejects_unterminated_strings_without_panicking() {
403        for body in [
404            "\"",
405            "{\"offsetTable",
406            "{\"offsetTable\":{{\"topic\":\"unfinished",
407            "{\"offsetTable\":{{\"topic\":\"topic\"",
408            "{\"offsetTable\":{\"",
409            "{\"offsetTable\":{\"key\":\"value",
410            "{\"offsetTable\":{\"key\":\"value\\",
411        ] {
412            assert!(ConsumeStats::decode(body.as_bytes()).is_err(), "{body:?}");
413            assert!(ConsumeStats::decode_strict(body.as_bytes()).is_err(), "{body:?}");
414            assert_eq!(normalize_nonstandard_offset_table_keys(body), body);
415        }
416    }
417
418    #[test]
419    fn decode_preserves_standard_json_and_rejects_unrelated_malformed_input() {
420        let decoded = ConsumeStats::decode(br#"{"offsetTable":{},"consumeTps":1.5}"#).unwrap();
421        assert!(decoded.get_offset_table().is_empty());
422        assert_eq!(decoded.get_consume_tps(), 1.5);
423        for body in [r#"{"offsetTable":{},"consumeTps":}"#, r#"{"offsetTable":{},}"#] {
424            assert_eq!(normalize_nonstandard_offset_table_keys(body), body);
425            assert!(ConsumeStats::decode(body.as_bytes()).is_err());
426        }
427    }
428
429    fn create_offset_wrapper(broker: i64, consumer: i64, pull: i64) -> OffsetWrapper {
430        let mut wrapper = OffsetWrapper::new();
431        wrapper.set_broker_offset(broker);
432        wrapper.set_consumer_offset(consumer);
433        wrapper.set_pull_offset(pull);
434        wrapper
435    }
436
437    fn create_mq(topic: &str, queue_id: i32) -> MessageQueue {
438        let json = format!(
439            r#"{{"topic": "{}", "brokerName": "broker_1", "queueId": {}}}"#,
440            topic, queue_id
441        );
442        serde_json::from_str(&json).expect("Failed to create MessageQueue for test")
443    }
444
445    #[test]
446    fn methods_compute_consumer_lag() {
447        let mut stats = ConsumeStats::new();
448        assert_eq!(stats.get_consume_tps(), 0.0);
449        assert!(stats.get_offset_table().is_empty());
450        assert_eq!(stats.compute_total_diff(), 0);
451        assert_eq!(stats.compute_inflight_total_diff(), 0);
452
453        stats.set_consume_tps(5.67);
454        let mut map = HashMap::new();
455        map.insert(create_mq("topic_a", 0), create_offset_wrapper(500, 400, 450));
456        stats.set_offset_table(map);
457        stats
458            .get_offset_table_mut()
459            .insert(create_mq("topic_a", 1), create_offset_wrapper(200, 150, 180));
460
461        assert_eq!(stats.get_consume_tps(), 5.67);
462        assert_eq!(stats.get_offset_table().len(), 2);
463        assert_eq!(stats.compute_total_diff(), 150);
464        assert_eq!(stats.compute_inflight_total_diff(), 80);
465    }
466
467    #[test]
468    fn test_serialization_json_any_key() {
469        let mut stats = ConsumeStats::new();
470        let mq = create_mq("order", 2);
471        stats.get_offset_table_mut().insert(mq, create_offset_wrapper(10, 5, 8));
472        stats.set_consume_tps(1.0);
473
474        let serialized = serde_json::to_string(&stats).expect("Serialization failed");
475
476        assert!(serialized.contains("offsetTable"));
477        assert!(serialized.contains("consumeTps"));
478
479        let deserialized: ConsumeStats = serde_json::from_str(&serialized).expect("Deserialization failed");
480        assert_eq!(deserialized.get_offset_table().len(), 1);
481        assert_eq!(deserialized.compute_total_diff(), 5);
482    }
483
484    #[test]
485    fn decode_accepts_nonstandard_offset_table_object_keys() {
486        let body = r#"{"offsetTable":{{"topic":"TopicTest","brokerName":"broker-a","queueId":1}:{"brokerOffset":120,"consumerOffset":20,"pullOffset":80,"lastTimestamp":1700000000000}},"consumeTps":1.5}"#;
487
488        let decoded = ConsumeStats::decode(body.as_bytes()).expect("consume stats should decode");
489
490        assert_eq!(decoded.get_offset_table().len(), 1);
491        assert_eq!(decoded.get_consume_tps(), 1.5);
492        let (queue, offset) = decoded.get_offset_table().iter().next().expect("one queue");
493        assert_eq!(queue.topic_str(), "TopicTest");
494        assert_eq!(queue.broker_name(), "broker-a");
495        assert_eq!(queue.queue_id(), 1);
496        assert_eq!(offset.get_broker_offset(), 120);
497        assert_eq!(offset.get_consumer_offset(), 20);
498        assert_eq!(offset.get_pull_offset(), 80);
499        assert_eq!(offset.get_last_timestamp(), 1_700_000_000_000);
500    }
501
502    #[test]
503    fn strict_decode_rejects_duplicate_logical_queue_keys() {
504        let queue = serde_json::to_string(&create_mq("TopicTest", 1)).expect("queue key");
505        let encoded_key = serde_json::to_string(&queue).expect("encoded key");
506        let offset = r#"{"brokerOffset":120,"consumerOffset":20,"pullOffset":80,"lastTimestamp":0}"#;
507        let body = format!(r#"{{"offsetTable":{{{encoded_key}:{offset},{encoded_key}:{offset}}},"consumeTps":1.5}}"#);
508
509        assert!(ConsumeStats::decode_strict(body.as_bytes()).is_err());
510    }
511
512    #[test]
513    fn normalize_nonstandard_offset_table_keys_stringifies_object_keys() {
514        let input = r#"{"offsetTable":{{"topic":"TopicTest","brokerName":"broker-a","queueId":1}:{"brokerOffset":120}},"consumeTps":1.5}"#;
515        let normalized = normalize_nonstandard_offset_table_keys(input);
516
517        assert_eq!(
518            normalized,
519            r#"{"offsetTable":{"{\"topic\":\"TopicTest\",\"brokerName\":\"broker-a\",\"queueId\":1}":{"brokerOffset":120}},"consumeTps":1.5}"#
520        );
521    }
522
523    #[test]
524    fn normalize_nonstandard_offset_table_keys_escapes_object_key_strings() {
525        let input = r#"{"offsetTable":{{"topic":"Topic\"Quoted","brokerName":"broker\\one","queueId":1}:{"brokerOffset":120}},"consumeTps":1.5}"#;
526        let normalized = normalize_nonstandard_offset_table_keys(input);
527        let value: serde_json::Value = serde_json::from_str(&normalized).expect("normalized JSON should parse");
528        let offset_table = value["offsetTable"].as_object().expect("offsetTable object");
529        let key = offset_table.keys().next().expect("one queue key");
530
531        assert_eq!(
532            key,
533            r#"{"topic":"Topic\"Quoted","brokerName":"broker\\one","queueId":1}"#
534        );
535    }
536
537    #[test]
538    fn java_compatible_encode_uses_object_keys_for_offset_table() {
539        let mut stats = ConsumeStats::new();
540        stats
541            .get_offset_table_mut()
542            .insert(create_mq("TopicTest", 1), create_offset_wrapper(120, 20, 80));
543        stats.set_consume_tps(1.5);
544
545        let encoded = String::from_utf8(stats.encode_java_compatible().expect("encode consume stats"))
546            .expect("utf8 consume stats");
547
548        assert!(encoded.contains(r#""offsetTable":{{"topic":"TopicTest","brokerName":"broker_1","queueId":1}:"#));
549        assert!(!encoded.contains(r#""{\"topic\""#));
550
551        let decoded = ConsumeStats::decode(encoded.as_bytes()).expect("decode java-compatible consume stats");
552        assert_eq!(decoded.get_offset_table().len(), 1);
553        assert_eq!(decoded.compute_total_diff(), 100);
554    }
555}
rocketmq_protocol/protocol/admin/consume_stats.rs

rocketmq_protocol/protocol/admin/
consume_stats.rs
