rocketmq_model/
consistent_hash.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
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use md5::Digest;
19use rocketmq_error::Result;
20
21/// A physical or virtual node that can be placed on a hash ring.
22pub trait Node {
23    fn get_key(&self) -> &str;
24}
25
26/// Converts a stable string key into a ring position.
27pub trait HashFunction {
28    fn hash(&self, key: &str) -> i64;
29}
30
31/// A virtual replica of a physical node.
32#[derive(Debug, Clone)]
33pub struct VirtualNode<T: Node> {
34    physical_node: T,
35    replica_index: i32,
36    key: String,
37}
38
39impl<T: Node> VirtualNode<T> {
40    pub fn new(physical_node: T, replica_index: i32) -> Self {
41        let key = format!("{}-{replica_index}", physical_node.get_key());
42        Self {
43            physical_node,
44            replica_index,
45            key,
46        }
47    }
48
49    pub fn is_virtual_node_of(&self, physical_node: &T) -> bool {
50        self.physical_node.get_key() == physical_node.get_key()
51    }
52
53    pub fn physical_node(&self) -> &T {
54        &self.physical_node
55    }
56
57    pub fn replica_index(&self) -> i32 {
58        self.replica_index
59    }
60}
61
62impl<T: Node> Node for VirtualNode<T> {
63    fn get_key(&self) -> &str {
64        &self.key
65    }
66}
67
68/// MD5-compatible hash used by the Java allocation algorithm.
69#[derive(Debug, Clone, Default)]
70pub struct MD5Hash;
71
72impl MD5Hash {
73    pub fn new() -> Self {
74        Self
75    }
76}
77
78impl HashFunction for MD5Hash {
79    fn hash(&self, key: &str) -> i64 {
80        let mut hasher = md5::Md5::new();
81        hasher.update(key.as_bytes());
82        let digest = hasher.finalize();
83        digest[..4]
84            .iter()
85            .fold(0_i64, |hash, byte| (hash << 8) | i64::from(*byte))
86    }
87}
88
89/// Consistent hash ring with a fixed number of virtual nodes per physical node.
90#[derive(Clone)]
91pub struct ConsistentHashRouter<T: Node + Clone> {
92    ring: BTreeMap<i64, VirtualNode<T>>,
93    hash_function: Arc<dyn HashFunction>,
94}
95
96impl<T: Node + Clone> ConsistentHashRouter<T> {
97    pub fn new(physical_nodes: Vec<T>, virtual_node_count: i32) -> Self {
98        Self::new_with_hash_function(physical_nodes, virtual_node_count, Arc::new(MD5Hash))
99    }
100
101    pub fn try_new(physical_nodes: Vec<T>, virtual_node_count: i32) -> Result<Self> {
102        Self::try_new_with_hash_function(physical_nodes, virtual_node_count, Arc::new(MD5Hash))
103    }
104
105    pub fn new_with_hash_function(
106        physical_nodes: Vec<T>,
107        virtual_node_count: i32,
108        hash_function: Arc<dyn HashFunction>,
109    ) -> Self {
110        if virtual_node_count < 0 {
111            return Self {
112                ring: BTreeMap::new(),
113                hash_function,
114            };
115        }
116        let mut router = Self {
117            ring: BTreeMap::new(),
118            hash_function,
119        };
120        for node in physical_nodes {
121            router.add_node(node, virtual_node_count);
122        }
123        router
124    }
125
126    pub fn try_new_with_hash_function(
127        physical_nodes: Vec<T>,
128        virtual_node_count: i32,
129        hash_function: Arc<dyn HashFunction>,
130    ) -> Result<Self> {
131        if virtual_node_count < 0 {
132            return Err(crate::error::invalid_argument(format!(
133                "illegal virtual node counts: {virtual_node_count}"
134            )));
135        }
136        let mut router = Self {
137            ring: BTreeMap::new(),
138            hash_function,
139        };
140        for node in physical_nodes {
141            router.try_add_node(node, virtual_node_count)?;
142        }
143        Ok(router)
144    }
145
146    pub fn add_node(&mut self, physical_node: T, virtual_node_count: i32) {
147        let _ = self.try_add_node(physical_node, virtual_node_count);
148    }
149
150    pub fn try_add_node(&mut self, physical_node: T, virtual_node_count: i32) -> Result<()> {
151        if virtual_node_count < 0 {
152            return Err(crate::error::invalid_argument(format!(
153                "illegal virtual node counts: {virtual_node_count}"
154            )));
155        }
156        let existing = self.get_existing_replicas(&physical_node);
157        for replica in 0..virtual_node_count {
158            let virtual_node = VirtualNode::new(physical_node.clone(), existing + replica);
159            self.ring
160                .insert(self.hash_function.hash(virtual_node.get_key()), virtual_node);
161        }
162        Ok(())
163    }
164
165    pub fn remove_node(&mut self, physical_node: &T) {
166        self.ring
167            .retain(|_, virtual_node| !virtual_node.is_virtual_node_of(physical_node));
168    }
169
170    pub fn route_node_ref(&self, object_key: &str) -> Option<&T> {
171        let hash = self.hash_function.hash(object_key);
172        self.ring
173            .range(hash..)
174            .next()
175            .or_else(|| self.ring.iter().next())
176            .map(|(_, node)| node.physical_node())
177    }
178
179    pub fn route_node(&self, object_key: &str) -> Option<T> {
180        self.route_node_ref(object_key).cloned()
181    }
182
183    pub fn get_existing_replicas(&self, physical_node: &T) -> i32 {
184        self.ring
185            .values()
186            .filter(|node| node.is_virtual_node_of(physical_node))
187            .count() as i32
188    }
189
190    pub fn size(&self) -> usize {
191        self.ring.len()
192    }
193
194    pub fn is_empty(&self) -> bool {
195        self.ring.is_empty()
196    }
197}
rocketmq_model/consistent_hash.rs

rocketmq_model/
consistent_hash.rs
