Skip to main content

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_common::common::message::message_queue::MessageQueue;
20use rocketmq_remoting::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}