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 472926757..489a7a7ee 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 @@ -23,17 +23,17 @@ import javax.annotation.PostConstruct; import brave.http.HttpTracing; import brave.spring.web.TracingClientHttpRequestInterceptor; -import org.springframework.beans.BeansException; import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.annotation.Autowired; -import org.springframework.beans.factory.config.BeanPostProcessor; import org.springframework.boot.autoconfigure.AutoConfigureAfter; import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; -import org.springframework.boot.web.client.RestTemplateBuilder; +import org.springframework.boot.web.client.RestTemplateCustomizer; import org.springframework.cloud.sleuth.instrument.web.TraceWebServletAutoConfiguration; 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.web.client.RestTemplate; import org.springframework.web.reactive.function.client.WebClient; @@ -73,40 +73,34 @@ public class TraceWebClientAutoConfiguration { public void init() { if (this.restTemplates != null) { for (RestTemplate restTemplate : this.restTemplates) { - List interceptors = new ArrayList( - restTemplate.getInterceptors()); - interceptors.add(this.clientInterceptor); - restTemplate.setInterceptors(interceptors); + new RestTemplateInterceptorInjector(this.clientInterceptor) + .inject(restTemplate); } } } } + @Autowired(required = false) + private Collection restTemplates; + + @Autowired + private TracingClientHttpRequestInterceptor traceRestTemplateInterceptor; + @Bean - public BeanPostProcessor traceRestTemplateBuilderBPP(BeanFactory beanFactory) { - return new TraceRestTemplateBuilderBPP(beanFactory); + @Order(Ordered.HIGHEST_PRECEDENCE) + RestTemplateCustomizer traceRestTemplateCustomizer() { + final TracingClientHttpRequestInterceptor interceptor = this.traceRestTemplateInterceptor; + return restTemplate -> + new RestTemplateInterceptorInjector(interceptor).inject(restTemplate); } - 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(TracingClientHttpRequestInterceptor.class)); + @PostConstruct + public void init() { + if (this.restTemplates != null) { + for (RestTemplate restTemplate : this.restTemplates) { + new RestTemplateInterceptorInjector( + this.traceRestTemplateInterceptor).inject(restTemplate); } - return o; } } } @@ -119,4 +113,32 @@ public class TraceWebClientAutoConfiguration { return new TraceWebClientBeanPostProcessor(beanFactory); } } +} + +class RestTemplateInterceptorInjector { + private final TracingClientHttpRequestInterceptor interceptor; + + RestTemplateInterceptorInjector(TracingClientHttpRequestInterceptor interceptor) { + this.interceptor = interceptor; + } + + void inject(RestTemplate restTemplate) { + if (hasTraceInterceptor(restTemplate)) { + return; + } + List interceptors = new ArrayList( + restTemplate.getInterceptors()); + interceptors.add(0, this.interceptor); + restTemplate.setInterceptors(interceptors); + } + + private boolean hasTraceInterceptor(RestTemplate restTemplate) { + for (ClientHttpRequestInterceptor interceptor : restTemplate + .getInterceptors()) { + if (interceptor instanceof TracingClientHttpRequestInterceptor) { + return true; + } + } + return false; + } } \ No newline at end of file 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 064375566..ce0dd59b5 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 @@ -16,7 +16,6 @@ package org.springframework.cloud.sleuth.instrument.web.client.integration; -import javax.servlet.http.HttpServletRequest; import java.lang.invoke.MethodHandles; import java.util.ArrayList; import java.util.Collections; @@ -26,6 +25,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import javax.servlet.http.HttpServletRequest; import brave.Span; import brave.Tracer; @@ -33,12 +33,12 @@ import brave.Tracing; import brave.propagation.SamplingFlags; import brave.propagation.TraceContextOrSamplingFlags; import brave.sampler.Sampler; +import brave.spring.web.TracingClientHttpRequestInterceptor; +import com.netflix.loadbalancer.BaseLoadBalancer; +import com.netflix.loadbalancer.ILoadBalancer; +import com.netflix.loadbalancer.Server; import junitparams.JUnitParamsRunner; import junitparams.Parameters; -import reactor.core.publisher.Hooks; -import reactor.core.scheduler.Schedulers; -import zipkin2.Annotation; -import zipkin2.reporter.Reporter; import org.apache.commons.logging.LogFactory; import org.assertj.core.api.BDDAssertions; import org.awaitility.Awaitility; @@ -55,13 +55,14 @@ import org.springframework.boot.autoconfigure.web.ServerProperties; import org.springframework.boot.autoconfigure.web.servlet.error.BasicErrorController; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.web.client.RestTemplateBuilder; +import org.springframework.boot.web.client.RestTemplateCustomizer; import org.springframework.boot.web.server.LocalServerPort; import org.springframework.boot.web.servlet.error.ErrorAttributes; -import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; import org.springframework.cloud.client.loadbalancer.LoadBalanced; import org.springframework.cloud.netflix.feign.EnableFeignClients; import org.springframework.cloud.netflix.feign.FeignClient; import org.springframework.cloud.netflix.ribbon.RibbonClient; +import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.http.HttpHeaders; @@ -77,10 +78,10 @@ import org.springframework.web.bind.annotation.RestController; import org.springframework.web.client.HttpClientErrorException; import org.springframework.web.client.RestTemplate; import org.springframework.web.reactive.function.client.WebClient; - -import com.netflix.loadbalancer.BaseLoadBalancer; -import com.netflix.loadbalancer.ILoadBalancer; -import com.netflix.loadbalancer.Server; +import reactor.core.publisher.Hooks; +import reactor.core.scheduler.Schedulers; +import zipkin2.Annotation; +import zipkin2.reporter.Reporter; import static org.assertj.core.api.Assertions.fail; import static org.assertj.core.api.BDDAssertions.then; @@ -114,6 +115,7 @@ public class WebClientTests { @Autowired RestTemplateBuilder restTemplateBuilder; @LocalServerPort int port; @Autowired FooController fooController; + @Autowired MyRestTemplateCustomizer customizer; @After public void close() { @@ -181,7 +183,7 @@ public class WebClientTests { Span span = this.tracer.nextSpan( TraceContextOrSamplingFlags.create(SamplingFlags.NOT_SAMPLED)) .name("foo").start(); - + try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { ResponseEntity> response = provider.get(this); @@ -190,7 +192,7 @@ public class WebClientTests { } finally { span.finish(); } - + then(this.reporter.getSpans()).isEmpty(); then(Tracing.current().tracer().currentSpan()).isNull(); } @@ -321,6 +323,7 @@ public class WebClientTests { span.finish(); } then(this.tracer.currentSpan()).isNull(); + then(this.customizer.isExecuted()).isTrue(); } private void assertThatSpanGotContinued(Span span) { @@ -388,6 +391,25 @@ public class WebClientTests { WebClient.Builder webClientBuilder() { return WebClient.builder(); } + + @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( + TracingClientHttpRequestInterceptor.class); + } + + public boolean isExecuted() { + return executed; + } } public static class TestErrorController extends BasicErrorController {