MAde tests pass, fixed webclient support

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

@@ -83,7 +83,8 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 			ExchangeFunction next) {
 		final ClientRequest.Builder builder = ClientRequest.from(request);
 
-		Mono<ClientResponse> exchange = next.exchange(builder.build())
+		Mono<ClientResponse> exchange = Mono
+				.defer(() -> next.exchange(builder.build()))
 				.cast(Object.class)
 				.onErrorResume(Mono::just)
 				.zipWith(Mono.subscriberContext())
@@ -124,10 +125,11 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 					}
 					Span parent = c.getOrDefault(Span.class, null);
 					Span clientSpan = createNewSpan(request, parent);
+					tracer().continueSpan(clientSpan);
 
 					httpSpanInjector().inject(clientSpan, new ClientRequestTextMap(request, builder));
 					if (log.isDebugEnabled()) {
-						log.debug("Headers got injected to the client span " + clientSpan);
+						log.debug("Headers got injected from the client span " + clientSpan);
 					}
 
 					if (parent == null) {
@@ -136,6 +138,9 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 							log.debug("Reactor Context got injected with the client span " + clientSpan);
 						}
 					}
+					if (clientSpan != null && clientSpan.equals(tracer().getCurrentSpan())) {
+						tracer().continueSpan(tracer().detach(clientSpan));
+					}
 					return c.put(CLIENT_SPAN_KEY, clientSpan);
 				});
 		return exchange;

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

@@ -37,6 +37,8 @@ import org.apache.commons.logging.LogFactory;
 import org.assertj.core.api.BDDAssertions;
 import org.awaitility.Awaitility;
 import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
@@ -60,9 +62,11 @@ 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.reactor.TraceReactorAutoConfiguration;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
 import org.springframework.cloud.sleuth.util.ArrayListSpanAccumulator;
+import org.springframework.cloud.sleuth.util.ExceptionUtils;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.HttpHeaders;
@@ -77,6 +81,9 @@ import org.springframework.web.bind.annotation.RequestMethod;
 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 reactor.core.publisher.Hooks;
+import reactor.core.scheduler.Schedulers;
 
 import static org.assertj.core.api.Assertions.fail;
 import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
@@ -84,7 +91,9 @@ import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
 @RunWith(JUnitParamsRunner.class)
 @SpringBootTest(classes = WebClientTests.TestConfiguration.class,
 		webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
-@TestPropertySource(properties = { "spring.application.name=fooservice", "feign.hystrix.enabled=false" })
+@TestPropertySource(properties = {
+		"spring.application.name=fooservice",
+		"feign.hystrix.enabled=false" })
 @DirtiesContext
 public class WebClientTests {
 
@@ -95,6 +104,8 @@ public class WebClientTests {
 
 	@Autowired TestFeignInterface testFeignInterface;
 	@Autowired @LoadBalanced RestTemplate template;
+	@Autowired WebClient webClient;
+	@Autowired WebClient.Builder webClientBuilder;
 	@Autowired ArrayListSpanAccumulator listener;
 	@Autowired Tracer tracer;
 	@Autowired TestErrorController testErrorController;
@@ -104,9 +115,17 @@ public class WebClientTests {
 
 	@After
 	public void close() {
+		ExceptionUtils.setFail(true);
 		TestSpanContextHolder.removeCurrentSpan();
 		this.listener.getSpans().clear();
 		this.testErrorController.clear();
+		this.fooController.clear();
+	}
+
+	@BeforeClass
+	public static void cleanup() {
+		Hooks.resetOnLastOperator();
+		Schedulers.resetFactory();
 	}
 
 	@Test
@@ -181,11 +200,7 @@ public class WebClientTests {
 	@SuppressWarnings("unchecked")
 	public void shouldAttachTraceIdWhenCallingAnotherService(
 			ResponseEntityProvider provider) {
-		Long currentTraceId = 1L;
-		Long currentParentId = 2L;
-		Long currentSpanId = 100L;
-		this.tracer.continueSpan(Span.builder().traceId(currentTraceId)
-				.spanId(currentSpanId).parent(currentParentId).build());
+		spanContinued();
 
 		ResponseEntity<String> response = provider.get(this);
 
@@ -196,6 +211,33 @@ public class WebClientTests {
 		thenRegisteredClientSentAndReceivedEvents(spanWithClientEvents());
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void shouldAttachTraceIdWhenCallingAnotherServiceViaWebClient() {
+		Span span = this.tracer.createSpan("foo");
+		try {
+			this.webClient.get()
+					.uri("http://localhost:" + this.port + "/traceid")
+					.retrieve()
+					.bodyToMono(String.class)
+					.block();
+
+			assertThatSpanGotContinued(span);
+		} finally {
+			this.tracer.close(span);
+		}
+		then(this.tracer.getCurrentSpan()).isNull();
+		thenRegisteredClientSentAndReceivedEvents(spanWithClientEvents());
+	}
+
+	private void spanContinued() {
+		Long currentTraceId = 1L;
+		Long currentParentId = 2L;
+		Long currentSpanId = 100L;
+		this.tracer.continueSpan(Span.builder().traceId(currentTraceId)
+				.spanId(currentSpanId).parent(currentParentId).build());
+	}
+
 	private Span spanWithClientEvents() {
 		List<Span> spans = new ArrayList<>(this.listener.getSpans());
 		for(Span span : spans) {
@@ -281,15 +323,19 @@ public class WebClientTests {
 
 			template.getForObject("http://localhost:" + this.port + "/traceid", String.class);
 
-			Span spanInController = this.fooController.getSpan();
-			BDDAssertions.then(spanInController).isNotNull();
-			then(spanInController.getTraceId()).isEqualTo(span.getTraceId());
+			assertThatSpanGotContinued(span);
 		} finally {
 			this.tracer.close(span);
 		}
 		then(this.tracer.getCurrentSpan()).isNull();
 	}
 
+	private void assertThatSpanGotContinued(Span span) {
+		Span spanInController = this.fooController.getSpan();
+		BDDAssertions.then(spanInController).isNotNull();
+		then(spanInController.getTraceId()).isEqualTo(span.getTraceId());
+	}
+
 	private void thenRegisteredClientSentAndReceivedEvents(Span span) {
 		then(span).hasLoggedAnEvent(Span.CLIENT_RECV);
 		then(span).hasLoggedAnEvent(Span.CLIENT_SEND);
@@ -320,7 +366,7 @@ public class WebClientTests {
 	}
 
 	@Configuration
-	@EnableAutoConfiguration
+	@EnableAutoConfiguration(exclude = TraceReactorAutoConfiguration.class)
 	@EnableFeignClients
 	@RibbonClient(value = "fooservice", configuration = SimpleRibbonClientConfiguration.class)
 	public static class TestConfiguration {
@@ -350,6 +396,16 @@ public class WebClientTests {
 		SpanReporter spanReporter() {
 			return new ArrayListSpanAccumulator();
 		}
+
+		@Bean
+		WebClient webClient() {
+			return WebClient.builder().build();
+		}
+
+		@Bean
+		WebClient.Builder webClientBuilder() {
+			return WebClient.builder();
+		}
 	}
 
 	public static class TestErrorController extends BasicErrorController {
@@ -429,6 +485,10 @@ public class WebClientTests {
 		public Span getSpan() {
 			return this.span;
 		}
+
+		public void clear() {
+			this.span = null;
+		}
 	}
 
 	@Configuration

spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java

@@ -111,8 +111,12 @@ public class MessagingApplicationTests extends AbstractIntegrationTest {
 				.stream()
 				.map(Span::toString)
 		.collect(Collectors.joining("\n")) + "\n]");
-		then(this.integrationTestSpanCollector.hashedSpans.stream()
-				.allMatch(span -> span.traceId().equals(traceIdHex))).describedAs("All spans have same trace id").isTrue();
+		then(this.integrationTestSpanCollector.hashedSpans
+				.stream()
+				.filter(span -> !span.traceId().equals(traceIdHex))
+				.collect(Collectors.toList()))
+				.describedAs("All spans have same trace id [" + traceIdHex + "]")
+				.isEmpty();
 	}
 
 	private void thenTheSpansHaveProperParentStructure() {