From a4dbfbee5f5e6287ab1086acd9bf9a104e73dbc8 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Wed, 15 Nov 2017 11:51:24 +0100 Subject: [PATCH] MAde tests pass, fixed webclient support --- .../TraceWebClientBeanPostProcessor.java | 9 ++- .../client/integration/WebClientTests.java | 80 ++++++++++++++++--- .../MessagingApplicationTests.java | 8 +- 3 files changed, 83 insertions(+), 14 deletions(-) diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java index d996e4c64..9163c7d15 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java +++ b/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 exchange = next.exchange(builder.build()) + Mono 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; 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 89402fa3b..591eb3244 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 @@ -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 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 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 diff --git a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java b/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java index cb792ff1d..d8cedc684 100644 --- a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java +++ b/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() {