Merge branch 'master' into 2.0.x

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java

@@ -21,12 +21,15 @@ import java.util.Collection;
 import java.util.List;
 import javax.annotation.PostConstruct;
 
+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.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.web.client.RestTemplateBuilder;
 import org.springframework.cloud.sleuth.ErrorParser;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.instrument.web.HttpSpanInjector;
@@ -84,6 +87,33 @@ public class TraceWebClientAutoConfiguration {
 				}
 			}
 		}
+
+		@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;
+			}
+		}
 	}
 
 	@ConditionalOnClass(WebClient.class)

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java

@@ -34,6 +34,7 @@ import com.netflix.loadbalancer.Server;
 import junitparams.JUnitParamsRunner;
 import junitparams.Parameters;
 import org.apache.commons.logging.LogFactory;
+import org.assertj.core.api.BDDAssertions;
 import org.awaitility.Awaitility;
 import org.junit.After;
 import org.junit.ClassRule;
@@ -46,6 +47,8 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 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.server.LocalServerPort;
 import org.springframework.boot.web.servlet.error.ErrorAttributes;
 import org.springframework.cloud.client.loadbalancer.LoadBalanced;
 import org.springframework.cloud.netflix.feign.EnableFeignClients;
@@ -95,6 +98,9 @@ public class WebClientTests {
 	@Autowired ArrayListSpanAccumulator listener;
 	@Autowired Tracer tracer;
 	@Autowired TestErrorController testErrorController;
+	@Autowired RestTemplateBuilder restTemplateBuilder;
+	@LocalServerPort int port;
+	@Autowired FooController fooController;
 
 	@After
 	public void close() {
@@ -267,6 +273,23 @@ public class WebClientTests {
 		then(this.testErrorController.getSpan()).isNull();
 	}
 
+	@Test
+	public void should_wrap_rest_template_builders() {
+		Span span = this.tracer.createSpan("foo");
+		try {
+			RestTemplate template = this.restTemplateBuilder.build();
+
+			template.getForObject("http://localhost:" + this.port + "/traceid", String.class);
+
+			Span spanInController = this.fooController.getSpan();
+			BDDAssertions.then(spanInController).isNotNull();
+			then(spanInController.getTraceId()).isEqualTo(span.getTraceId());
+		} finally {
+			this.tracer.close(span);
+		}
+		then(this.tracer.getCurrentSpan()).isNull();
+	}
+
 	private void thenRegisteredClientSentAndReceivedEvents(Span span) {
 		then(span).hasLoggedAnEvent(Span.CLIENT_RECV);
 		then(span).hasLoggedAnEvent(Span.CLIENT_SEND);
@@ -361,6 +384,8 @@ public class WebClientTests {
 		@Autowired
 		Tracer tracer;
 
+		Span span;
+
 		@RequestMapping(value = "/notrace", method = RequestMethod.GET)
 		public String notrace(
 				@RequestHeader(name = Span.TRACE_ID_NAME, required = false) String traceId) {
@@ -375,6 +400,7 @@ public class WebClientTests {
 			then(traceId).isNotEmpty();
 			then(parentId).isNotEmpty();
 			then(spanId).isNotEmpty();
+			this.span = this.tracer.getCurrentSpan();
 			return traceId;
 		}
 
@@ -399,6 +425,10 @@ public class WebClientTests {
 			then(parentId).isNotEmpty();
 			then(spanId).isNotEmpty();
 		}
+
+		public Span getSpan() {
+			return this.span;
+		}
 	}
 
 	@Configuration