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 6b4fa85a1..b9f06b981 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 @@ -75,25 +75,38 @@ final class TraceWebClientBeanPostProcessor implements BeanPostProcessor { throws BeansException { if (bean instanceof WebClient) { WebClient webClient = (WebClient) bean; - return webClient.mutate() - .filters(addTraceExchangeFilterFunctionIfNotPresent()).build(); + return wrapBuilder(webClient.mutate()).build(); } else if (bean instanceof WebClient.Builder) { WebClient.Builder webClientBuilder = (WebClient.Builder) bean; - return webClientBuilder.filters(addTraceExchangeFilterFunctionIfNotPresent()); + return wrapBuilder(webClientBuilder); } return bean; } + private WebClient.Builder wrapBuilder(WebClient.Builder webClientBuilder) { + return webClientBuilder.filters(addTraceExchangeFilterFunctionIfNotPresent()); + } + private Consumer> addTraceExchangeFilterFunctionIfNotPresent() { return functions -> { - if (functions.stream() - .noneMatch(f -> f instanceof TraceExchangeFilterFunction)) { + boolean noneMatch = noneMatchTraceExchangeFunction(functions); + if (noneMatch) { functions.add(new TraceExchangeFilterFunction(this.beanFactory)); } }; } + private boolean noneMatchTraceExchangeFunction( + List functions) { + for (ExchangeFilterFunction function : functions) { + if (function instanceof TraceExchangeFilterFunction) { + return false; + } + } + return true; + } + } final class TraceExchangeFilterFunction implements ExchangeFilterFunction { @@ -129,7 +142,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { HttpTracing httpTracing; - HttpClientHandler handler; + HttpClientHandler handler; TraceContext.Injector injector; @@ -145,25 +158,26 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { @Override public Mono filter(ClientRequest request, ExchangeFunction next) { - ClientRequest.Builder builder = ClientRequest.from(request); + HttpClientRequest wrapper = new HttpClientRequest(request); if (log.isDebugEnabled()) { log.debug("Instrumenting WebClient call"); } - Span span = handler().handleSend(injector(), builder, request, - tracer().nextSpan()); + Span parentSpan = tracer().currentSpan(); + Span span = handler().handleSend(wrapper); if (log.isDebugEnabled()) { log.debug("Handled send of " + span); } - - return new MonoWebClientTrace(next, builder.build(), this, span); + MonoWebClientTrace trace = new MonoWebClientTrace(next, wrapper.buildRequest(), + this, span); + tracer().withSpanInScope(parentSpan); + return trace; } @SuppressWarnings("unchecked") - HttpClientHandler handler() { + HttpClientHandler handler() { if (this.handler == null) { - this.handler = HttpClientHandler.create( - this.beanFactory.getBean(HttpTracing.class), - new TraceExchangeFilterFunction.HttpAdapter()); + this.handler = HttpClientHandler + .create(this.beanFactory.getBean(HttpTracing.class)); } return this.handler; } @@ -198,7 +212,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { final Tracer tracer; - final HttpClientHandler handler; + final HttpClientHandler handler; final TraceContext.Injector injector; @@ -238,9 +252,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { final Span span; - final Tracer.SpanInScope ws; - - final HttpClientHandler handler; + final HttpClientHandler handler; final Function, ? extends Publisher> scopePassingTransformer; @@ -265,8 +277,6 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { } this.context = context.put(CLIENT_SPAN_KEY, span); - this.ws = parent.tracer.withSpanInScope(span); - } @Override @@ -274,51 +284,68 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { this.actual.onSubscribe(new Subscription() { @Override public void request(long n) { - subscription.request(n); + try (Tracer.SpanInScope ws = tracing.tracer() + .withSpanInScope(span)) { + if (log.isTraceEnabled()) { + log.trace("Request"); + } + subscription.request(n); + } } @Override public void cancel() { - terminateSpanOnCancel(); - subscription.cancel(); + try (Tracer.SpanInScope ws = tracing.tracer() + .withSpanInScope(span)) { + if (log.isTraceEnabled()) { + log.trace("Cancel"); + } + terminateSpanOnCancel(); + subscription.cancel(); + } } }); } @Override public void onNext(ClientResponse response) { - this.done = true; - try { - // decorate response body - this.actual - .onNext(ClientResponse.from(response) - .body(response.bodyToFlux(DataBuffer.class) - .transform(this.scopePassingTransformer)) - .build()); - } - finally { - terminateSpan(response, null); + try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) { + this.done = true; + try { + // decorate response body + this.actual.onNext(ClientResponse.from(response) + .body(response.bodyToFlux(DataBuffer.class) + .transform(this.scopePassingTransformer)) + .build()); + } + finally { + terminateSpan(response, null); + } } } @Override public void onError(Throwable t) { - try { - this.actual.onError(t); - } - finally { - terminateSpan(null, t); + try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) { + try { + this.actual.onError(t); + } + finally { + terminateSpan(null, t); + } } } @Override public void onComplete() { - try { - this.actual.onComplete(); - } - finally { - if (!this.done) { - terminateSpan(null, null); + try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) { + try { + this.actual.onComplete(); + } + finally { + if (!this.done) { + terminateSpan(null, null); + } } } } @@ -328,10 +355,16 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { return this.context; } - void handleReceive(Span clientSpan, Tracer.SpanInScope ws, - ClientResponse clientResponse, Throwable throwable) { - this.handler.handleReceive(clientResponse, throwable, clientSpan); - ws.close(); + void handleReceive(Span clientSpan, ClientResponse clientResponse, + Throwable throwable) { + if (log.isTraceEnabled()) { + log.trace("Handling receive"); + } + this.handler.handleReceive(new HttpClientResponse(clientResponse), + throwable, clientSpan); + if (log.isTraceEnabled()) { + log.trace("Closed scope"); + } } void terminateSpanOnCancel() { @@ -341,7 +374,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { } this.span.tag("error", CANCELLED_SUBSCRIPTION_ERROR); - handleReceive(this.span, this.ws, null, null); + handleReceive(this.span, null, null); } void terminateSpan(@Nullable ClientResponse clientResponse, @@ -351,7 +384,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { log.debug("No response was returned. Will close the span [" + this.span + "]"); } - handleReceive(this.span, this.ws, clientResponse, throwable); + handleReceive(this.span, clientResponse, throwable); return; } int statusCode = clientResponse.rawStatusCode(); @@ -365,42 +398,77 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction { throwable = new RestClientException( "Status code of the response is [" + statusCode + "]"); } - handleReceive(this.span, this.ws, clientResponse, throwable); + handleReceive(this.span, clientResponse, throwable); } } } - static final class HttpAdapter - extends brave.http.HttpClientAdapter { + static final class HttpClientRequest extends brave.http.HttpClientRequest { - @Override - public String method(ClientRequest request) { - return request.method().name(); + private final ClientRequest delegate; + + private final ClientRequest.Builder builder; + + HttpClientRequest(ClientRequest delegate) { + this.delegate = delegate; + this.builder = ClientRequest.from(delegate); } @Override - public String url(ClientRequest request) { - return request.url().toString(); + public Object unwrap() { + return this.delegate; } @Override - public String requestHeader(ClientRequest request, String name) { - Object result = request.headers().getFirst(name); - return result != null ? result.toString() : null; + public String method() { + return this.delegate.method().name(); } @Override - public Integer statusCode(ClientResponse response) { - int result = statusCodeAsInt(response); - return result != 0 ? result : null; + public String path() { + return this.delegate.url().getPath(); } @Override - public int statusCodeAsInt(ClientResponse response) { + public String url() { + return this.delegate.url().toString(); + } + + @Override + public String header(String name) { + return this.delegate.headers().getFirst(name); + } + + @Override + public void header(String name, String value) { + this.builder.header(name, value); + } + + ClientRequest buildRequest() { + return this.builder.build(); + } + + } + + static final class HttpClientResponse extends brave.http.HttpClientResponse { + + private final ClientResponse delegate; + + HttpClientResponse(ClientResponse delegate) { + this.delegate = delegate; + } + + @Override + public Object unwrap() { + return this.delegate; + } + + @Override + public int statusCode() { try { - return response.rawStatusCode(); + return delegate.rawStatusCode(); } catch (Exception dontCare) { return 0; diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunctionHttpAdapterTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunctionHttpAdapterTests.java deleted file mode 100644 index 9295fbc6b..000000000 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunctionHttpAdapterTests.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright 2013-2019 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.cloud.sleuth.instrument.web.client; - -import org.assertj.core.api.BDDAssertions; -import org.junit.Test; -import org.mockito.BDDMockito; - -import org.springframework.web.reactive.function.client.ClientResponse; - -public class TraceExchangeFilterFunctionHttpAdapterTests { - - @Test - public void should_return_0_when_invalid_status_code_is_returned() { - ClientResponse clientResponse = BDDMockito.mock(ClientResponse.class); - BDDMockito.given(clientResponse.rawStatusCode()) - .willThrow(new IllegalStateException("Boom")); - TraceExchangeFilterFunction.HttpAdapter adapter = new TraceExchangeFilterFunction.HttpAdapter(); - - Integer statusCode = adapter.statusCodeAsInt(clientResponse); - - BDDAssertions.then(statusCode).isZero(); - } - - @Test - public void should_return_status_code_when_valid_status_code_is_returned() { - ClientResponse clientResponse = BDDMockito.mock(ClientResponse.class); - BDDMockito.given(clientResponse.rawStatusCode()).willReturn(200); - TraceExchangeFilterFunction.HttpAdapter adapter = new TraceExchangeFilterFunction.HttpAdapter(); - - Integer statusCode = adapter.statusCodeAsInt(clientResponse); - - BDDAssertions.then(statusCode).isEqualTo(200); - } - -}