Skip to main content

rocketmq_protocol/protocol/body/response/
reset_offset_body.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::common::message::message_queue_for_c::MessageQueueForC;
18use rocketmq_model::message::MessageQueue;
19use serde::Deserialize;
20use serde::Serialize;
21use serde_json_any_key::*;
22
23use crate::protocol::admin::consume_stats::append_message_queue_object_key;
24use crate::protocol::admin::consume_stats::normalize_nonstandard_offset_table_keys;
25use crate::protocol::RemotingDeserializable;
26
27/// Response body for reset consumer offset operation.
28#[derive(Serialize, Deserialize, Debug, Default, Clone)]
29#[serde(rename_all = "camelCase")]
30pub struct ResetOffsetBody {
31    #[serde(with = "any_key_map")]
32    pub offset_table: HashMap<MessageQueue, i64>,
33}
34
35impl ResetOffsetBody {
36    pub fn new() -> Self {
37        Self {
38            offset_table: HashMap::new(),
39        }
40    }
41
42    pub fn encode(&self) -> Vec<u8> {
43        self.encode_java_compatible().unwrap_or_default()
44    }
45
46    pub fn encode_java_compatible(&self) -> rocketmq_error::Result<Vec<u8>> {
47        Ok(self.to_java_compatible_json()?.into_bytes())
48    }
49
50    pub fn to_java_compatible_json(&self) -> rocketmq_error::Result<String> {
51        let mut body = String::new();
52        body.push_str("{\"offsetTable\":{");
53
54        for (index, (queue, offset)) in self.offset_table.iter().enumerate() {
55            if index > 0 {
56                body.push(',');
57            }
58            append_message_queue_object_key(&mut body, queue)?;
59            body.push(':');
60            body.push_str(&offset.to_string());
61        }
62
63        body.push_str("}}");
64        Ok(body)
65    }
66
67    pub fn decode(body: &[u8]) -> Option<Self> {
68        match <Self as RemotingDeserializable>::decode(body) {
69            Ok(decoded) => Some(decoded),
70            Err(_) => {
71                let raw = std::str::from_utf8(body).ok()?;
72                let normalized = normalize_nonstandard_offset_table_keys(raw);
73                if normalized == raw {
74                    return None;
75                }
76                <Self as RemotingDeserializable>::decode_str(&normalized).ok()
77            }
78        }
79    }
80}
81
82/// Response body for reset consumer offset operation for C++ clients.
83/// Uses a list format with offset embedded in MessageQueueForC.
84#[derive(Serialize, Deserialize, Debug, Default, Clone)]
85#[serde(rename_all = "camelCase")]
86pub struct ResetOffsetBodyForC {
87    pub offset_table: Vec<MessageQueueForC>,
88}
89
90impl ResetOffsetBodyForC {
91    pub fn new() -> Self {
92        Self {
93            offset_table: Vec::new(),
94        }
95    }
96
97    /// Convert from standard offset table to C++ client format.
98    pub fn from_offset_table(offset_table: &HashMap<MessageQueue, i64>) -> Self {
99        let offset_list: Vec<MessageQueueForC> = offset_table
100            .iter()
101            .map(|(mq, offset)| MessageQueueForC::from_message_queue(mq, *offset))
102            .collect();
103        Self {
104            offset_table: offset_list,
105        }
106    }
107
108    pub fn encode(&self) -> Vec<u8> {
109        serde_json::to_vec(self).unwrap_or_default()
110    }
111
112    pub fn decode(body: &[u8]) -> Option<Self> {
113        serde_json::from_slice(body).ok()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn reset_offset_body_encode_uses_java_object_keys() {
123        let mut body = ResetOffsetBody::new();
124        body.offset_table
125            .insert(MessageQueue::from_parts("test_topic", "broker-a", 1), 100);
126
127        let encoded = String::from_utf8(body.encode()).expect("reset offset body should be utf8");
128
129        assert!(encoded.contains(r#""offsetTable":{{"topic":"test_topic","brokerName":"broker-a","queueId":1}:100}"#));
130        assert!(!encoded.contains(r#""{\"topic\""#));
131    }
132
133    #[test]
134    fn reset_offset_body_decode_accepts_java_object_keys() {
135        let body = r#"{"offsetTable":{{"topic":"test_topic","brokerName":"broker-a","queueId":1}:100}}"#;
136
137        let decoded = ResetOffsetBody::decode(body.as_bytes()).expect("decode reset offset body");
138
139        let queue = MessageQueue::from_parts("test_topic", "broker-a", 1);
140        assert_eq!(decoded.offset_table.get(&queue), Some(&100));
141    }
142
143    #[test]
144    fn test_reset_offset_body_for_c_from_offset_table() {
145        let mut offset_table = HashMap::new();
146        let mq1 = MessageQueue::from_parts("topic1", "broker-a", 0);
147        let mq2 = MessageQueue::from_parts("topic1", "broker-a", 1);
148        offset_table.insert(mq1, 100);
149        offset_table.insert(mq2, 200);
150
151        let body = ResetOffsetBodyForC::from_offset_table(&offset_table);
152        assert_eq!(body.offset_table.len(), 2);
153
154        // Verify offsets are included
155        for mq_for_c in &body.offset_table {
156            assert!(mq_for_c.offset == 100 || mq_for_c.offset == 200);
157        }
158    }
159
160    #[test]
161    fn test_reset_offset_body_for_c_encode_decode() {
162        let mut body = ResetOffsetBodyForC::new();
163        body.offset_table
164            .push(MessageQueueForC::new("test_topic", "broker-a", 1, 100));
165
166        let encoded = body.encode();
167        let decoded = ResetOffsetBodyForC::decode(&encoded).unwrap();
168        assert_eq!(decoded.offset_table.len(), 1);
169        assert_eq!(decoded.offset_table[0].offset, 100);
170    }
171}