cassandra_protocol/
token.rs

1use crate::error::Error;
2use bytes::Buf;
3use derive_more::Constructor;
4use std::cmp::min;
5use std::convert::TryFrom;
6use std::num::Wrapping;
7
8const C1: Wrapping<i64> = Wrapping(0x87c3_7b91_1142_53d5_u64 as i64);
9const C2: Wrapping<i64> = Wrapping(0x4cf5_ad43_2745_937f_u64 as i64);
10
11/// A token on the ring. Only Murmur3 tokens are supported for now.
12#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Default, Debug, Hash, Constructor)]
13pub struct Murmur3Token {
14    pub value: i64,
15}
16
17impl Murmur3Token {
18    // based on buggy Cassandra implementation
19    pub fn generate(mut routing_key: &[u8]) -> Self {
20        let length = routing_key.len();
21
22        let mut h1: Wrapping<i64> = Wrapping(0);
23        let mut h2: Wrapping<i64> = Wrapping(0);
24
25        while routing_key.len() >= 16 {
26            let mut k1 = Wrapping(routing_key.get_i64_le());
27            let mut k2 = Wrapping(routing_key.get_i64_le());
28
29            k1 *= C1;
30            k1 = rotl64(k1, 31);
31            k1 *= C2;
32            h1 ^= k1;
33
34            h1 = rotl64(h1, 27);
35            h1 += h2;
36            h1 = h1 * Wrapping(5) + Wrapping(0x52dce729);
37
38            k2 *= C2;
39            k2 = rotl64(k2, 33);
40            k2 *= C1;
41            h2 ^= k2;
42
43            h2 = rotl64(h2, 31);
44            h2 += h1;
45            h2 = h2 * Wrapping(5) + Wrapping(0x38495ab5);
46        }
47
48        let mut k1 = Wrapping(0_i64);
49        let mut k2 = Wrapping(0_i64);
50
51        debug_assert!(routing_key.len() < 16);
52
53        if routing_key.len() > 8 {
54            for i in (8..routing_key.len()).rev() {
55                k2 ^= Wrapping(routing_key[i] as i8 as i64) << ((i - 8) * 8);
56            }
57
58            k2 *= C2;
59            k2 = rotl64(k2, 33);
60            k2 *= C1;
61            h2 ^= k2;
62        }
63
64        if !routing_key.is_empty() {
65            for i in (0..min(8, routing_key.len())).rev() {
66                k1 ^= Wrapping(routing_key[i] as i8 as i64) << (i * 8);
67            }
68
69            k1 *= C1;
70            k1 = rotl64(k1, 31);
71            k1 *= C2;
72            h1 ^= k1;
73        }
74
75        h1 ^= Wrapping(length as i64);
76        h2 ^= Wrapping(length as i64);
77
78        h1 += h2;
79        h2 += h1;
80
81        h1 = fmix(h1);
82        h2 = fmix(h2);
83
84        h1 += h2;
85
86        Murmur3Token::new(h1.0)
87    }
88}
89
90impl TryFrom<String> for Murmur3Token {
91    type Error = Error;
92
93    fn try_from(value: String) -> Result<Self, Self::Error> {
94        value
95            .parse()
96            .map_err(|error| format!("Error parsing token: {error}").into())
97            .map(Murmur3Token::new)
98    }
99}
100
101impl From<i64> for Murmur3Token {
102    fn from(value: i64) -> Self {
103        Murmur3Token::new(value)
104    }
105}
106
107#[inline]
108fn rotl64(v: Wrapping<i64>, n: u32) -> Wrapping<i64> {
109    Wrapping((v.0 << n) | (v.0 as u64 >> (64 - n)) as i64)
110}
111
112#[inline]
113fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
114    k ^= Wrapping((k.0 as u64 >> 33) as i64);
115    k *= Wrapping(0xff51afd7ed558ccd_u64 as i64);
116    k ^= Wrapping((k.0 as u64 >> 33) as i64);
117    k *= Wrapping(0xc4ceb9fe1a85ec53_u64 as i64);
118    k ^= Wrapping((k.0 as u64 >> 33) as i64);
119
120    k
121}
122
123#[cfg(test)]
124mod test {
125    use super::*;
126
127    #[test]
128    fn test_generate_murmur3_token() {
129        for s in [
130            ("testvalue", 5965290492934326460),
131            ("testvalue123", 1518494936189046133),
132            ("example_key", -7813763279771224608),
133            ("château", 9114062196463836094),
134        ] {
135            let generated_token = Murmur3Token::generate(s.0.as_bytes());
136            assert_eq!(generated_token.value, s.1);
137        }
138    }
139}