|
1 | 1 | package io.github.dunwu.javatech; |
2 | 2 |
|
3 | | -import io.github.dunwu.javatech.support.HashStrategy; |
4 | | -import io.github.dunwu.javatech.support.MurmurHashStrategy; |
| 3 | +import java.nio.charset.StandardCharsets; |
| 4 | +import java.security.MessageDigest; |
| 5 | +import java.security.NoSuchAlgorithmException; |
| 6 | +import java.util.List; |
| 7 | +import java.util.Map; |
| 8 | +import java.util.TreeMap; |
| 9 | +import java.util.concurrent.ConcurrentHashMap; |
| 10 | +import java.util.concurrent.ConcurrentMap; |
5 | 11 |
|
6 | | -import java.util.*; |
| 12 | +public class ConsistentHashLoadBalance<N extends Node> extends BaseLoadBalance<N> implements LoadBalance<N> { |
7 | 13 |
|
8 | | -public class ConsistentHashLoadBalance<V extends Node> implements LoadBalance<V> { |
| 14 | + private final ConcurrentMap<String, ConsistentHashSelector<?>> selectors = new ConcurrentHashMap<>(); |
9 | 15 |
|
10 | | - private HashStrategy hashStrategy = new MurmurHashStrategy(); |
11 | | - |
12 | | - private final static int VIRTUAL_NODE_SIZE = 1000; |
| 16 | + @SuppressWarnings("unchecked") |
| 17 | + @Override |
| 18 | + protected N doSelect(List<N> nodes, String ip) { |
| 19 | + // 获取 nodes 原始的 hashcode |
| 20 | + int identityHashCode = System.identityHashCode(nodes); |
| 21 | + ConsistentHashSelector<N> selector = (ConsistentHashSelector<N>) selectors.get(ip); |
| 22 | + if (selector == null || selector.identityHashCode != identityHashCode) { |
| 23 | + selectors.put(ip, new ConsistentHashSelector<>(nodes, identityHashCode, 1000)); |
| 24 | + selector = (ConsistentHashSelector<N>) selectors.get(ip); |
| 25 | + } |
| 26 | + return selector.select(ip); |
| 27 | + } |
13 | 28 |
|
14 | | - private final static String VIRTUAL_NODE_SUFFIX = "&&"; |
| 29 | + /** |
| 30 | + * @param <N> |
| 31 | + */ |
| 32 | + private static final class ConsistentHashSelector<N extends Node> { |
15 | 33 |
|
16 | | - private Set<V> nodes = new LinkedHashSet<>(); |
| 34 | + /** |
| 35 | + * 使用 TreeMap 存储 Node 虚拟节点 |
| 36 | + */ |
| 37 | + private final TreeMap<Long, N> virtualNodes; |
17 | 38 |
|
18 | | - private TreeMap<Integer, V> hashRing = new TreeMap<>(); |
| 39 | + private final int identityHashCode; |
19 | 40 |
|
20 | | - @Override |
21 | | - public void buildNodes(final Collection<V> collection) { |
22 | | - this.nodes = new LinkedHashSet<>(collection); |
23 | | - this.hashRing = buildConsistentHashRing(this.nodes); |
24 | | - } |
| 41 | + ConsistentHashSelector(List<N> nodes, int identityHashCode, Integer replicaNum) { |
| 42 | + this.virtualNodes = new TreeMap<>(); |
| 43 | + this.identityHashCode = identityHashCode; |
| 44 | + // 获取虚拟节点数,默认为 100 |
| 45 | + if (replicaNum == null) { |
| 46 | + replicaNum = 100; |
| 47 | + } |
| 48 | + for (N node : nodes) { |
| 49 | + for (int i = 0; i < replicaNum / 4; i++) { |
| 50 | + // 对 url 进行 md5 运算,得到一个长度为16的字节数组 |
| 51 | + byte[] digest = md5(node.getUrl()); |
| 52 | + // 对 digest 部分字节进行 4 次 hash 运算,得到四个不同的 long 型正整数 |
| 53 | + for (int j = 0; j < 4; j++) { |
| 54 | + // h = 0 时,取 digest 中下标为 0 ~ 3 的4个字节进行位运算 |
| 55 | + // h = 1 时,取 digest 中下标为 4 ~ 7 的4个字节进行位运算 |
| 56 | + // h = 2, h = 3 时过程同上 |
| 57 | + long m = hash(digest, j); |
| 58 | + // 将 hash 到 node 的映射关系存储到 virtualNodes 中, |
| 59 | + // virtualNodes 需要提供高效的查询操作,因此选用 TreeMap 作为存储结构 |
| 60 | + virtualNodes.put(m, node); |
| 61 | + } |
| 62 | + } |
| 63 | + } |
| 64 | + } |
25 | 65 |
|
26 | | - @Override |
27 | | - public void addNode(V node) { |
28 | | - this.nodes.add(node); |
29 | | - this.hashRing = buildConsistentHashRing(this.nodes); |
30 | | - } |
| 66 | + public N select(String key) { |
| 67 | + byte[] digest = md5(key); |
| 68 | + return selectForKey(hash(digest, 0)); |
| 69 | + } |
31 | 70 |
|
32 | | - @Override |
33 | | - public void removeNode(V node) { |
34 | | - this.nodes.removeIf(v -> v.equals(node)); |
35 | | - this.hashRing = buildConsistentHashRing(this.nodes); |
36 | | - } |
| 71 | + private N selectForKey(long hash) { |
| 72 | + Map.Entry<Long, N> entry = virtualNodes.ceilingEntry(hash); |
| 73 | + if (entry == null) { |
| 74 | + entry = virtualNodes.firstEntry(); |
| 75 | + } |
| 76 | + return entry.getValue(); |
| 77 | + } |
37 | 78 |
|
38 | | - @Override |
39 | | - public V select() { |
40 | | - return next(UUID.randomUUID().toString()); |
41 | 79 | } |
42 | 80 |
|
43 | | - public V next(String key) { |
44 | | - int hashCode = hashStrategy.hashCode(key); |
45 | | - // 向右找到第一个 key |
46 | | - Map.Entry<Integer, V> entry = hashRing.ceilingEntry(hashCode); |
47 | | - if (entry == null) { |
48 | | - // 想象成一个环,超过尾部则取第一个 key |
49 | | - entry = hashRing.firstEntry(); |
50 | | - } |
51 | | - return entry.getValue(); |
| 81 | + /** |
| 82 | + * 计算 hash 值 |
| 83 | + */ |
| 84 | + public static long hash(byte[] digest, int number) { |
| 85 | + return (((long) (digest[3 + number * 4] & 0xFF) << 24) |
| 86 | + | ((long) (digest[2 + number * 4] & 0xFF) << 16) |
| 87 | + | ((long) (digest[1 + number * 4] & 0xFF) << 8) |
| 88 | + | (digest[number * 4] & 0xFF)) |
| 89 | + & 0xFFFFFFFFL; |
52 | 90 | } |
53 | 91 |
|
54 | | - private TreeMap<Integer, V> buildConsistentHashRing(Set<V> nodes) { |
55 | | - TreeMap<Integer, V> hashRing = new TreeMap<>(); |
56 | | - for (V node : nodes) { |
57 | | - for (int i = 0; i < VIRTUAL_NODE_SIZE; i++) { |
58 | | - // 新增虚拟节点的方式如果有影响,也可以抽象出一个由物理节点扩展虚拟节点的类 |
59 | | - hashRing.put(hashStrategy.hashCode(node + VIRTUAL_NODE_SUFFIX + i), node); |
60 | | - } |
| 92 | + /** |
| 93 | + * 计算 MD5 值 |
| 94 | + */ |
| 95 | + public static byte[] md5(String value) { |
| 96 | + MessageDigest md5; |
| 97 | + try { |
| 98 | + md5 = MessageDigest.getInstance("MD5"); |
| 99 | + } catch (NoSuchAlgorithmException e) { |
| 100 | + throw new IllegalStateException(e.getMessage(), e); |
61 | 101 | } |
62 | | - return hashRing; |
| 102 | + md5.reset(); |
| 103 | + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); |
| 104 | + md5.update(bytes); |
| 105 | + return md5.digest(); |
63 | 106 | } |
64 | 107 |
|
65 | 108 | } |
0 commit comments