单独使用Ribbon进行负载均衡

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

使用时调用getServer（）方法传key标识和对应的服务列表即可。

import com.netflix.loadbalancer.*;
import kl.cds.rule.CommonRoundRobinRule;
import kl.cds.utils.PingUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 密管机负载均衡
 *
 * @Author ncx
 * @Date 2021/5/12
 */
public class LoadBalancerMgr {

    private static String defaultKey = "default";

    public static ConcurrentHashMap<String, BaseLoadBalancer> balancerMap = new ConcurrentHashMap<>();

    private static LoadBalancerMgr instance = new LoadBalancerMgr();

    public static LoadBalancerMgr getInstance() {
        return instance;
    }

    public Server
    getServer(List<Server> serverList) {
        return getServer(defaultKey, serverList);
    }

	// 通过hash表维护多个负载均衡列表
    public Server getServer(String key, List<Server> serverList) {
        // 读取key对应的loadbalancer
        BaseLoadBalancer baseLoadBalancer = balancerMap.get(key);
        if (baseLoadBalancer == null) {
            synchronized (LoadBalancerMgr.class) {
                // loadbalancer为空则新建
                baseLoadBalancer = LoadBalancerBuilder.newBuilder()
                        .withRule(new CommonRoundRobinRule())
                        .buildFixedServerListLoadBalancer(serverList);
                balancerMap.put(key, baseLoadBalancer);
                return getServerFromLoadBalancer(key, baseLoadBalancer);
            }
        } else {
            synchronized (LoadBalancerMgr.class) {
                List<Server> aliveServer = baseLoadBalancer.getServerList(AbstractLoadBalancer.ServerGroup.ALL);
                List<Server> finalServerList = new ArrayList<>();
                finalServerList.addAll(aliveServer);
                if (PingUtil.errorMap.get(key) != null && !PingUtil.errorMap.get(key).isEmpty()) {
                    PingUtil.errorMap.forEach((str, servers) -> finalServerList.addAll(servers));
                }
                // 传过来的是相同的serverList,直接调用负载获取server
                if (!isListEqual(finalServerList, serverList)) {
                    baseLoadBalancer = LoadBalancerBuilder.newBuilder()
                            .withRule(new CommonRoundRobinRule())
                            .buildFixedServerListLoadBalancer(serverList);
                    balancerMap.put(key, baseLoadBalancer);
                }
                return getServerFromLoadBalancer(key, baseLoadBalancer);
            }
        }
    }

    /**
     * 判断list内容是否相同
     * @param l0
     * @param l1
     * @return
     */
    private static boolean isListEqual(List<Server> l0, List<Server> l1) {
        boolean flag1 = l0.containsAll(l1);
        boolean flag2 = l1.containsAll(l0);
        return flag1 && flag2;
    }

    /**
     * 负载均衡实现
     *
     * @return
     */
    private Server getServerFromLoadBalancer(String key, BaseLoadBalancer baseLoadBalancer) {
        CommonRoundRobinRule rule = (CommonRoundRobinRule) baseLoadBalancer.getRule();
        Server server = rule.choose(key, baseLoadBalancer);
        return server;
    }

    public void shutdown() {
        shutdown(defaultKey);
    }

    public void shutdown(String key) {
        BaseLoadBalancer baseLoadBalancer = balancerMap.get(key);
        baseLoadBalancer.shutdown();
    }

}

/*
 *
 * Copyright 2013 Netflix, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.BaseLoadBalancer;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import kl.cds.utils.PingUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Ribbon的负载策略，单独使用时都不太好，自定义负载策略
 * The most well known and basic load balancing strategy, i.e. Round Robin Rule.
 *
 * @author stonse
 * @author Nikos Michalakis <nikos@netflix.com>
 *
 */
public class CommonRoundRobinRule extends AbstractLoadBalancerRule {

    private AtomicInteger nextServerCyclicCounter;
    private PingUtil pingUtil = new PingUtil();

    private static Logger log = LoggerFactory.getLogger(CommonRoundRobinRule.class);

    public CommonRoundRobinRule() {
        nextServerCyclicCounter = new AtomicInteger(0);
    }

    public CommonRoundRobinRule(ILoadBalancer lb) {
        this();
        setLoadBalancer(lb);
    }

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

        Server server = null;
        int count = 0;

        List<Server> reachableServers = lb.getReachableServers();
        List<Server> allServers = lb.getAllServers();

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

        pingUtil.setKey(key);
        server = getAliveServer(allServers, count);

        if (server == null) {
            log.warn("No available alive servers");
            return null;
        } else {
            return server;
        }
    }

    /**
     * Inspired by the implementation of {@link AtomicInteger#incrementAndGet()}.
     *
     * @param modulo The modulo to bound the value of the counter.
     * @return The next value.
     */
    private int incrementAndGetModulo(int modulo) {
        for (;;) {
            int current = nextServerCyclicCounter.get();
            int next = (current + 1) % modulo;
            if (nextServerCyclicCounter.compareAndSet(current, next)) {
                return next;
            }
        }
    }

    @Override
    public Server choose(Object key) {
        return choose(key.toString(), (BaseLoadBalancer) getLoadBalancer());
    }

    /**
     * 获得可用服务
     * @param allServers
     * @return
     */
    private Server getAliveServer(List<Server> allServers, int count) {
        // 当查询超过可用服务列表数量还没有可用服务时返回null
        if (count > allServers.size()) {
            log.warn("No available alive servers");
            return null;
        }
        int nextServerIndex = incrementAndGetModulo(allServers.size());

        // 通过cas获取可用服务地址
        boolean alive;
        Server server;
        try {
            server = allServers.get(nextServerIndex);
        } catch (Exception e) {
            nextServerIndex = incrementAndGetModulo(allServers.size());
            server = allServers.get(nextServerIndex);
        }

        long start = System.currentTimeMillis();
        alive = pingUtil.isTelnetAlive(server);
        long end = System.currentTimeMillis();
        System.out.println("telnet: " + (end-start));

        // 服务可用返回服务，不可用查找下一个可用服务
        if (alive) {
            return server;
        } else {
            count++;
            return getAliveServer(allServers, count);
        }
    }

}

import com.netflix.loadbalancer.AbstractLoadBalancer;
import com.netflix.loadbalancer.Server;
import kl.cds.loadbalancer.LoadBalancerMgr;
import org.apache.http.HttpResponse;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.impl.client.DefaultHttpClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketAddress;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @Author ncx
 * @Date 2021/8/19
 */
public class PingUtil {
    private static Logger log = LoggerFactory.getLogger(PingUtil.class);

    private String key;
    public static Map<String, List<Server>> errorMap = new HashMap<>();

    // 启动定时任务
    static {
        scheduleTask();
    }

    /**
     * 定时任务检测不可用服务是否可用，可以则更新到可用Map中
     */
    public static void scheduleTask() {
        Runnable scheduleRunnable = () -> {
            try {
                if (errorMap.isEmpty()) {
                    return;
                } else {
                    errorMap.forEach((key, servers) -> {
                        Socket socket = new Socket();
                        servers.forEach(server -> {
                            SocketAddress address = new InetSocketAddress(server.getHost(), server.getPort());
                            try {
                                // 连接成功则更新到可用Map中
                                socket.connect(address, 5000);
                                LoadBalancerMgr.balancerMap.get(key).getServerList(AbstractLoadBalancer.ServerGroup.ALL)
                                        .add(new Server(server.getHost(), server.getPort()));
                                errorMap.get(key).remove(server);
                            } catch (Exception e) {
                                log.error(e.getMessage(), e);
                            } finally {
                                if (socket != null) {
                                    try {
                                        socket.close();
                                    } catch (IOException e) {
                                    }
                                }
                            }
                        });
                    });
                }
            } catch (Exception e) {
            }
        };
        ScheduledThreadPoolExecutor scheduled = new ScheduledThreadPoolExecutor(1);
        scheduled.scheduleAtFixedRate(scheduleRunnable ,0, 30, TimeUnit.SECONDS);
    }

    /**
     * 发送http请求验证服务是否可用
     * @param server
     * @return
     */
    public boolean isAlive(Server server) {
        String urlStr = "http://";
        urlStr += server.getId();

        boolean isAlive = false;

        HttpClient httpClient = new DefaultHttpClient();
        HttpUriRequest getRequest = new HttpGet(urlStr);
        try {
            HttpResponse response = httpClient.execute(getRequest);
            isAlive = (response.getStatusLine().getStatusCode() == 200);
        } catch (IOException e) {
            log.error(e.getMessage(), e);
        }finally{
            // Release the connection.
            getRequest.abort();
        }

        return isAlive;
    }

    /**
     * 验证服务是否可用
     * @param server
     * @return
     */
    public boolean isTelnetAlive(Server server) {
        Socket socketClient = null;
        try {
            SocketAddress address = new InetSocketAddress(server.getHost(), server.getPort());
            socketClient = new Socket();
            socketClient.connect(address, 2000);
        } catch (Exception e) {
            LoadBalancerMgr.balancerMap.get(key).getServerList(AbstractLoadBalancer.ServerGroup.ALL).remove(server);
            if (errorMap.get(key) != null && !errorMap.get(key).isEmpty()) {
                errorMap.get(key).add(server);
            } else {
                ArrayList<Server> errorServerList = new ArrayList<>();
                errorServerList.add(server);
                errorMap.put(key, errorServerList);
            }
            return false;
        } finally {
            try {
                if (socketClient != null) {
                    socketClient.close();
                }
            }
            catch (Exception e) {
            }
        }

        return true;
    }

    public String getKey() {
        return key;
    }

    public void setKey(String key) {
        this.key = key;
    }
}