Spring RestTemplate 设置每次请求的 Timeout

前言

在实现这个功能之前，我也上网搜索了一下方案。大多数的解决方法都是定义多个 RestTemplate 设置不同的超时时间。有没有更好的方式呢？带着这个问题，我们一起来深入一下 RestTemplate 的源码

提示：本文包含了大量的源码分析，如果想直接看笔者是如何实现的，直接跳到最后的改造思路

版本

SpringBoot：2.3.4.RELEASE

RestTemplate

RestTemplate#doExecute

RestTemplate 发送请求的方法，随便找一个最后都会走到上图的 doExecute。

从上图来看，这个方法做的就是这几件事

1. createRequest
2. 执行 RequestCallback
3. 执行 Request
4. 处理响应，将响应转换成用户声明的类型

RequestCallback 做了什么

1. 根据 RestTemplate 中的定义 HttpMessageConverter 填充 Header Accept（支持的响应类型）
2. 通过 HttpMessageConverter 转换 HttpBody

这里我们需要重点关注的是，createRequest 和 执行 Request 部分

createRequest

RestTemplate 中的 Request 是由 RequestFactory 完成创建。所以我们先来看下获取 RequestFactory 的逻辑

如果 RestTemplate 配置了 ClientHttpRequestInterceptor（拦截器）的话，则创建 InterceptingClientHttpRequestFactory，反之则直接获取 RequestFactory

1. 我们可以通过 RestTemplate#setInterceptors 手动添加拦截器；
2. 当使用 @LoadBalanced 标记 RestTemplate 时，RestTemplate 中也会被加入拦截器，具体原理可以参考 org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration

我们先来看下 InterceptingClientHttpRequestFactory 是什么逻辑

InterceptingClientHttpRequestFactory

createRequest 方法直接返回了 InterceptingClientHttpRequest，参考 doExecute 的逻辑，接下来会执行 InterceptingClientHttpRequest#execute，其内部会执行到 InterceptingRequestExecution#execute

这里随便找一个拦截器的实现配合着来看

逻辑梳理一下：

1. InterceptingRequestExecution 会先去执行所有的拦截器
2. 拦截器在执行完逻辑之后，再次 InterceptingRequestExecution#execute。InterceptingRequestExecution 再次调用下一个拦截器
3. 在拦截器逻辑执行完之后，会去调用真正的 RequestFactory 创建请求，然后执行请求

在阅读完 InterceptingRequestExecution#execute 的代码之后，我们可以发现。这里仅仅是将 request 的 uri，method，header，body 复制到了 delegate 中。说明拦截器只能对这些属性进行处理，并不能在拦截器层面添加 timeout 的相关处理。

默认情况的 RequestFactory

默认情况下 RestTemplate 会使用 SimpleClientHttpRequestFactory 来创建请求，我们也可以在这个类中看到 setReadTimeout 方法。但是 SimpleClientHttpRequestFactory 并没有提供可以拓展的点，只能设置一个针对所有请求的超时时间。感兴趣的同学可以自己阅读下源码，这里就不贴出来了

HttpComponentsClientHttpRequestFactory

在阅读 HttpComponentsClientHttpRequestFactory 时，发现了可以扩展的地方。每次在创建 Request 的时候，都需要在 HttpContext 这个类中设置 RequestConfig，使用过 apache http client 的同学可能知道 RequestConfig 这个类，这个类包含了大量的属性可以定义请求的行为，这其中有一个属性 socketTimeout 正是我们所需要的。

这个类中我们可以扩展的地方就在 createHttpContext 方法中

默认情况下 createHttpContext 返回 null，然后会尝试从 HttpUriRequest 和 HttpClient 中获取 RequestConfig 赋值到 HttpContext 中。

createHttpContext 这个方法我们也来看一下

@Nullable
private BiFunction<HttpMethod, URI, HttpContext> httpContextFactory;

/**
 * Configure a factory to pre-create the {@link HttpContext} for each request.
 * <p>This may be useful for example in mutual TLS authentication where a
 * different {@code RestTemplate} for each client certificate such that
 * all calls made through a given {@code RestTemplate} instance as associated
 * for the same client identity. {@link HttpClientContext#setUserToken(Object)}
 * can be used to specify a fixed user token for all requests.
 * @param httpContextFactory the context factory to use
 * @since 5.2.7
 */
public void setHttpContextFactory(BiFunction<HttpMethod, URI, HttpContext> httpContextFactory) {
	this.httpContextFactory = httpContextFactory;
}

/**
 * Template methods that creates a {@link HttpContext} for the given HTTP method and URI.
 * <p>The default implementation returns {@code null}.
 * @param httpMethod the HTTP method
 * @param uri the URI
 * @return the http context
 */
@Nullable
protected HttpContext createHttpContext(HttpMethod httpMethod, URI uri) {
	return (this.httpContextFactory != null ? this.httpContextFactory.apply(httpMethod, uri) : null);
}

至此，已经很清晰了。我们可以通过调用 setHttpContextFactory 来改变 createHttpContext 的结果。

改造思路

我们可以开始进行改造了，思路如下

1. 默认的超时时间等属性，我们可以通过 HttpComponentsClientHttpRequestFactory#setHttpClient 或者 HttpComponentsClientHttpRequestFactory#setReadTimeout 来决定
2. 在需要自定义 RequsetConfig 的场景，将 RequsetConfig 存储在 ThreadLocal 中
3. 我们自定义的 HttpContextFactory 在读取到 ThreadLocal 中的 RequsetConfig 后，会生成一个 HttpContext，其他情况返回 null（走原来的逻辑）

代码如下

public class CustomHttpContextFactory implements BiFunction<HttpMethod, URI, HttpContext> {
    @Override
    public HttpContext apply(HttpMethod httpMethod, URI uri) {
        RequestConfig requestConfig = RequestConfigHolder.get();
        if (requestConfig != null) {
            HttpContext context = HttpClientContext.create();
            context.setAttribute(HttpClientContext.REQUEST_CONFIG, requestConfig);
            return context;
        }
        return null;
    }
}

public class RequestConfigHolder {

    private static final ThreadLocal<RequestConfig> threadLocal = new ThreadLocal<>();

    public static void bind(RequestConfig requestConfig) {
        threadLocal.set(requestConfig);
    }

    public static RequestConfig get() {
        return threadLocal.get();
    }

    public static void clear() {
        threadLocal.remove();
    }
}

配置类

@Configuration
public class RestTemplateConfiguration {

    @Bean("customTimeoutRestTemplate")
    public RestTemplate customTimeout() {
        RestTemplate restTemplate = new RestTemplate();
        HttpComponentsClientHttpRequestFactory requestFactory = new HttpComponentsClientHttpRequestFactory();
        requestFactory.setHttpContextFactory(new CustomHttpContextFactory());
        requestFactory.setReadTimeout(3000);
        restTemplate.setRequestFactory(requestFactory);
        return restTemplate;
    }

}

使用案例

@GetMapping("custom/setTimeout")
public String customSetTimeout(Integer timeout) {
    RequestConfig requestConfig = RequestConfig.custom().setSocketTimeout(timeout).build();
    try {
        RequestConfigHolder.bind(requestConfig);
        customTimeoutRestTemplate.getForObject("https://www.baidu.com", String.class);
    } finally {
        RequestConfigHolder.clear();
    }
    return "OK";
}

思路就是这样，可以将这个使用方式封装为 注解 + AOP，这样用起来会更简单。

Demo

本文完整 demo：https://github.com/TavenYin/taven-springboot-learning/tree/master/springboot-restTemplate

最后

如果觉得我的文章对你有帮助，动动小手点下关注或者喜欢，你的支持是对我最大的帮助