diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java index 167e8ab2a..fd8ac3c84 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java @@ -30,7 +30,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.autoconfigure.security.oauth2.resource.UserInfoRestTemplateCustomizer; -import org.springframework.boot.web.client.RestTemplateBuilder; +import org.springframework.boot.web.client.RestTemplateCustomizer; import org.springframework.cloud.sleuth.ErrorParser; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.instrument.web.HttpSpanInjector; @@ -38,6 +38,8 @@ import org.springframework.cloud.sleuth.instrument.web.HttpTraceKeysInjector; import org.springframework.cloud.sleuth.instrument.web.TraceWebAutoConfiguration; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.core.Ordered; +import org.springframework.core.annotation.Order; import org.springframework.http.client.ClientHttpRequestInterceptor; import org.springframework.security.oauth2.client.OAuth2RestTemplate; import org.springframework.web.client.RestTemplate; @@ -66,33 +68,6 @@ public class TraceWebClientAutoConfiguration { httpTraceKeysInjector, errorParser); } - @Bean - public BeanPostProcessor traceRestTemplateBuilderBPP(BeanFactory beanFactory) { - return new TraceRestTemplateBuilderBPP(beanFactory); - } - - private static class TraceRestTemplateBuilderBPP implements BeanPostProcessor { - private final BeanFactory beanFactory; - - private TraceRestTemplateBuilderBPP(BeanFactory beanFactory) { - this.beanFactory = beanFactory; - } - - @Override public Object postProcessBeforeInitialization(Object o, String s) - throws BeansException { - return o; - } - - @Override public Object postProcessAfterInitialization(Object o, String s) - throws BeansException { - if (o instanceof RestTemplateBuilder) { - RestTemplateBuilder builder = (RestTemplateBuilder) o; - return builder.additionalInterceptors(this.beanFactory.getBean(TraceRestTemplateInterceptor.class)); - } - return o; - } - } - @Configuration protected static class TraceInterceptorConfiguration { @@ -102,6 +77,17 @@ public class TraceWebClientAutoConfiguration { @Autowired private TraceRestTemplateInterceptor traceRestTemplateInterceptor; + @Bean + @Order(Ordered.HIGHEST_PRECEDENCE) + RestTemplateCustomizer traceRestTemplateCustomizer() { + return new RestTemplateCustomizer() { + @Override public void customize(RestTemplate restTemplate) { + new RestTemplateInterceptorInjector(TraceInterceptorConfiguration.this.traceRestTemplateInterceptor) + .inject(restTemplate); + } + }; + } + @PostConstruct public void init() { if (this.restTemplates != null) { @@ -166,11 +152,24 @@ class RestTemplateInterceptorInjector { } void inject(RestTemplate restTemplate) { + if (hasTraceInterceptor(restTemplate)) { + return; + } List interceptors = new ArrayList( restTemplate.getInterceptors()); - interceptors.add(this.interceptor); + interceptors.add(0, this.interceptor); restTemplate.setInterceptors(interceptors); } + + private boolean hasTraceInterceptor(RestTemplate restTemplate) { + for (ClientHttpRequestInterceptor interceptor : restTemplate + .getInterceptors()) { + if (interceptor instanceof TraceRestTemplateInterceptor) { + return true; + } + } + return false; + } } class TraceUserInfoRestTemplateCustomizer implements UserInfoRestTemplateCustomizer { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java index d757f5554..ef3a6fc43 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java @@ -44,6 +44,7 @@ import org.springframework.boot.autoconfigure.web.ServerProperties; import org.springframework.boot.context.embedded.LocalServerPort; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.web.client.RestTemplateBuilder; +import org.springframework.boot.web.client.RestTemplateCustomizer; import org.springframework.cloud.client.loadbalancer.LoadBalanced; import org.springframework.cloud.netflix.feign.EnableFeignClients; import org.springframework.cloud.netflix.feign.FeignClient; @@ -54,6 +55,7 @@ import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.SpanReporter; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.assertions.ListOfSpans; +import org.springframework.cloud.sleuth.instrument.web.client.TraceRestTemplateInterceptor; import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.cloud.sleuth.trace.TestSpanContextHolder; import org.springframework.cloud.sleuth.util.ArrayListSpanAccumulator; @@ -103,6 +105,7 @@ public class WebClientTests { @Autowired RestTemplateBuilder restTemplateBuilder; @LocalServerPort int port; @Autowired FooController fooController; + @Autowired MyRestTemplateCustomizer customizer; @After public void close() { @@ -286,6 +289,7 @@ public class WebClientTests { Span spanInController = this.fooController.getSpan(); BDDAssertions.then(spanInController).isNotNull(); then(spanInController.getTraceId()).isEqualTo(span.getTraceId()); + then(this.customizer.isExecuted()).isTrue(); } finally { this.tracer.close(span); } @@ -352,6 +356,25 @@ public class WebClientTests { SpanReporter spanReporter() { return new ArrayListSpanAccumulator(); } + + @Bean + RestTemplateCustomizer myRestTemplateCustomizer() { + return new MyRestTemplateCustomizer(); + } + } + + static class MyRestTemplateCustomizer implements RestTemplateCustomizer { + boolean executed; + + @Override public void customize(RestTemplate restTemplate) { + this.executed = true; + then(restTemplate.getInterceptors().get(0)).isInstanceOf( + TraceRestTemplateInterceptor.class); + } + + public boolean isExecuted() { + return executed; + } } public static class TestErrorController extends BasicErrorController {