apigate_core/balancing/
consistent_hash.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use super::{BalanceCtx, Balancer};
4
5pub struct ConsistentHash {
6    offset: AtomicUsize,
7}
8
9impl ConsistentHash {
10    pub fn new() -> Self {
11        Self {
12            offset: AtomicUsize::new(0),
13        }
14    }
15}
16
17impl Default for ConsistentHash {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl Balancer for ConsistentHash {
24    #[inline]
25    fn pick(&self, ctx: &BalanceCtx) -> Option<usize> {
26        let candidate_len = ctx.candidate_len();
27        if candidate_len == 0 {
28            return None;
29        }
30
31        match ctx.affinity {
32            Some(key) => {
33                let hash = xxhash_rust::xxh3::xxh3_64(key.as_str().as_bytes());
34                pick_candidate(hash, ctx)
35            }
36            None => {
37                let n = self.offset.fetch_add(1, Ordering::Relaxed);
38
39                let nth = if candidate_len.is_power_of_two() {
40                    n & (candidate_len - 1)
41                } else {
42                    n % candidate_len
43                };
44
45                ctx.candidate_index(nth)
46            }
47        }
48    }
49}
50
51#[inline]
52fn pick_candidate(hash: u64, ctx: &BalanceCtx<'_>) -> Option<usize> {
53    let bucket = jump_consistent_hash(hash, ctx.candidate_len())?;
54    ctx.candidate_index(bucket)
55}
56
57#[inline]
58fn jump_consistent_hash(mut key: u64, buckets: usize) -> Option<usize> {
59    if buckets == 0 {
60        return None;
61    }
62
63    let mut b: i64 = -1;
64    let mut j: i64 = 0;
65    let buckets = buckets as i64;
66
67    while j < buckets {
68        b = j;
69        key = key.wrapping_mul(2862933555777941757).wrapping_add(1);
70        j = ((b + 1) * (1_i64 << 31)) / (((key >> 33) + 1) as i64);
71    }
72
73    Some(b as usize)
74}
apigate_core/balancing/consistent_hash.rs

apigate_core/balancing/
consistent_hash.rs
