diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java index 1e1cac58a..3f1f2e301 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java @@ -60,27 +60,27 @@ ApplicationEventPublisherAware { @Override public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException { + if (getCurrentSpan() == null) { + return execution.execute(request, body); + } setHeader(request, SPAN_ID_NAME, getCurrentSpan().getSpanId()); setHeader(request, TRACE_ID_NAME, getCurrentSpan().getTraceId()); setHeader(request, SPAN_NAME_NAME, getCurrentSpan().getName()); - String parentId = getParentId(getCurrentSpan()); - if (parentId != null) { - setHeader(request, PARENT_ID_NAME, parentId); - } - String processId = getCurrentSpan().getProcessId(); - if (processId != null) { - setHeader(request, PROCESS_ID_NAME, processId); - } + setHeader(request, PARENT_ID_NAME, getParentId(getCurrentSpan())); + setHeader(request, PROCESS_ID_NAME, getCurrentSpan().getProcessId()); publish(new ClientSentEvent(this, getCurrentSpan())); return new TraceHttpResponse(this, execution.execute(request, body)); } public void close() { + if (getCurrentSpan() == null) { + return; + } publish(new ClientReceivedEvent(this, getCurrentSpan())); } private void publish(ApplicationEvent event) { - if (this.publisher !=null) { + if (this.publisher != null) { this.publisher.publishEvent(event); } } @@ -91,7 +91,7 @@ ApplicationEventPublisherAware { } public void setHeader(HttpRequest request, String name, String value) { - if (!request.getHeaders().containsKey(name) && isTracing()) { + if (value != null && !request.getHeaders().containsKey(name) && isTracing()) { request.getHeaders().add(name, value); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java index cebbe3314..5ce9362d6 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java @@ -22,8 +22,8 @@ import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME; import org.junit.After; import org.junit.Test; -import org.springframework.boot.SpringApplication; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.builder.SpringApplicationBuilder; import org.springframework.cloud.sleuth.Trace; import org.springframework.cloud.sleuth.TraceScope; import org.springframework.cloud.sleuth.sampler.AlwaysSampler; @@ -47,18 +47,18 @@ public class TraceContextPropagationChannelInterceptorTests { @After public void close() { - if (context != null) { - context.close(); + if (this.context != null) { + this.context.close(); } } @Test public void testSpanPropagation() { - context = SpringApplication.run(App.class); + this.context = new SpringApplicationBuilder(App.class).web(false).run(); - PollableChannel channel = context.getBean("channel", PollableChannel.class); + PollableChannel channel = this.context.getBean("channel", PollableChannel.class); - Trace trace = context.getBean(Trace.class); + Trace trace = this.context.getBean(Trace.class); TraceScope traceScope = trace.startSpan("testSendMessage", new AlwaysSampler(), null); channel.send(MessageBuilder.withPayload("hi").build()); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptorTests.java new file mode 100644 index 000000000..4edf41091 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptorTests.java @@ -0,0 +1,108 @@ +/* + * Copyright 2015 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.springframework.cloud.sleuth.MilliSpan; +import org.springframework.cloud.sleuth.Trace; +import org.springframework.cloud.sleuth.TraceContextHolder; +import org.springframework.http.HttpHeaders; +import org.springframework.http.client.ClientHttpRequestInterceptor; +import org.springframework.test.web.client.MockMvcClientHttpRequestFactory; +import org.springframework.test.web.servlet.MockMvc; +import org.springframework.test.web.servlet.setup.MockMvcBuilders; +import org.springframework.web.bind.annotation.RequestHeader; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.client.RestTemplate; + +/** + * @author Dave Syer + * + */ +public class TraceRestTemplateInterceptorTests { + + private MockMvc mockMvc = MockMvcBuilders.standaloneSetup(new TestController()) + .build(); + + private RestTemplate template = new RestTemplate(new MockMvcClientHttpRequestFactory( + this.mockMvc)); + + @Before + public void setup() { + this.template + .setInterceptors(Arrays + . asList(new TraceRestTemplateInterceptor())); + } + + @After + public void clean() { + TraceContextHolder.setCurrentSpan(null); + } + + @Test + public void headersAddedWhenTracing() { + TraceContextHolder.setCurrentSpan(MilliSpan.builder().traceId("foo") + .spanId("bar").build()); + @SuppressWarnings("unchecked") + Map headers = this.template.getForEntity("/", Map.class) + .getBody(); + assertEquals("bar", headers.get(Trace.SPAN_ID_NAME)); + assertEquals("foo", headers.get(Trace.TRACE_ID_NAME)); + } + + @Test + public void headersNotAddedWhenNotTracing() { + @SuppressWarnings("unchecked") + Map headers = this.template.getForEntity("/", Map.class) + .getBody(); + assertFalse("Wrong headers: " + headers, headers.containsKey(Trace.SPAN_ID_NAME)); + } + + @RestController + public static class TestController { + @RequestMapping("/") + public Map home(@RequestHeader HttpHeaders headers) { + Map map = new HashMap(); + addHeaders(map, headers, Trace.SPAN_ID_NAME, Trace.TRACE_ID_NAME, + Trace.PARENT_ID_NAME, Trace.SPAN_NAME_NAME, Trace.PROCESS_ID_NAME); + return map; + } + + private void addHeaders(Map map, HttpHeaders headers, + String... names) { + if (headers != null) { + for (String name : names) { + String value = headers.getFirst(name); + if (value != null) { + map.put(name, value); + } + } + } + } + } + +}