rocketmq_client_rust/common/
thread_local_index.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements.  See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License.  You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18#![allow(clippy::missing_const_for_thread_local)]
19use std::cell::RefCell;
20use std::fmt;
21
22use rand::Rng;
23
24thread_local! {
25    static THREAD_LOCAL_INDEX: RefCell<Option<i32>> = const {RefCell::new(None)};
26}
27
28const POSITIVE_MASK: i32 = 0x7FFFFFFF;
29const MAX: i32 = i32::MAX;
30
31#[derive(Default, Clone)]
32pub struct ThreadLocalIndex;
33
34impl ThreadLocalIndex {
35    pub fn increment_and_get(&self) -> i32 {
36        THREAD_LOCAL_INDEX.with(|index| {
37            let mut index = index.borrow_mut();
38            let new_value = match *index {
39                Some(val) => val.wrapping_add(1) & POSITIVE_MASK,
40                None => rand::rng().random_range(0..=MAX) & POSITIVE_MASK,
41            };
42            *index = Some(new_value);
43            new_value
44        })
45    }
46
47    pub fn reset(&self) {
48        let new_value = rand::rng().random_range(0..=MAX).abs();
49        THREAD_LOCAL_INDEX.with(|index| {
50            *index.borrow_mut() = Some(new_value);
51        });
52    }
53}
54
55impl fmt::Display for ThreadLocalIndex {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        THREAD_LOCAL_INDEX.with(|index| {
58            write!(
59                f,
60                "ThreadLocalIndex {{ thread_local_index={} }}",
61                index.borrow().unwrap_or(0)
62            )
63        })
64    }
65}