spring cloud gateway自定义负载均衡, 优先转发到指定IP

示例为优先将请求分配到指定IP, 如果指定IP没有服务, 则随机分配.

GatewayLoadBalancerConfiguration.java

挂载配置类, 可以通过@org.springframework.boot.autoconfigure.condition.ConditionalOnXXX或者@Profile注解或者@ComponentScan过滤等方式来配置是否启用

在任意一个@Configuration注解的类上配置均可

import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
import org.springframework.context.annotation.Configuration;

@Configuration
@LoadBalancerClients(defaultConfiguration = {LyLoadBalancerClientConfiguration.class})
public class GatewayLoadBalancerConfiguration {
}

LyLoadBalancerClientConfiguration.java

配置类, 用于配置自定义负载均衡策略

import kim.nzxy.gateway.loadbalance.IpPriorityLoadBalancer;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.core.env.Environment;

/**
 * @author ly-chn
 */
public class LyLoadBalancerClientConfiguration {
    @Bean
    public ReactorLoadBalancer<ServiceInstance> customLoadBalancer(Environment environment,
                                                                   LoadBalancerClientFactory loadBalancerClientFactory) {
        String name = environment.getProperty(LoadBalancerClientFactory.PROPERTY_NAME);
        return new IpPriorityLoadBalancer(name,
            loadBalancerClientFactory.getLazyProvider(name, ServiceInstanceListSupplier.class));
    }
}

IpPriorityLoadBalancer.java

自定义负载均衡实现方式, 如果请求头中包含client-ip, 并且某个服务注册地址匹配, 则转到此服务, 否则转发到随机一个服务, 参考: org.springframework.cloud.loadbalancer.core.RandomLoadBalancer

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Objects;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 优先匹配指定系统IP
 *
 * @author ly-chn
 */
@Slf4j
@AllArgsConstructor
public class IpPriorityLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;


    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        // 建议将"client-ip"设定为常量
        List<String> clientIpList = ((RequestDataContext) request.getContext()).getClientRequest().getHeaders().get("client-ip");
        String clientIp = CollectionUtils.isNotEmpty(clientIpList) ? clientIpList.get(0) : null;
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map(serviceInstances -> processInstanceResponse(supplier, serviceInstances, clientIp));
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances, String clientIp) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(serviceInstances, clientIp);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, String clientIp) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("No servers available for service: " + serviceId);
            }
            return new EmptyResponse();
        }
        if (StringUtils.isEmpty(clientIp)) {
            return new DefaultResponse(instances.get(ThreadLocalRandom.current().nextInt(instances.size())));
        }
        for (ServiceInstance instance : instances) {
            if (Objects.equals(clientIp, instance.getHost())) {
                return new DefaultResponse(instance);
            }
        }
        return new DefaultResponse(instances.get(ThreadLocalRandom.current().nextInt(instances.size())));
    }

}

前端使用时:

GET http://网关IP:网关端口/目标资源
client-ip: 192.168.1.168

此时请求将优先转发到192.168.1.168上