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 Balancer for ConsistentHash {
18    #[inline]
19    fn pick(&self, ctx: &BalanceCtx) -> Option<usize> {
20        let candidate_len = ctx.candidate_len();
21        if candidate_len == 0 {
22            return None;
23        }
24
25        match ctx.affinity {
26            Some(key) => {
27                let hash = xxhash_rust::xxh3::xxh3_64(key.as_str().as_bytes());
28                pick_candidate(hash, ctx)
29            }
30            None => {
31                let n = self.offset.fetch_add(1, Ordering::Relaxed);
32
33                let nth = if candidate_len.is_power_of_two() {
34                    n & (candidate_len - 1)
35                } else {
36                    n % candidate_len
37                };
38
39                ctx.candidate_index(nth)
40            }
41        }
42    }
43}
44
45#[inline]
46fn pick_candidate(hash: u64, ctx: &BalanceCtx<'_>) -> Option<usize> {
47    let bucket = jump_consistent_hash(hash, ctx.candidate_len())?;
48    ctx.candidate_index(bucket)
49}
50
51#[inline]
52fn jump_consistent_hash(mut key: u64, buckets: usize) -> Option<usize> {
53    if buckets == 0 {
54        return None;
55    }
56
57    let mut b: i64 = -1;
58    let mut j: i64 = 0;
59    let buckets = buckets as i64;
60
61    while j < buckets {
62        b = j;
63        key = key.wrapping_mul(2862933555777941757).wrapping_add(1);
64        j = ((b + 1) * (1_i64 << 31)) / (((key >> 33) + 1) as i64);
65    }
66
67    Some(b as usize)
68}
apigate_core/balancing/consistent_hash.rs

apigate_core/balancing/
consistent_hash.rs
