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 4bfa9dbbf..c117661ae 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 @@ -16,6 +16,7 @@ package org.springframework.cloud.sleuth.instrument.web.client; +import java.util.Collections; import java.util.List; import java.util.function.Consumer; @@ -27,6 +28,8 @@ import brave.propagation.Propagation; import brave.propagation.TraceContext; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import reactor.core.publisher.Mono; + import org.springframework.beans.BeansException; import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.config.BeanPostProcessor; @@ -36,7 +39,6 @@ import org.springframework.web.reactive.function.client.ClientResponse; import org.springframework.web.reactive.function.client.ExchangeFilterFunction; import org.springframework.web.reactive.function.client.ExchangeFunction; import org.springframework.web.reactive.function.client.WebClient; -import reactor.core.publisher.Mono; /** * {@link BeanPostProcessor} to wrap a {@link WebClient} instance into its trace @@ -87,10 +89,20 @@ class TraceWebClientBeanPostProcessor implements BeanPostProcessor { class TraceExchangeFilterFunction implements ExchangeFilterFunction { + private static final Log log = LogFactory.getLog(TraceExchangeFilterFunction.class); + + private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan"; + static final Propagation.Setter SETTER = new Propagation.Setter() { @Override public void put(ClientRequest.Builder carrier, String key, String value) { - carrier.header(key, value); + carrier.headers(httpHeaders -> { + if (log.isTraceEnabled()) { + log.trace("Replacing [" + key + "] with value [" + value + "]"); + } + httpHeaders.merge(key, Collections.singletonList(value), + (oldValue, newValue) -> newValue); + }); } @Override @@ -98,21 +110,10 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { return "ClientRequest.Builder::header"; } }; - 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"; - } - }; - - private static final Log log = LogFactory.getLog(TraceExchangeFilterFunction.class); - - private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan"; + public static ExchangeFilterFunction create(BeanFactory beanFactory) { + return new TraceExchangeFilterFunction(beanFactory); + } final BeanFactory beanFactory; @@ -128,16 +129,15 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction { this.beanFactory = beanFactory; } - public static ExchangeFilterFunction create(BeanFactory beanFactory) { - return new TraceExchangeFilterFunction(beanFactory); - } - @Override public Mono filter(ClientRequest request, ExchangeFunction next) { final ClientRequest.Builder builder = ClientRequest.from(request); Mono exchange = Mono.defer(() -> next.exchange(builder.build())) .cast(Object.class).onErrorResume(Mono::just) .zipWith(Mono.subscriberContext()).flatMap(anyAndContext -> { + if (log.isDebugEnabled()) { + log.debug("Wrapping the context [" + anyAndContext + "]"); + } Object any = anyAndContext.getT1(); Span clientSpan = anyAndContext.getT2().get(CLIENT_SPAN_KEY); Mono continuation; diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH1102Tests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH1102Tests.java new file mode 100644 index 000000000..09b8f9ec7 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH1102Tests.java @@ -0,0 +1,129 @@ +/* + * Copyright 2013-2018 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 + * + * http://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 brave.ScopedSpan; +import brave.Tracer; +import brave.sampler.Sampler; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.assertj.core.api.BDDAssertions; +import org.junit.Test; +import org.junit.runner.RunWith; +import reactor.core.publisher.Mono; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.test.context.SpringBootTest.WebEnvironment; +import org.springframework.boot.web.server.LocalServerPort; +import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.test.context.junit4.SpringRunner; +import org.springframework.util.MultiValueMap; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RequestHeader; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.reactive.function.client.WebClient; +import org.springframework.web.reactive.function.client.WebClientResponseException; + +@SpringBootTest(webEnvironment = WebEnvironment.RANDOM_PORT) +@RunWith(SpringRunner.class) +public class GH1102Tests { + + @Autowired + Tracer tracer; + + @Autowired + WebClient webClient; + + @Autowired + TestRetry testRetry; + + @Autowired + ArrayListSpanReporter reporter; + + @LocalServerPort + int port; + + @Test + public void should_store_retries_as_separate_spans() throws Exception { + ScopedSpan foo = this.tracer.startScopedSpan("foo"); + try { + this.webClient.get().uri("http://localhost:" + this.port + "/test").retrieve() + .bodyToMono(String.class).retry(1).block(); + BDDAssertions.fail("should throw exception"); + } + catch (WebClientResponseException ex) { + + } + finally { + foo.finish(); + } + + BDDAssertions.then(this.testRetry.getHttpHeaders().get("x-b3-traceid")) + .hasSize(1); + } + + @EnableAutoConfiguration + @Configuration + static class WebConfig { + + @Bean + Sampler sampler() { + return Sampler.ALWAYS_SAMPLE; + } + + @Bean + ArrayListSpanReporter reporter() { + return new ArrayListSpanReporter(); + } + + @Bean + WebClient webClient() { + return WebClient.builder().build(); + } + + @Bean + TestRetry testRetry() { + return new TestRetry(); + } + + } + + @RestController + static class TestRetry { + + private static final Log log = LogFactory.getLog(TestRetry.class); + + private MultiValueMap httpHeaders; + + @GetMapping("test") + Mono test(@RequestHeader MultiValueMap map) { + this.httpHeaders = map; + log.info("Processing test. Headers [" + this.httpHeaders + "]"); + return Mono.error(new RuntimeException("BOOM!")); + } + + MultiValueMap getHttpHeaders() { + return this.httpHeaders; + } + + } + +} diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Test.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Tests.java similarity index 95% rename from spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Test.java rename to spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Tests.java index 4bc50114d..b90769394 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Test.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/GH846Tests.java @@ -31,9 +31,9 @@ import org.springframework.http.client.ClientHttpRequestInterceptor; import org.springframework.test.context.junit4.SpringRunner; import org.springframework.web.client.RestTemplate; -@SpringBootTest(classes = GH846Test.App.class, webEnvironment = WebEnvironment.NONE) +@SpringBootTest(classes = GH846Tests.App.class, webEnvironment = WebEnvironment.NONE) @RunWith(SpringRunner.class) -public class GH846Test { +public class GH846Tests { @Autowired private MyBean myBean; diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java index c41db1898..1b04f158c 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessorTest.java @@ -16,12 +16,18 @@ package org.springframework.cloud.sleuth.instrument.web.client; +import java.net.URI; + +import brave.propagation.Propagation; import org.assertj.core.api.BDDAssertions; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.Mock; import org.mockito.junit.MockitoJUnitRunner; + import org.springframework.beans.factory.BeanFactory; +import org.springframework.http.HttpMethod; +import org.springframework.web.reactive.function.client.ClientRequest; import org.springframework.web.reactive.function.client.WebClient; /**