rocketmq_client_rust/consumer/store/
offset_serialize_wrapper.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;
16use std::sync::atomic::AtomicI64;
17use std::sync::atomic::Ordering;
18
19use rocketmq_model::common::message::message_queue::MessageQueue;
20use rocketmq_protocol::protocol::RemotingDeserializable;
21use serde::Deserialize;
22use serde::Serialize;
23use tracing::warn;
24
25use crate::consumer::store::offset_serialize::OffsetSerialize;
26
27#[derive(Serialize, Deserialize, Debug, Default)]
28#[serde(rename_all = "camelCase")]
29pub struct OffsetSerializeWrapper {
30    #[serde(serialize_with = "serialize_atomic_i64", deserialize_with = "deserialize_atomic_i64")]
31    pub offset_table: HashMap<MessageQueue, AtomicI64>,
32}
33
34fn serialize_atomic_i64<S>(map: &HashMap<MessageQueue, AtomicI64>, serializer: S) -> Result<S::Ok, S::Error>
35where
36    S: serde::Serializer,
37{
38    let map_as_i64: HashMap<_, _> = map.iter().map(|(k, v)| (k, v.load(Ordering::Relaxed))).collect();
39    map_as_i64.serialize(serializer)
40}
41
42fn deserialize_atomic_i64<'de, D>(deserializer: D) -> Result<HashMap<MessageQueue, AtomicI64>, D::Error>
43where
44    D: serde::Deserializer<'de>,
45{
46    let map_as_i64: HashMap<MessageQueue, i64> = HashMap::deserialize(deserializer)?;
47    let map_as_atomic: HashMap<_, _> = map_as_i64.into_iter().map(|(k, v)| (k, AtomicI64::new(v))).collect();
48    Ok(map_as_atomic)
49}
50
51impl From<OffsetSerialize> for OffsetSerializeWrapper {
52    fn from(offset_serialize: OffsetSerialize) -> Self {
53        let offset_table = offset_serialize
54            .offset_table
55            .into_iter()
56            .filter_map(|(k, v)| match MessageQueue::decode(k.as_bytes()) {
57                Ok(message_queue) => Some((message_queue, AtomicI64::new(v))),
58                Err(err) => {
59                    warn!("skip offset entry because message queue decode failed: {err}");
60                    None
61                }
62            })
63            .collect();
64        Self { offset_table }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use rocketmq_protocol::protocol::RemotingSerializable;
72
73    fn queue_key(topic: &str, broker: &str, queue_id: i32) -> String {
74        MessageQueue::from_parts(topic, broker, queue_id)
75            .serialize_json()
76            .expect("queue should serialize to JSON")
77    }
78
79    #[test]
80    fn offset_table_round_trips_through_json_with_signed_offsets_and_keys() {
81        let queue_a = MessageQueue::from_parts("topic-a", "broker-a", 0);
82        let queue_b = MessageQueue::from_parts("topic-b", "broker-b", 1);
83
84        let mut wrapper = OffsetSerializeWrapper::default();
85        wrapper.offset_table.insert(queue_a.clone(), AtomicI64::new(-42));
86        wrapper.offset_table.insert(queue_b.clone(), AtomicI64::new(100));
87
88        let serialize: OffsetSerialize = wrapper.into();
89        let json = serialize.serialize_json().expect("offset serialize should encode");
90        let decoded = OffsetSerialize::decode_str(&json).expect("offset serialize should decode");
91        let round_tripped: OffsetSerializeWrapper = decoded.into();
92
93        assert_eq!(round_tripped.offset_table.len(), 2);
94        assert_eq!(
95            round_tripped
96                .offset_table
97                .get(&queue_a)
98                .map(|v| v.load(Ordering::Relaxed)),
99            Some(-42)
100        );
101        assert_eq!(
102            round_tripped
103                .offset_table
104                .get(&queue_b)
105                .map(|v| v.load(Ordering::Relaxed)),
106            Some(100)
107        );
108    }
109
110    #[test]
111    fn offset_larger_than_i32_max_is_not_narrowed() {
112        let queue = MessageQueue::from_parts("topic-a", "broker-a", 0);
113        let large_offset = i64::from(i32::MAX) + 1000;
114
115        let mut wrapper = OffsetSerializeWrapper::default();
116        wrapper.offset_table.insert(queue.clone(), AtomicI64::new(large_offset));
117
118        let serialize: OffsetSerialize = wrapper.into();
119        let json = serialize.serialize_json().expect("offset serialize should encode");
120        let decoded = OffsetSerialize::decode_str(&json).expect("offset serialize should decode");
121        let round_tripped: OffsetSerializeWrapper = decoded.into();
122
123        assert_eq!(
124            round_tripped
125                .offset_table
126                .get(&queue)
127                .map(|v| v.load(Ordering::Relaxed)),
128            Some(large_offset)
129        );
130    }
131
132    #[test]
133    fn invalid_queue_key_is_skipped_while_valid_entries_are_kept() {
134        let queue = MessageQueue::from_parts("topic-a", "broker-a", 0);
135        let mut offset_table = HashMap::new();
136        offset_table.insert(queue_key("topic-a", "broker-a", 0), 7);
137        offset_table.insert("not a valid message queue key".to_string(), 99);
138        let serialize = OffsetSerialize { offset_table };
139
140        let wrapper: OffsetSerializeWrapper = serialize.into();
141
142        assert_eq!(wrapper.offset_table.len(), 1);
143        assert_eq!(
144            wrapper.offset_table.get(&queue).map(|v| v.load(Ordering::Relaxed)),
145            Some(7)
146        );
147    }
148
149    #[test]
150    fn empty_offset_table_round_trips_with_persisted_field_name() {
151        let wrapper = OffsetSerializeWrapper::default();
152
153        let serialize: OffsetSerialize = wrapper.into();
154        let json = serialize.serialize_json().expect("offset serialize should encode");
155
156        assert!(json.contains("\"offsetTable\""));
157        let decoded = OffsetSerialize::decode_str(&json).expect("offset serialize should decode");
158        let round_tripped: OffsetSerializeWrapper = decoded.into();
159
160        assert!(round_tripped.offset_table.is_empty());
161    }
162}
rocketmq_client_rust/consumer/store/offset_serialize_wrapper.rs

rocketmq_client_rust/consumer/store/
offset_serialize_wrapper.rs
