自定义策略&配置

前面介绍了 Spring Cloud LoadBalancer 提供的轮训和随机策略，显然在实际开发中不能满足业务的需要，这时我们可以选择第三方实现的策略，或者自定义策略。

下面将通过一个简单的示例来演示如何自定义负载均衡策略，步骤如下：

实现 ReactorServiceInstanceLoadBalancer 接口

实现 ReactorServiceInstanceLoadBalancer 接口的 choose(Request request) 方法，实现自己的负载均衡策略。

这里我们将基于自定义请求头（custom_route）来实现，如果传递了custom_route 请求头，则通过请求的哈希码与实例列表数量模运算，选择一个节点执行请求。如果没有传递 custom_route 请求头，则采用随机策略，随机返回一个实例。

代码如下：

package com.hxstrive.springcloud.loadbalancer_demo.demo7.loadbalancer;

import org.apache.commons.lang.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.ServiceInstanceListSupplier;
import org.springframework.util.function.SingletonSupplier;
import reactor.core.publisher.Mono;

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

/**
 * 自定义负载均衡算法
 * 判断用户是否传递 custom_route 请求头，如果传递了，且不为空，则根据请求头的值的哈希值返回服务实例
 * 如果没有传递 custom_route 请求头，或者请求为空，则返回一个随机的服务实例
 */
public class SimpleRoundRobinLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    public static final String CUSTOM_HEADER = "custom_route";
    // 服务ID（当前负载均衡的服务名称）
    private final String serviceId;
    // 服务实例供应器：用于获取可用服务实例列表
    private final SingletonSupplier<ServiceInstanceListSupplier> serviceInstanceListSingletonSupplier;

    public SimpleRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                        String serviceId) {
        this.serviceId = serviceId;
        this.serviceInstanceListSingletonSupplier = SingletonSupplier
                .of(() -> serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new));
    }

    /**
     * 核心方法：选择服务实例
     */
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        // 获取用户自定义的头
        String customHeader = "";
        Object content = request.getContext();
        if(content instanceof RequestDataContext) {
            List<String> list = ((RequestDataContext)content).getClientRequest().getHeaders().get(CUSTOM_HEADER);
            if(Objects.nonNull(list) && !list.isEmpty()) {
                customHeader = StringUtils.join(list, ",");
            }
        }

        // 从供应器获取服务实例列表（响应式操作）
        final String finalCustomHeader = customHeader;
        final ServiceInstanceListSupplier supplier = serviceInstanceListSingletonSupplier.obtain();
        return supplier.get().next()
                .map(instances -> {
                    // 无可用实例时返回空响应
                    if (instances.isEmpty()) {
                        return new EmptyResponse();
                    }

                    if(StringUtils.isEmpty(finalCustomHeader)) {
                        System.out.println("->> " + serviceId + " :: 使用随机策略");
                        // 随机返回一个
                        ThreadLocalRandom random = ThreadLocalRandom.current();
                        int index = ThreadLocalRandom.current().nextInt(instances.size());
                        return new DefaultResponse(instances.get(index));
                    } else {
                        System.out.println("->> " + serviceId + " :: 使用自定义哈希策略");
                        // 根据固定请求头信息进行哈希返回
                        int index = Math.abs(finalCustomHeader.hashCode()%instances.size());
                        return new DefaultResponse(instances.get(index));
                    }
                });
    }

}

注意：如果你不知道如何自定义策略，可以直接将 Spring Cloud LoadBalancer 内置的轮训或随机策略源码靠过来修改修改就可以了。

配置自定义策略

创建一个配置类（不要使用 @Configuration 修饰），创建一个自定义策略的 Bean，代码如下：

package com.hxstrive.springcloud.loadbalancer_demo.demo7.config;

import com.hxstrive.springcloud.loadbalancer_demo.demo7.loadbalancer.SimpleRoundRobinLoadBalancer;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ConditionalOnDiscoveryEnabled;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientConfiguration;
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.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.core.env.Environment;

/**
 * 简单轮询策略配置类
 */
@Import(LoadBalancerClientConfiguration.class) // 导入LoadBalancer默认配置
@ConditionalOnDiscoveryEnabled // 服务发现启用时才生效
public class SimpleRoundRobinConfig {

    /**
     * 注册自定义轮询负载均衡器
     */
    @Bean
    public ReactorLoadBalancer<ServiceInstance> simpleRoundRobinLoadBalancer(
            Environment environment,
            LoadBalancerClientFactory loadBalancerClientFactory) {
        // 获取当前服务ID（从环境变量中读取，由LoadBalancer自动注入）
        String serviceId = environment.getProperty(LoadBalancerClientFactory.PROPERTY_NAME);

        ObjectProvider<ServiceInstanceListSupplier> supplierProvider =
                loadBalancerClientFactory.getLazyProvider(serviceId, ServiceInstanceListSupplier.class);

        // 返回自定义负载均衡器实例
        return new SimpleRoundRobinLoadBalancer(supplierProvider, serviceId);
    }
}

为 RestTemplate 开启负载均衡

创建一个 @Configuration 配置类，使用 @LoadBalanced 注解为 RestTemplate 赋予负载均衡能力。如下：

package com.hxstrive.springcloud.loadbalancer_demo.demo7.config;

import org.springframework.cloud.client.loadbalancer.LoadBalanced;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.client.RestTemplate;

@Configuration
public class AppConfig {

    // 注入 RestTemplate，并添加 @LoadBalanced 注解启用负载均衡
    @Bean
    @LoadBalanced
    public RestTemplate restTemplate() {
        return new RestTemplate();
    }

}

spring 广告位

设置负载均衡策略

在启动类上面使用 @LoadBalancerClients 或者 @LoadBalancerClient 注解为全局或指定的服务配置负载均衡策略，如下：

package com.hxstrive.springcloud.loadbalancer_demo.demo7;

import com.hxstrive.springcloud.loadbalancer_demo.demo7.config.SimpleRoundRobinConfig;
import com.hxstrive.springcloud.loadbalancer_demo.demo7.loadbalancer.SimpleRoundRobinLoadBalancer;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.cloud.client.discovery.EnableDiscoveryClient;
import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClient;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

@RestController
@SpringBootApplication
@EnableDiscoveryClient
// 设置全局的负载均衡策略
//@LoadBalancerClients(defaultConfiguration = SimpleRoundRobinConfig.class)
@LoadBalancerClient(value = "USER-SERVICE", configuration = SimpleRoundRobinConfig.class)
public class LoadbalancerDemoApplication {

    @Autowired
    private RestTemplate restTemplate;

    public static void main(String[] args) {
        SpringApplication.run(LoadbalancerDemoApplication.class, args);
    }

    @GetMapping("/getData")
    public String getData(HttpServletRequest request) throws Exception {
        StringBuilder result = new StringBuilder();

        // 1. 创建请求头对象，并添加自定义头信息
        HttpHeaders headers = new HttpHeaders();
        String customHeader = request.getHeader(SimpleRoundRobinLoadBalancer.CUSTOM_HEADER);
        if(StringUtils.hasText(customHeader)) {
            // 自定义请求头，用来触发负载均衡
            headers.add(SimpleRoundRobinLoadBalancer.CUSTOM_HEADER, customHeader);
        }

        // 2. 封装请求头（无请求体时，第二个参数传null）
        HttpEntity<String> requestEntity = new HttpEntity<>(null, headers);

        // 关键：用服务名代替具体 IP:端口，LoadBalancer 会自动处理
        final String url = "http://USER-SERVICE/api/users";  // 服务名与提供者的 spring.application.name 一致

        // 发起10次调用，仅截取每次返回结果的钱40个字符
        for(int i=0; i<10;i++) {
            // 3. 使用exchange()发送GET请求，携带请求头
            ResponseEntity<String> response = restTemplate.exchange(
                    url,           // 请求URL
                    HttpMethod.GET, // HTTP方法（GET）
                    requestEntity,  // 封装了请求头的对象
                    String.class    // 响应数据类型
            );

            // 4. 获取响应结果
            String value = response.getBody();
            if(StringUtils.hasText(value)) {
                result.append(value.replaceAll("\n", "")
                        .substring(0, Math.min(40, value.length())) + "...").append("<br/>");
            }
        }

        return result.toString();
    }

}

注意，上面在 getData() 中将发起 10 次对服务的调用，每次调用都会进行负载均衡，方便我们观察负载均衡的规律，是否是我们预期的规律。

运行&验证

到这里就可以直接启动应用，通过浏览器或者类似 ApiFox 的工具来进行验证，如下：

（1）使用浏览器访问 http://localhost:8080/getData, 没有传递自定义的请求头，则按照随机策略进行分配，如下图：

（2）使用 ApiFox 工具，自定义 cuteom_route 请求头，负载均衡策略将根据请求头的哈希值进行负载选择，如下图：

注意，上图中所有的服务全部调用在一个服务上，这也是基于哈希负载均衡的特点。