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 72c19ab33..342761488 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 @@ -19,7 +19,6 @@ package org.springframework.cloud.sleuth.instrument.web.client; import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -31,7 +30,6 @@ import brave.httpasyncclient.TracingHttpAsyncClientBuilder; import brave.httpclient.TracingHttpClientBuilder; import brave.propagation.Propagation; import brave.propagation.TraceContext; -import brave.propagation.TraceContextOrSamplingFlags; import brave.spring.web.TracingClientHttpRequestInterceptor; import io.netty.handler.codec.http.HttpHeaders; import io.netty.handler.codec.http.HttpMethod; @@ -259,7 +257,14 @@ class NettyAspect { public Object wrapHttpClientRequestSending(ProceedingJoinPoint pjp, HttpMethod method, String url, Function> handler) throws Throwable { - return this.instrumentation.wrapHttpClientRequestSending(pjp, method, url, handler); + return Mono.defer(() -> { + try { + return this.instrumentation.wrapHttpClientRequestSending(pjp, method, url, handler); + } + catch (Throwable e) { + return Mono.error(e); + } + }); } } @@ -304,27 +309,17 @@ class TracingHttpClientInstrumentation { this.httpTracing = httpTracing; } - Object wrapHttpClientRequestSending(ProceedingJoinPoint pjp, + Mono wrapHttpClientRequestSending(ProceedingJoinPoint pjp, HttpMethod method, String url, Function> handler) throws Throwable { // add headers and set CS final Span currentSpan = this.tracer.currentSpan(); final AtomicReference span = new AtomicReference<>(); - final AtomicBoolean requestAlreadyInstrumented = new AtomicBoolean(); Function> combinedFunction = req -> { try (Tracer.SpanInScope spanInScope = this.tracer.withSpanInScope(currentSpan)) { io.netty.handler.codec.http.HttpHeaders headers = req .requestHeaders(); - TraceContextOrSamplingFlags flags = this.httpTracing.tracing() - .propagation().extractor(GETTER).extract(headers); - if (flags != TraceContextOrSamplingFlags.EMPTY) { - requestAlreadyInstrumented.set(true); - if (log.isDebugEnabled()) { - log.debug("Request already instrumented. Skipping"); - } - return handle(handler, req); - } span.set(this.handler.handleSend(this.injector, headers, req)); try (Tracer.SpanInScope clientInScope = this.tracer.withSpanInScope(span.get())) { if (log.isDebugEnabled()) { @@ -339,12 +334,6 @@ class TracingHttpClientInstrumentation { (Mono) pjp.proceed(new Object[] { method , url, combinedFunction }); // get response return responseMono.doOnSuccessOrError((httpClientResponse, throwable) -> { - if (requestAlreadyInstrumented.get()) { - if (log.isDebugEnabled()) { - log.debug("Request already instrumented. Skipping"); - return; - } - } try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span.get())) { // status codes and CR this.handler.handleReceive(httpClientResponse, throwable, span.get()); 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 0642fc54e..09eff9014 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 @@ -25,6 +25,7 @@ import brave.http.HttpClientHandler; import brave.http.HttpTracing; import brave.propagation.Propagation; import brave.propagation.TraceContext; +import brave.propagation.TraceContextOrSamplingFlags; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.springframework.beans.BeansException; @@ -89,6 +90,7 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { private static final Log log = LogFactory.getLog( TraceExchangeFilterFunction.class); private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan"; + private static final String CLIENT_SPAN_ALREADY_PROCESSED_KEY = "sleuth.webclient.clientSpanAlreadyProcessed"; static final Propagation.Setter SETTER = new Propagation.Setter() { @@ -101,12 +103,23 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { } }; + static final Propagation.Getter GETTER = new Propagation.Getter() { + @Override public String get(ClientRequest carrier, String key) { + return carrier.headers().getFirst(key); + } + + @Override public String toString() { + return "HttpHeaders::getFirst"; + } + }; + public static ExchangeFilterFunction create(BeanFactory beanFactory) { return new TraceExchangeFilterFunction(beanFactory); } final BeanFactory beanFactory; Tracer tracer; + HttpTracing httpTracing; HttpClientHandler handler; TraceContext.Injector injector; @@ -125,6 +138,7 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { .flatMap(anyAndContext -> { Object any = anyAndContext.getT1(); Span clientSpan = anyAndContext.getT2().get(CLIENT_SPAN_KEY); + boolean clientSpanAlreadyProcessed = anyAndContext.getT2().get(CLIENT_SPAN_ALREADY_PROCESSED_KEY); Mono continuation; final Tracer.SpanInScope ws = tracer().withSpanInScope(clientSpan); if (any instanceof Throwable) { @@ -134,6 +148,13 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { } return continuation.doAfterSuccessOrError( (clientResponse, throwable1) -> { + if (clientSpanAlreadyProcessed) { + if (log.isDebugEnabled()) { + log.debug("Another component will process the response. Skipping"); + ws.close(); + return; + } + } Throwable throwable = throwable1; boolean error = clientResponse.statusCode().is4xxClientError() || clientResponse.statusCode().is5xxServerError(); @@ -156,6 +177,14 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { if (log.isDebugEnabled()) { log.debug("Instrumenting WebClient call"); } + TraceContextOrSamplingFlags flags = httpTracing().tracing() + .propagation().extractor(GETTER).extract(request); + if (flags != TraceContextOrSamplingFlags.EMPTY) { + if (log.isDebugEnabled()) { + log.debug("The request was already instrumented. Will not do it again"); + } + return c.put(CLIENT_SPAN_ALREADY_PROCESSED_KEY, true); + } Span parent = c.getOrDefault(Span.class, null); Span clientSpan = handler().handleSend(injector(), builder, request, tracer().nextSpan()); @@ -168,7 +197,8 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { log.debug("Reactor Context got injected with the client span " + clientSpan); } } - return c.put(CLIENT_SPAN_KEY, clientSpan); + return c.put(CLIENT_SPAN_ALREADY_PROCESSED_KEY, false) + .put(CLIENT_SPAN_KEY, clientSpan); }); return exchange; } @@ -184,11 +214,18 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { Tracer tracer() { if (this.tracer == null) { - this.tracer = this.beanFactory.getBean(HttpTracing.class).tracing().tracer(); + this.tracer = httpTracing().tracing().tracer(); } return this.tracer; } + HttpTracing httpTracing() { + if (this.httpTracing == null) { + this.httpTracing = this.beanFactory.getBean(HttpTracing.class); + } + return this.httpTracing; + } + TraceContext.Injector injector() { if (this.injector == null) { this.injector = this.beanFactory.getBean(HttpTracing.class) 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 c16c155ec..568effa75 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 @@ -69,6 +69,7 @@ import org.springframework.cloud.client.loadbalancer.LoadBalanced; import org.springframework.cloud.netflix.ribbon.RibbonClient; import org.springframework.cloud.openfeign.EnableFeignClients; import org.springframework.cloud.openfeign.FeignClient; +import org.springframework.cloud.sleuth.instrument.web.TraceWebServletAutoConfiguration; import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -248,15 +249,16 @@ public class WebClientTests { .get("http://localhost:" + port).block(); then(response).isNotNull(); - } finally { - span.finish(); } then(this.tracer.currentSpan()).isNull(); - then(this.reporter.getSpans()).isNotEmpty(); then(this.reporter.getSpans()) + .isNotEmpty() .extracting("traceId", String.class) .containsOnly(span.context().traceIdString()); + then(this.reporter.getSpans()) + .extracting("kind.name") + .contains("CLIENT"); } @Test @@ -270,15 +272,16 @@ public class WebClientTests { new BasicResponseHandler()); then(response).isNotEmpty(); - } finally { - span.finish(); } then(this.tracer.currentSpan()).isNull(); - then(this.reporter.getSpans()).isNotEmpty(); then(this.reporter.getSpans()) + .isNotEmpty() .extracting("traceId", String.class) .containsOnly(span.context().traceIdString()); + then(this.reporter.getSpans()) + .extracting("kind.name") + .contains("CLIENT"); } @Test @@ -306,15 +309,17 @@ public class WebClientTests { }); then(future.get()).isNotNull(); } finally { - span.finish(); client.close(); } then(this.tracer.currentSpan()).isNull(); - then(this.reporter.getSpans()).isNotEmpty(); then(this.reporter.getSpans()) + .isNotEmpty() .extracting("traceId", String.class) .containsOnly(span.context().traceIdString()); + then(this.reporter.getSpans()) + .extracting("kind.name") + .contains("CLIENT"); } @Test @@ -328,13 +333,14 @@ public class WebClientTests { .retrieve() .bodyToMono(String.class) .block(); - - assertThatSpanGotContinued(span); } finally { span.finish(); } then(this.tracer.currentSpan()).isNull(); - then(this.reporter.getSpans()).isNotEmpty(); + then(this.reporter.getSpans()) + .isNotEmpty() + .extracting("kind.name") + .contains("CLIENT"); } Object[] parametersForShouldAttachTraceIdWhenCallingAnotherService() { @@ -389,7 +395,11 @@ public class WebClientTests { log.info("logs " + span.annotations()); then(initialSize).as("there are no duplicate log entries").isEqualTo(distinctSize); }); - then(this.testErrorController.getSpan()).isNotNull(); + + then(this.reporter.getSpans()) + .isNotEmpty() + .extracting("kind.name") + .contains("CLIENT"); } @Test @@ -408,19 +418,14 @@ public class WebClientTests { RestTemplate template = this.restTemplateBuilder.build(); template.getForObject("http://localhost:" + this.port + "/traceid", String.class); - - assertThatSpanGotContinued(span); } finally { span.finish(); } then(this.tracer.currentSpan()).isNull(); then(this.customizer.isExecuted()).isTrue(); - } - - private void assertThatSpanGotContinued(Span span) { - Span spanInController = this.fooController.getSpan(); - BDDAssertions.then(spanInController).isNotNull(); - then(spanInController.context().traceId()).isEqualTo(span.context().traceId()); + then(this.reporter.getSpans()) + .extracting("kind.name") + .contains("CLIENT"); } private String getHeader(ResponseEntity response, String name) { @@ -444,7 +449,7 @@ public class WebClientTests { } @Configuration - @EnableAutoConfiguration + @EnableAutoConfiguration(exclude = TraceWebServletAutoConfiguration.class) @EnableFeignClients @RibbonClient(value = "fooservice", configuration = SimpleRibbonClientConfiguration.class) public static class TestConfiguration {