自定义实现restTemplate负载均衡功能

依赖

 <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
            <version>2.6.10</version>
        </dependency>

创建注解EnableLoadBalance

@Target({ElementType.METHOD, ElementType.FIELD,ElementType.ANNOTATION_TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
@Qualifier
public @interface EnableLoadBalance {
}

创建规则接口Rule

public interface Rule {
    /**
     * 获取目标url
     */
    String getTargetUrl(List<String> urls);
}

实现一个随机规则RandomRule

@Component
public class RandomRule implements Rule{
    @Override
    public String getTargetUrl(List<String> urls) {
        if (CollectionUtils.isEmpty(urls)){
            return null;
        }
        int index = ThreadLocalRandom.current().nextInt(0, urls.size());
        return urls.get(index);
    }
}

创建Request包装类ModifiedHttpRequest

public class ModifiedHttpRequest implements HttpRequest {

    private final HttpRequest originalRequest;
    private final URI modifiedUri;

    public ModifiedHttpRequest(HttpRequest originalRequest, String modifiedUrl) {
        this.originalRequest = originalRequest;
        // 保存原始请求的查询参数和路径参数
        MultiValueMap<String, String> originalParams = UriComponentsBuilder.fromUri(originalRequest.getURI()).build().getQueryParams();
        // 构建新的 URI，保留原始查询参数和路径参数
        this.modifiedUri = UriComponentsBuilder.fromUriString(modifiedUrl)
                .queryParams(originalParams)
                .build()
                .toUri();
    }

    @Override
    public String getMethodValue() {
        return originalRequest.getMethodValue();
    }

    @Override
    public URI getURI() {
        return modifiedUri;
    }

    @Override
    public HttpHeaders getHeaders() {
        return originalRequest.getHeaders();
    }
}

创建拦截器LoadBalancerInterceptor

public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
    /**
     * 服务名对应的可用服务集合
     */
    private final Map<String, List<String>> serverNameToUrlMap;
    /**
     * 负载均衡规则
     */
    private final Rule rule;

    public LoadBalancerInterceptor(Map<String, List<String>> serverNameToUrlMap, Rule rule) {
        this.serverNameToUrlMap = serverNameToUrlMap;
        this.rule = rule;
    }
    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
        URI uri = request.getURI();
        String originHost = uri.getHost();
        String fullUrl = uri.toString();
        if (serverNameToUrlMap.containsKey(originHost)) {
            List<String> urlList = serverNameToUrlMap.get(originHost);
            String host = rule.getTargetUrl(urlList);
            fullUrl = fullUrl.replace(originHost, host);
            request = new ModifiedHttpRequest(request, fullUrl);
        }
        return execution.execute(request, body);
    }
}

创建配置类RestTemplateProcessor

@Component
public class RestTemplateBeanPostProcessor_Inject {

    @EnableLoadBalance
    @Autowired
    private List<RestTemplate> restTemplates;

    @Autowired(required = false)
    private Rule rule;
    private final Map<String, List<String>> serverNameToUrlMap = new HashMap<>();

    @PostConstruct
    public void init() {
        if (Objects.isNull(rule)) {
            rule = new RandomRule();
        }
        serverNameToUrlMap.put("servername", Arrays.asList("10.21.36.211:8002", "10.21.36.211:8001"));
        restTemplates.forEach(it -> it.getInterceptors().add(new LoadBalancerInterceptor(serverNameToUrlMap, rule)));
        ResponseEntity<String> test = restTemplates.get(0).getForEntity("http://servername/test", String.class);
        System.out.println(test);
    }
}

写主类并启动测试

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


    @EnableLoadBalance
    @Bean
    RestTemplate restTemplate(){
        return new RestTemplate();
    }

    @Bean
    RestTemplate restTemplate1(){
        return new RestTemplate();
    }

    @Bean
    RestTemplate restTemplate2(){
        return new RestTemplate();
    }

    @Bean
    RestTemplate restTemplate3(){
        return new RestTemplate();
    }
    @Bean
    RestTemplate restTemplate4(){
        return new RestTemplate();
    }
    @EnableLoadBalance
    @Bean
    RestTemplate restTemplate5(){
        return new RestTemplate();
    }
}

要点：

  注解上加上@Qualifier注解，然后注入字段的时候加上这个自定义注解，则之后注入使用这个注解的bean

如下面的代码，只会注入标了EnableLoadBalance注解的resttemplate

   @EnableLoadBalance
    @Autowired
    private List<RestTemplate> restTemplates;

resttemplate执行的时候有很多拦截器，我们可以加入拦截器修改执行方法