rocketmq_model/common/message/
message_client_id_setter.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::net::IpAddr;
16use std::sync::atomic::AtomicI32;
17use std::sync::atomic::AtomicI64;
18use std::sync::atomic::Ordering;
19use std::sync::LazyLock;
20
21use bytes::BufMut;
22use bytes::BytesMut;
23use cheetah_string::CheetahString;
24use chrono::DateTime;
25use chrono::Datelike;
26use chrono::Local;
27use chrono::LocalResult;
28use chrono::Months;
29use chrono::TimeZone;
30use rocketmq_error::Result;
31
32use crate::codec::bytes_to_string;
33use crate::codec::string_to_bytes;
34use crate::codec::write_int;
35use crate::codec::write_short;
36use crate::common::hasher::string_hasher::JavaStringHasher;
37use crate::common::message::MessageConst;
38use crate::common::message::MessageTrait;
39use crate::time::current_millis;
40
41static COUNTER: LazyLock<AtomicI32> = LazyLock::new(|| AtomicI32::new(0));
42static TIME_WINDOW: LazyLock<TimeWindow> = LazyLock::new(|| {
43    let (start_time, next_start_time) = compute_time_window(current_millis() as i64);
44    TimeWindow {
45        start_time: AtomicI64::new(start_time),
46        next_start_time: AtomicI64::new(next_start_time),
47    }
48});
49static LEN: LazyLock<usize> = LazyLock::new(|| {
50    let ip = get_ip().unwrap_or_else(create_fake_ip);
51    ip.len() + 2 + 4 + 4 + 2
52});
53static FIX_STRING: LazyLock<Vec<char>> = LazyLock::new(|| {
54    let ip = get_ip().unwrap_or_else(create_fake_ip);
55    let pid = std::process::id() as i16;
56    let class_loader_hash = JavaStringHasher::hash_str("MessageClientIDSetter");
57    let mut bytes = BytesMut::with_capacity(ip.len() + 2 + 4);
58    bytes.put(ip.as_slice());
59    bytes.put_i16(pid);
60    bytes.put_i32(class_loader_hash);
61    let data = bytes_to_string(bytes.freeze().as_ref()).chars().collect::<Vec<char>>();
62    data
63});
64
65struct TimeWindow {
66    start_time: AtomicI64,
67    next_start_time: AtomicI64,
68}
69
70pub fn create_fake_ip() -> Vec<u8> {
71    current_millis().to_be_bytes()[4..].to_vec()
72}
73
74fn get_ip() -> Option<Vec<u8>> {
75    match local_ip_address::local_ip().ok()? {
76        IpAddr::V4(address) => Some(address.octets().to_vec()),
77        IpAddr::V6(address) => Some(address.octets().to_vec()),
78    }
79}
80
81fn local_timestamp_millis_or_epoch(millis: i64) -> DateTime<Local> {
82    Local
83        .timestamp_millis_opt(millis)
84        .single()
85        .or_else(|| Local.timestamp_millis_opt(0).single())
86        .unwrap_or_else(Local::now)
87}
88
89fn local_month_start(dt: DateTime<Local>) -> DateTime<Local> {
90    match Local.with_ymd_and_hms(dt.year(), dt.month(), 1, 0, 0, 0) {
91        LocalResult::Single(value) => value,
92        LocalResult::Ambiguous(earliest, _) => earliest,
93        LocalResult::None => dt,
94    }
95}
96
97fn compute_time_window(millis: i64) -> (i64, i64) {
98    let start = local_month_start(local_timestamp_millis_or_epoch(millis));
99    let next = start.checked_add_months(Months::new(1)).unwrap_or(start);
100    (start.timestamp_millis(), next.timestamp_millis())
101}
102
103fn decode_uniq_id_bytes(msg_id: &str) -> Result<Vec<u8>> {
104    let msg_id = msg_id.trim();
105    if msg_id.len() != 32 && msg_id.len() != 56 {
106        return Err(crate::error::invalid_argument(format!(
107            "Invalid uniq id length: {}. Expected 32 characters (IPv4) or 56 characters (IPv6)",
108            msg_id.len()
109        )));
110    }
111    if !msg_id.bytes().all(|b| b.is_ascii_hexdigit()) {
112        return Err(crate::error::invalid_argument(
113            "Invalid uniq id: expected hexadecimal characters",
114        ));
115    }
116    string_to_bytes(msg_id).ok_or_else(|| crate::error::invalid_argument("Invalid uniq id hex"))
117}
118
119pub struct MessageClientIDSetter;
120
121impl MessageClientIDSetter {
122    #[inline]
123    pub fn get_uniq_id<T>(message: &T) -> Option<CheetahString>
124    where
125        T: MessageTrait,
126    {
127        message.property(&CheetahString::from_static_str(
128            MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX,
129        ))
130    }
131
132    fn set_start_time(millis: i64) {
133        let (start_time, next_start_time) = compute_time_window(millis);
134        TIME_WINDOW.start_time.store(start_time, Ordering::Release);
135        TIME_WINDOW.next_start_time.store(next_start_time, Ordering::Release);
136    }
137
138    pub fn create_uniq_id() -> String {
139        let mut sb = vec!['\0'; *LEN * 2];
140        sb[..FIX_STRING.len()].copy_from_slice(&FIX_STRING);
141
142        let current = current_millis() as i64;
143        if current >= TIME_WINDOW.next_start_time.load(Ordering::Acquire) {
144            Self::set_start_time(current);
145        }
146
147        let mut diff = (current - TIME_WINDOW.start_time.load(Ordering::Acquire)) as i32;
148        if diff < 0 && diff > -1_000_000 {
149            // May cause by NTP
150            diff = 0;
151        }
152
153        let pos = FIX_STRING.len();
154        write_int(&mut sb, pos, diff);
155
156        let counter_val = COUNTER.fetch_add(1, Ordering::SeqCst) as i16;
157        write_short(&mut sb, pos + 8, counter_val);
158        sb.into_iter().collect()
159    }
160
161    pub fn get_nearly_time_from_id(msg_id: &str) -> Result<DateTime<Local>> {
162        let bytes = decode_uniq_id_bytes(msg_id)?;
163        let ip_length = if bytes.len() == 28 { 16 } else { 4 };
164        let diff_start = ip_length + 2 + 4;
165        let span_ms = u32::from_be_bytes([
166            bytes[diff_start],
167            bytes[diff_start + 1],
168            bytes[diff_start + 2],
169            bytes[diff_start + 3],
170        ]) as i64;
171        let now = Local::now();
172        let mut month_start = local_month_start(now).timestamp_millis();
173        if month_start.saturating_add(span_ms) >= now.timestamp_millis() {
174            month_start = local_month_start(
175                now.checked_sub_months(Months::new(1))
176                    .unwrap_or_else(|| local_timestamp_millis_or_epoch(0)),
177            )
178            .timestamp_millis();
179        }
180        let timestamp = month_start.saturating_add(span_ms);
181        Local
182            .timestamp_millis_opt(timestamp)
183            .single()
184            .ok_or_else(|| crate::error::invalid_argument("Invalid uniq id timestamp"))
185    }
186
187    pub fn get_ip_from_id(msg_id: &str) -> Result<Vec<u8>> {
188        let bytes = decode_uniq_id_bytes(msg_id)?;
189        let ip_length = if bytes.len() == 28 { 16 } else { 4 };
190        Ok(bytes[..ip_length].to_vec())
191    }
192
193    pub fn get_ip_str_from_id(msg_id: &str) -> Result<String> {
194        let ip_bytes = Self::get_ip_from_id(msg_id)?;
195        match ip_bytes.len() {
196            4 => Ok(std::net::Ipv4Addr::new(ip_bytes[0], ip_bytes[1], ip_bytes[2], ip_bytes[3]).to_string()),
197            16 => {
198                let mut segments = [0u8; 16];
199                segments.copy_from_slice(&ip_bytes);
200                Ok(std::net::Ipv6Addr::from(segments).to_string())
201            }
202            len => Err(crate::error::invalid_argument(format!(
203                "Invalid uniq id IP length: {len}"
204            ))),
205        }
206    }
207
208    pub fn get_pid_from_id(msg_id: &str) -> Result<i32> {
209        let bytes = decode_uniq_id_bytes(msg_id)?;
210        let pid_start = bytes.len() - 2 - 4 - 4 - 2;
211        let pid = u16::from_be_bytes([bytes[pid_start], bytes[pid_start + 1]]);
212        Ok(i32::from(pid))
213    }
214
215    pub fn set_uniq_id<T>(message: &mut T)
216    where
217        T: MessageTrait,
218    {
219        if message
220            .property(&CheetahString::from_static_str(
221                MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX,
222            ))
223            .is_none()
224        {
225            let uniq_id = Self::create_uniq_id();
226            message.put_property(
227                CheetahString::from_static_str(MessageConst::PROPERTY_UNIQ_CLIENT_MESSAGE_ID_KEYIDX),
228                CheetahString::from_string(uniq_id),
229            );
230        }
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use chrono::Local;
237
238    use super::*;
239    use crate::common::message::message_single::Message;
240
241    #[test]
242    fn unique_id_is_generated_correctly() {
243        let _message = Message::default();
244        let unique_id = MessageClientIDSetter::create_uniq_id();
245        assert!(!unique_id.is_empty());
246    }
247
248    #[test]
249    fn unique_id_counter_increments_properly() {
250        let first_id = MessageClientIDSetter::create_uniq_id();
251        let second_id = MessageClientIDSetter::create_uniq_id();
252        let first_counter = &first_id[FIX_STRING.len() + 8..];
253        let second_counter = &second_id[FIX_STRING.len() + 8..];
254        assert_ne!(first_counter, second_counter);
255    }
256
257    #[test]
258    fn get_uniq_id_returns_none_when_not_set() {
259        let message = Message::default();
260        assert_eq!(MessageClientIDSetter::get_uniq_id(&message), None);
261    }
262
263    #[test]
264    fn create_fake_ip_generates_valid_ip() {
265        let fake_ip = create_fake_ip();
266        assert_eq!(fake_ip.len(), 4);
267        assert!(fake_ip.iter().all(|&_byte| true));
268    }
269
270    #[test]
271    fn start_time_is_set_correctly_for_new_month() {
272        let millis = Local::now().timestamp_millis();
273        MessageClientIDSetter::set_start_time(millis);
274        let locked_start_time = TIME_WINDOW.start_time.load(Ordering::Acquire);
275        assert!(locked_start_time <= millis);
276    }
277
278    #[test]
279    fn next_start_time_is_set_correctly_for_new_month() {
280        let millis = Local::now().timestamp_millis();
281        MessageClientIDSetter::set_start_time(millis);
282        let locked_next_start_time = TIME_WINDOW.next_start_time.load(Ordering::Acquire);
283        assert!(locked_next_start_time > millis);
284    }
285
286    #[test]
287    fn nearly_time_from_id_is_close_to_now() {
288        let before = Local::now().timestamp_millis();
289        let unique_id = MessageClientIDSetter::create_uniq_id();
290        std::thread::sleep(std::time::Duration::from_millis(2));
291        let nearly_time = MessageClientIDSetter::get_nearly_time_from_id(&unique_id)
292            .expect("generated uniq id should decode time")
293            .timestamp_millis();
294        let after = Local::now().timestamp_millis();
295
296        assert!(nearly_time >= before - 1_000);
297        assert!(nearly_time <= after + 1_000);
298    }
299
300    #[test]
301    fn ip_and_pid_decode_from_generated_id() {
302        let unique_id = MessageClientIDSetter::create_uniq_id();
303
304        let ip = MessageClientIDSetter::get_ip_from_id(&unique_id).expect("generated uniq id should decode IP");
305        assert!(ip.len() == 4 || ip.len() == 16);
306        let ip_str =
307            MessageClientIDSetter::get_ip_str_from_id(&unique_id).expect("generated uniq id should decode IP string");
308        assert!(!ip_str.is_empty());
309        let pid = MessageClientIDSetter::get_pid_from_id(&unique_id).expect("generated uniq id should decode PID");
310        assert_eq!(pid, i32::from(std::process::id() as u16));
311    }
312
313    #[test]
314    fn malformed_id_decoders_return_error() {
315        assert!(MessageClientIDSetter::get_nearly_time_from_id("bad-id").is_err());
316        assert!(MessageClientIDSetter::get_ip_from_id("bad-id").is_err());
317        assert!(MessageClientIDSetter::get_ip_str_from_id("bad-id").is_err());
318        assert!(MessageClientIDSetter::get_pid_from_id("bad-id").is_err());
319    }
320}
rocketmq_model/common/message/message_client_id_setter.rs

rocketmq_model/common/message/
message_client_id_setter.rs
