epmet-cloud-rizhao0920/epmet-commons/epmet-commons-service-call/src/main/java/com/epmet/loadbalancer/EpmetRequestLoadBalancerRule.java

package com.epmet.loadbalancer;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.env.Environment;
import org.springframework.util.CollectionUtils;

import java.util.List;

/**
 * Epmet负载均衡算法规则类。当前支持读取环境变量，将请求重定向到指定的HOST:PORT
 *
 * @RibbonClient(name = "MS-PROXY", configuration = EpmetRequestLoadBalancerRule.class)
 *
 * @Bean
 * @LoadBalanced
 * public RestTemplate getRestTemplate() {
 *     return new RestTemplate();
 * }
 *
 * @Bean
 * public IRule getMyLoadBalancerRule(Environment env) {
 *     EpmetRequestLoadBalancerRule rule = new EpmetRequestLoadBalancerRule();
 *     rule.setEnv(env);
 *     return rule;
 * }
 *
 */
public class EpmetRequestLoadBalancerRule extends ZoneAvoidanceRule {

    private static Logger logger = LoggerFactory.getLogger(EpmetRequestLoadBalancerRule.class);

    private Environment env;

    public EpmetRequestLoadBalancerRule() {
    }

    public void setEnv(Environment env) {
        this.env = env;
    }

    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {

    }

    /**
     * 此处会取环境变量中"测试服务器"的值，
     * @param key
     * @return
     */
    @Override
    public Server choose(Object key) {
        ILoadBalancer loadBalancer = getLoadBalancer();
        List<Server> servers = loadBalancer.getReachableServers();
        if (CollectionUtils.isEmpty(servers)) {
            logger.error("自定义负载均衡器:ReachableServers列表为空");
            return super.choose(key);
        }
        Server server = servers.get(0);
        String appName = server.getMetaInfo().getAppName();
        String serviceName = appName.split("@@")[1];
        String hostEnvKey = getHostEnvKey(serviceName);
        String portEnvKey = getPortEnvKey(serviceName);
        String epmetGovOrgHost = env.getProperty(hostEnvKey);
        String epmetGovOrgPort = env.getProperty(portEnvKey);
        if (StringUtils.isAnyBlank(epmetGovOrgHost, epmetGovOrgPort)) {
            // 没有配置，走父类均衡器
            return super.choose(key);
        }

        server.setHost(epmetGovOrgHost);
        server.setPort(Integer.valueOf(epmetGovOrgPort));
        return server;
    }

    /**
     * 获取host环境变量Key
     * @param serviceName
     * @return
     */
    public String getHostEnvKey(String serviceName) {
        return serviceName.replace("-", "_").concat("_HOST").toUpperCase();
    }

    public String getPortEnvKey(String serviceName) {
        return serviceName.replace("-", "_").concat("_PORT").toUpperCase();
    }
}