模仿dubbo的ConsistentHashLoadBalance实现一个自定义负载均衡案例

通过了解一致性hash原理，这篇文章就够了，可以了解一致性哈希算法的原理，下面就通过一个小案例来事件一下。

背景

假设有个系统，用户量非常大，需要针对不同用户访问不同的ip地址进行分流，且要求同一个用户多次访问均为同一个ip。

这种场景就是典型的哈希一致性问题，下面就通过模仿dubbo的ConsistentHashLoadBalance，来尝试来实现一下。

原理回顾

要用一致性哈希算法实现上述需求，需要经过如下几个步骤：

• 构建一个哈希环
• 将ip地址均匀的映射到哈希环上
• 将用户请求映射到哈希环上，从而定位相应的访问IP

难点分析

主要的难点在于找到合适的hash算法，hash算法不好，会影响元素的散列程度，这里就直接使用dubbo中的hash算法。

        private long hash(byte[] digest, int number) {
            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;
        }

构建hash环

java中可以使用TreeMap来构建哈希环

private static TreeMap<Long, String> virtualInvokers = new TreeMap<Long, String>();

将 ip地址映射到哈希环上

由于是本地实验，就直接在静态代码块中进行，参考ConsistentHashLoadBalance，为了避免ip分布不均，引入虚拟节点。

    static {
        for (String ip : IP_LIST) {
            for (int i = 0; i < VIRTUAL_NODES; i++) {// 虚拟节点，避免分布不均的情况
                byte[] digest = Bytes.getMD5(ip + i);
                for (int h = 0; h < 4; h++) {
                    long m = hash(digest, h);
                    virtualInvokers.put(m, ip);// 映射ip地址
                }


            }
        }
    }

用户请求定位IP

主要是利用 TreeMap的ceilingEntry 方法，返回与大于或等于给定键元素(ele)的最小键元素链接的键值对，即哈希环顺时针方向第一个ip元素。

  public static String getServer(String invocation) {

        byte[] digest = Bytes.getMD5(invocation);
        long hash = hash(digest, 0);

        /**
         * 利用 TreeMap的ceilingEntry 方法
         * 返回与大于或等于给定键元素(ele)的最小键元素链接的键值对
         * 即哈希环顺时针方向第一个ip元素。
         */
        Map.Entry<Long, String> entry = virtualInvokers.ceilingEntry(hash);
        if (entry == null) {// 如果变量整个环都没有，则返回第一个
            entry = virtualInvokers.firstEntry();
        }


        return entry.getValue();
    }

完整测试代码

package com.fast.alibaba.service;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Map;

import java.util.TreeMap;

/**
 * 照搬dubbo的Bytes类中的部分方法，为了方便计算hash
 */
class Bytes {
    private static final String C64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="; //default base64.

    private static ThreadLocal<MessageDigest> MD = new ThreadLocal<MessageDigest>();

    private Bytes() {
    }

    /**
     * get md5.
     *
     * @param str input string.
     * @return MD5 byte array.
     */
    public static byte[] getMD5(String str) {
        return getMD5(str.getBytes());
    }

    /**
     * get md5.
     *
     * @param source byte array source.
     * @return MD5 byte array.
     */
    public static byte[] getMD5(byte[] source) {
        MessageDigest md = getMessageDigest();
        return md.digest(source);
    }

    private static MessageDigest getMessageDigest() {
        MessageDigest ret = MD.get();
        if (ret == null) {
            try {
                ret = MessageDigest.getInstance("MD5");
                MD.set(ret);
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException(e);
            }
        }
        return ret;
    }

}


public class MyConsistentHashLoadBalance {


    public static String[] IP_LIST = {"36.4.192.0", "39.153.127.255", "42.55.181.0", "59.44.141.43", "59.44.143.38",
            "59.44.143.11", "59.44.143.43", "59.44.143.53", "59.44.143.31", "59.44.143.33", "2.16.10.0",
            "2.16.154.0", "2.16.159.0", "2.16.168.0", "2.20.254.0", "2.21.1.0", "2.21.90.0", "2.22.237.0",
            "2.23.167.0", "2.56.240.0", "2.59.214.0", "2.60.0.0", "2.92.0.0", "5.1.48.0", "5.2.32.0",
            "5.3.0.0", "5.8.0.0", "5.8.19.0", "5.8.36.0"};

    private static final int VIRTUAL_NODES = 160;// 虚拟节点个数

    /**
     * 用 TreeMap构造一个哈希环
     */
    private static TreeMap<Long, String> virtualInvokers = new TreeMap<Long, String>();

    /**
     * 将 ip地址映射到哈希环上
     */
    static {
        for (String ip : IP_LIST) {
            for (int i = 0; i < VIRTUAL_NODES; i++) {// 虚拟节点，避免分布不均的情况
                byte[] digest = Bytes.getMD5(ip + i);
                for (int h = 0; h < 4; h++) {
                    long m = hash(digest, h);
                    virtualInvokers.put(m, ip);// 映射ip地址
                }


            }
        }
    }


    public static String getServer(String invocation) {

        byte[] digest = Bytes.getMD5(invocation);
        long hash = hash(digest, 0);

        /**
         * 利用 TreeMap的ceilingEntry 方法
         * 返回与大于或等于给定键元素(ele)的最小键元素链接的键值对
         * 即哈希环顺时针方向第一个ip元素。
         */
        Map.Entry<Long, String> entry = virtualInvokers.ceilingEntry(hash);
        if (entry == null) {// 如果变量整个环都没有，则返回第一个
            entry = virtualInvokers.firstEntry();
        }


        return entry.getValue();
    }


    public static void main(String[] args) {
        for (int i = 0; i < 200; i++) {
            System.out.println(getServer("userId" + i) + "---" + i);
        }
    }

    /**
     * hash 算法，照抄dubbo中ConsistentHashLoadBalance的
     * @param digest
     * @param number
     * @return
     */
    private static long hash(byte[] digest, int number) {
        return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                | (digest[number * 4] & 0xFF))
                & 0xFFFFFFFFL;
    }
}