rocketmq_client_rust/common/
thread_local_index.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
15#![allow(clippy::missing_const_for_thread_local)]
16use std::cell::RefCell;
17use std::fmt;
18
19use rand::RngExt;
20use rand::SeedableRng;
21
22thread_local! {
23    static THREAD_LOCAL_INDEX: RefCell<Option<i32>> = const {RefCell::new(None)};
24    static THREAD_LOCAL_RNG: RefCell<rand::rngs::SmallRng> = RefCell::new(rand::rngs::SmallRng::seed_from_u64(rand::random()));
25}
26
27const POSITIVE_MASK: i32 = 0x7FFFFFFF;
28const MAX: i32 = i32::MAX;
29
30#[derive(Default, Clone)]
31pub struct ThreadLocalIndex;
32
33impl ThreadLocalIndex {
34    pub fn increment_and_get(&self) -> i32 {
35        THREAD_LOCAL_INDEX.with(|index| {
36            let mut index = index.borrow_mut();
37            let new_value = match *index {
38                Some(val) => val.wrapping_add(1),
39                None => THREAD_LOCAL_RNG.with(|rng| rng.borrow_mut().random::<i32>()),
40            };
41            *index = Some(new_value);
42            new_value & POSITIVE_MASK
43        })
44    }
45
46    pub fn reset(&self) {
47        THREAD_LOCAL_INDEX.with(|index| {
48            let new_value = THREAD_LOCAL_RNG.with(|rng| rng.borrow_mut().random_range(0..MAX));
49            *index.borrow_mut() = Some(new_value);
50        });
51    }
52}
53
54impl fmt::Display for ThreadLocalIndex {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        THREAD_LOCAL_INDEX.with(|index| {
57            write!(
58                f,
59                "ThreadLocalIndex {{ thread_local_index={} }}",
60                index.borrow().unwrap_or(0)
61            )
62        })
63    }
64}
rocketmq_client_rust/common/thread_local_index.rs

rocketmq_client_rust/common/
thread_local_index.rs
