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;
18use std::sync::atomic::AtomicUsize;
19use std::sync::atomic::Ordering;
20
21use rand::RngExt;
22use rand::SeedableRng;
23
24thread_local! {
25    static THREAD_LOCAL_INDEXES: RefCell<Vec<Option<i32>>> = const { RefCell::new(Vec::new()) };
26    static THREAD_LOCAL_RNG: RefCell<rand::rngs::SmallRng> = RefCell::new(rand::rngs::SmallRng::seed_from_u64(rand::random()));
27}
28
29static NEXT_INDEX_ID: AtomicUsize = AtomicUsize::new(0);
30
31const POSITIVE_MASK: i32 = 0x7FFFFFFF;
32const MAX: i32 = i32::MAX;
33
34#[derive(Clone)]
35pub struct ThreadLocalIndex {
36    id: usize,
37}
38
39impl ThreadLocalIndex {
40    pub fn new() -> Self {
41        Self {
42            id: NEXT_INDEX_ID.fetch_add(1, Ordering::Relaxed),
43        }
44    }
45
46    pub fn increment_and_get(&self) -> i32 {
47        THREAD_LOCAL_INDEXES.with(|indexes| {
48            let mut indexes = indexes.borrow_mut();
49            let slot = thread_local_slot(&mut indexes, self.id);
50            let random = || THREAD_LOCAL_RNG.with(|rng| rng.borrow_mut().random::<i32>());
51            let new_value = next_java_index_value(*slot, random);
52            *slot = Some(new_value);
53            new_value & POSITIVE_MASK
54        })
55    }
56
57    pub fn reset(&self) {
58        THREAD_LOCAL_INDEXES.with(|indexes| {
59            let mut indexes = indexes.borrow_mut();
60            let slot = thread_local_slot(&mut indexes, self.id);
61            let new_value = THREAD_LOCAL_RNG.with(|rng| rng.borrow_mut().random_range(0..MAX));
62            *slot = Some(new_value);
63        });
64    }
65}
66
67impl Default for ThreadLocalIndex {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73#[inline]
74fn thread_local_slot(indexes: &mut Vec<Option<i32>>, id: usize) -> &mut Option<i32> {
75    if indexes.len() <= id {
76        indexes.resize(id + 1, None);
77    }
78    &mut indexes[id]
79}
80
81#[inline]
82fn next_java_index_value<F>(current: Option<i32>, mut random: F) -> i32
83where
84    F: FnMut() -> i32,
85{
86    current.unwrap_or_else(&mut random).wrapping_add(1)
87}
88
89impl fmt::Display for ThreadLocalIndex {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        THREAD_LOCAL_INDEXES.with(|indexes| {
92            write!(
93                f,
94                "ThreadLocalIndex {{ thread_local_index={} }}",
95                indexes.borrow().get(self.id).and_then(|value| *value).unwrap_or(0)
96            )
97        })
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn next_java_index_value_increments_existing_value_like_java() {
107        let value = next_java_index_value(Some(41), || 100);
108
109        assert_eq!(value, 42);
110    }
111
112    #[test]
113    fn next_java_index_value_increments_initial_random_value_like_java() {
114        let value = next_java_index_value(None, || 41);
115
116        assert_eq!(value, 42);
117    }
118
119    #[test]
120    fn next_java_index_value_wraps_before_positive_mask_like_java_int_overflow() {
121        let value = next_java_index_value(None, || i32::MAX);
122
123        assert_eq!(value, i32::MIN);
124        assert_eq!(value & POSITIVE_MASK, 0);
125    }
126
127    #[test]
128    fn different_thread_local_index_instances_keep_independent_java_state() {
129        let first = ThreadLocalIndex::new();
130        let second = ThreadLocalIndex::new();
131
132        first.reset();
133        second.reset();
134        let first_value = first.increment_and_get();
135        let second_value = second.increment_and_get();
136
137        assert_eq!(first.increment_and_get(), first_value.wrapping_add(1) & POSITIVE_MASK);
138        assert_eq!(second.increment_and_get(), second_value.wrapping_add(1) & POSITIVE_MASK);
139    }
140
141    #[test]
142    fn cloned_thread_local_index_shares_the_same_java_instance_state() {
143        let original = ThreadLocalIndex::new();
144        original.reset();
145        let cloned = original.clone();
146
147        let first = original.increment_and_get();
148        let second = cloned.increment_and_get();
149
150        assert_eq!(second, first.wrapping_add(1) & POSITIVE_MASK);
151    }
152}
rocketmq_client_rust/common/thread_local_index.rs

rocketmq_client_rust/common/
thread_local_index.rs
