Soul 网关源码阅读系列之权重轮询的负载均衡策略

之前在讲divide插件时，只讲了HashLoadBalance，RandomLoadBalance两种负载均衡方式，这次我们一起来看一下RoundRobinLoadBalance，轮询。

事实上，soul 提供的负载均衡方式是可设置权重的，所以代码比普通的负载均衡负载一点。

RoundRobinLoadBalance 维护了一个 methodWeightMap 作为各个upstream的缓存，同时维护了AtomicBoolean类型变量updateLock作为轻量级锁来保证一致性。

接下来看代码：

    @Override
    public DivideUpstream doSelect(final List<DivideUpstream> upstreamList, final String ip) {
        // 根据当前的url，获取WeightedRoundRobin缓存
        String key = upstreamList.get(0).getUpstreamUrl();
        ConcurrentMap<String, WeightedRoundRobin> map = methodWeightMap.get(key);
        if (map == null) {
            methodWeightMap.putIfAbsent(key, new ConcurrentHashMap<>(16));
            map = methodWeightMap.get(key);
        }
        int totalWeight = 0;
        long maxCurrent = Long.MIN_VALUE;
        long now = System.currentTimeMillis();
        DivideUpstream selectedInvoker = null;
        WeightedRoundRobin selectedWRR = null;
        // 遍历所有可选的主机
        for (DivideUpstream upstream : upstreamList) {
            String rKey = upstream.getUpstreamUrl();
            // 获取 WeightedRoundRobin
            WeightedRoundRobin weightedRoundRobin = map.get(rKey);
            // 获取当前主机的权重
            int weight = getWeight(upstream);
            // 设置自己的权重
            if (weightedRoundRobin == null) {
                weightedRoundRobin = new WeightedRoundRobin();
                weightedRoundRobin.setWeight(weight);
                map.putIfAbsent(rKey, weightedRoundRobin);
            }
            if (weight != weightedRoundRobin.getWeight()) {
                //weight changed
                weightedRoundRobin.setWeight(weight);
            }
            // 加上自己的权重
            long cur = weightedRoundRobin.increaseCurrent();
            weightedRoundRobin.setLastUpdate(now);
            // 对比，如果当前主机优先级最高，则设置为 selected
            if (cur > maxCurrent) {
                maxCurrent = cur;
                selectedInvoker = upstream;
                selectedWRR = weightedRoundRobin;
            }
            // 总权重累加
            totalWeight += weight;
        }
        // 以CAS操作来实现轻量级的同步锁
        if (!updateLock.get() && upstreamList.size() != map.size() && updateLock.compareAndSet(false, true)) {
            try {
                // copy -> modify -> update reference
                // 重设缓存
                ConcurrentMap<String, WeightedRoundRobin> newMap = new ConcurrentHashMap<>(map);
                newMap.entrySet().removeIf(item -> now - item.getValue().getLastUpdate() > recyclePeriod);
                methodWeightMap.put(key, newMap);
            } finally {
                updateLock.set(false);
            }
        }
        if (selectedInvoker != null) {
            // 如果选出了主机，则使其权重减去总权重，以使其在下一次调用中权重降低
            selectedWRR.sel(totalWeight);
            return selectedInvoker;
        }
        // should not happen here
        return upstreamList.get(0);
    }

作为对比，我们再来看一下 Spring ribbon 的轮询关键代码，com.netflix.loadbalancer.RoundRobinRule：

    public Server choose(ILoadBalancer lb, Object key) {
        if (lb == null) {
            log.warn("no load balancer");
            return null;
        }

        Server server = null;
        int count = 0;
        // 最多尝试 10 次，10 次后仍然没有可用主机时，直接失败
        while (server == null && count++ < 10) {
            List<Server> reachableServers = lb.getReachableServers();
            List<Server> allServers = lb.getAllServers();
            int upCount = reachableServers.size();
            int serverCount = allServers.size();

            if ((upCount == 0) || (serverCount == 0)) {
                log.warn("No up servers available from load balancer: " + lb);
                return null;
            }

            // 计数器加一并按照当前主机数取模
            int nextServerIndex = incrementAndGetModulo(serverCount);
            server = allServers.get(nextServerIndex);

            if (server == null) {
                /* Transient. */
                Thread.yield();
                continue;
            }

            // 主机探活
            if (server.isAlive() && (server.isReadyToServe())) {
                return (server);
            }

            // Next.
            server = null;
        }

        if (count >= 10) {
            log.warn("No available alive servers after 10 tries from load balancer: "
                    + lb);
        }
        return server;
    }
    
    
    private int incrementAndGetModulo(int modulo) {
        for (;;) {
            // 计数器加一
            int current = nextServerCyclicCounter.get();
            // 取模
            int next = (current + 1) % modulo;
            // CAS 操作保证一致性
            if (nextServerCyclicCounter.compareAndSet(current, next))
                return next;
        }
    }