RestTemplate自定义负载均衡-可扩展自定义策略

前言

在现今微服务的天下中，负载在如今开发过程中是无法避免的。现在比较流行的微服务框架中基本都集成了负载模块，比较流行的有Ribbon，以及后来的spring-cloud-loadbalancer。这些组件我们在日常开发过程中可以很简单的集成在自己的项目中，但是这些都需要服务注册。之前在开发项目时，涉及了多种开发语言协同配合及需要的负载策略开源的暂不支持，最终走向了自发的套路。

RestTemplate简介

RestTemplate 是 Spring Resources 中一个访问第三方 RESTful API 接口的网络请求框架。 RestTemplate 的设计原则和其他 Spring Template （例如 JdbcTemplate、 RedisTemplate）类似，都 是为执行复杂任务提供了一个具有默认行为的简单方法。对于第三方请求提供的API非常齐全，在Ribbion和CloudLoadBalancer中内部请求都使用了RestTemplate。

RestTemplate执行流程

RestTemplate的API请求在最后都会执行doExecute方法，此方法的主要执行逻辑可以观看try这块，主要为构建一个ClientHttpRequest，然后执行execute方法获取返回。因此我们需要看下如何创建Reuest的，以及ClientHttpRequest做了什么？

@Nullable
    protected <T> T doExecute(URI url, @Nullable HttpMethod method, @Nullable RequestCallback requestCallback, @Nullable ResponseExtractor<T> responseExtractor) throws RestClientException {
        Assert.notNull(url, "URI is required");
        Assert.notNull(method, "HttpMethod is required");
        ClientHttpResponse response = null;

        Object var14;
        try {
            ClientHttpRequest request = this.createRequest(url, method);
            if (requestCallback != null) {
                requestCallback.doWithRequest(request);
            }

            response = request.execute();
            this.handleResponse(url, method, response);
            var14 = responseExtractor != null ? responseExtractor.extractData(response) : null;
        } catch (IOException var12) {
            String resource = url.toString();
            String query = url.getRawQuery();
            resource = query != null ? resource.substring(0, resource.indexOf(63)) : resource;
            throw new ResourceAccessException("I/O error on " + method.name() + " request for \"" + resource + "\": " + var12.getMessage(), var12);
        } finally {
            if (response != null) {
                response.close();
            }

        }

        return var14;
    }

CilentHttpRequest介绍

public interface ClientHttpRequest extends HttpRequest, HttpOutputMessage {
    ClientHttpResponse execute() throws IOException;
}

这是一个接口通过实现可以获取请求结果，具体的实现类通过查找自行查看吧。我们主要说下InterceptingClientHttpRequest，代码如下，因为RestTemplate其实是继承了InterceptingHttpAccessor，因此在获取Client创建工厂时会根据有无拦截进行判断，如下

public ClientHttpRequestFactory getRequestFactory() {
        List<ClientHttpRequestInterceptor> interceptors = this.getInterceptors();
        if (!CollectionUtils.isEmpty(interceptors)) {
            ClientHttpRequestFactory factory = this.interceptingRequestFactory;
            if (factory == null) {
                factory = new InterceptingClientHttpRequestFactory(super.getRequestFactory(), interceptors);
                this.interceptingRequestFactory = (ClientHttpRequestFactory)factory;
            }

            return (ClientHttpRequestFactory)factory;
        } else {
            return super.getRequestFactory();
        }
    }

若没有拦截则会执行默认的

因此当我们在RestTemplate中定义了拦截后ClientHttpRequest最终会实现此类，这里就会执行我们的拦截器，在拦截器里我们就可以做很多事情了包括负载。

public ClientHttpResponse execute(HttpRequest request, byte[] body) throws IOException {
            if (this.iterator.hasNext()) {
                ClientHttpRequestInterceptor nextInterceptor = (ClientHttpRequestInterceptor)this.iterator.next();
                return nextInterceptor.intercept(request, body, this);
            } else {
                HttpMethod method = request.getMethod();
                Assert.state(method != null, "No standard HTTP method");
                ClientHttpRequest delegate = InterceptingClientHttpRequest.this.requestFactory.createRequest(request.getURI(), method);
                request.getHeaders().forEach((key, value) -> {
                    delegate.getHeaders().addAll(key, value);
                });
                if (body.length > 0) {
                    if (delegate instanceof StreamingHttpOutputMessage) {
                        StreamingHttpOutputMessage streamingOutputMessage = (StreamingHttpOutputMessage)delegate;
                        streamingOutputMessage.setBody((outputStream) -> {
                            StreamUtils.copy(body, outputStream);
                        });
                    } else {
                        StreamUtils.copy(body, delegate.getBody());
                    }
                }

                return delegate.execute();
            }
        }

基于RestTemplate自定义负载实现

首先自定义一个@LoadBalance的注解，这个注解中主要需要@Qualifier，然后实现自定义相关配置，代码如下

@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Qualifier
public @interface LoadBalance {
}

@Slf4j
@Configuration
@ConditionalOnClass(RestTemplate.class) //针对restTemplate
@EnableConfigurationProperties(LoadBalancerProperties.class) // 使注解生效必须满足配置
public class LoadBalancerAutoConfiguration {

    @Autowired(required = false)
    @LoadBalance
    private final List<RestTemplate> restTemplates = Collections.emptyList();

    @Autowired(required = false)
    private LoadBalancerProperties properties;


    @Bean
    public LoadBalanceClient LoadBalanceClient() throws ClassNotFoundException, InstantiationException, IllegalAccessException {
        return (LoadBalanceClient) Class.forName(properties.getType()).newInstance();
    }

    @Bean
    public LoadBalanceRequestFactory factory(LoadBalanceClient client) {
        return new LoadBalanceRequestFactory(client);
    }

    @Bean
    public RestRequestInterceptor restRequestInterceptor(LoadBalanceClient client, LoadBalanceRequestFactory factory) {
        return new RestRequestInterceptor(client, factory);
    }

    @Bean
    public SmartInitializingSingleton lbInitializing(final RestRequestInterceptor restRequestInterceptor){
        return () -> {
            for(RestTemplate restTemplate : restTemplates){
                restTemplate.getInterceptors().add(restRequestInterceptor);
            }
        };
    }
}

之后就是我们自定义的拦截器

public class RestRequestInterceptor implements ClientHttpRequestInterceptor {
    private final LoadBalanceClient client;
    private final LoadBalanceRequestFactory factory;


    public RestRequestInterceptor(LoadBalanceClient client, LoadBalanceRequestFactory factory) {
        this.client=client;
        this.factory=factory;
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest httpRequest, byte[] body, ClientHttpRequestExecution execution) throws IOException {
        ClientHttpResponse resp = this.client.execute(this.factory.createRequest(execution, httpRequest), body);
        return resp;
    }
}

以及执行策略

public interface LoadBalanceClient {
    ServiceInstance choose(String serviceId);

    ClientHttpResponse execute(LoadBalanceHttpRequest request, byte[] body) throws IOException;
}


public class RoundLoadBalanceClient implements LoadBalanceClient {

    @Override
    public ServiceInstance choose(String serviceId) {
        return null;
    }

    @Override
    public ClientHttpResponse execute(LoadBalanceHttpRequest request, byte[] body) throws IOException {
        log.info("轮询执行策略");
        return request.getExecution().execute(request.getRequest(), body);
    }
}


@Slf4j
public class RandomLoadBalanceClient implements LoadBalanceClient {

    @Override
    public ServiceInstance choose(String serviceId) {
        return null;
    }

    @Override
    public ClientHttpResponse execute(LoadBalanceHttpRequest request, byte[] body) throws IOException {
        log.info("随机执行策略");
        return request.getExecution().execute(request.getRequest(), body);
    }
}

策略配置类，这里可以配置自定义的策略，逻辑自定如随机，轮询等，当然此类还可以拆解，目前实现时还未拆分。

@Data
@ConfigurationProperties(prefix = "spring.load.balancer")
public class LoadBalancerProperties {
    private String type;
}

工厂创建类

public class LoadBalanceRequestFactory {
    private final LoadBalanceClient client;
    public LoadBalanceRequestFactory(LoadBalanceClient client) {
        this.client=client;
    }

    public LoadBalanceHttpRequest createRequest(ClientHttpRequestExecution execution, HttpRequest httpRequest) {
        try {
            ClientHttpRequestFactory requestFactory = new OkHttp3ClientHttpRequestFactory();
            ServiceInstance server = this.client.choose(httpRequest.getURI().getHost());
            return new LoadBalanceHttpRequest(requestFactory.createRequest(new URI("http://127.0.0.1:9999/api"), HttpMethod.GET),execution);
        } catch (Exception e) {
            return null;
        }
    }
}

整体项目结构

测试

更换配置后

再通过此地方访问百度

结语

不足之处望各位见谅，自知本篇还有不足之处尚待完善，感谢！