RestTemplate: register Trace interceptor *first* in the chain; fixes #822

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<ClientHttpRequestInterceptor> interceptors = new ArrayList<ClientHttpRequestInterceptor>(
 				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 {

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 {