diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java index 276839df6..46f70540e 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java @@ -26,9 +26,11 @@ import static org.springframework.util.StringUtils.hasText; import org.springframework.cloud.sleuth.MilliSpan; import org.springframework.cloud.sleuth.MilliSpan.MilliSpanBuilder; +import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Trace; import org.springframework.cloud.sleuth.TraceContextHolder; import org.springframework.cloud.sleuth.TraceScope; +import org.springframework.integration.channel.AbstractMessageChannel; import org.springframework.integration.context.IntegrationObjectSupport; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; @@ -40,14 +42,30 @@ import org.springframework.messaging.support.ChannelInterceptorAdapter; */ public class TraceChannelInterceptor extends ChannelInterceptorAdapter { + private ThreadLocal traceScopeHolder = new ThreadLocal(); + + private ThreadLocal spanHolder = new ThreadLocal(); + private final Trace trace; public TraceChannelInterceptor(Trace trace) { this.trace = trace; } + @Override + public void postSend(Message message, MessageChannel channel, boolean sent) { + TraceScope traceScope = this.traceScopeHolder.get(); + if (traceScope != null) { + traceScope.close(); + } + this.traceScopeHolder.set(null); + // TODO: Maybe the TraceScope could handle this + TraceContextHolder.setCurrentSpan(this.spanHolder.get()); + } + @Override public Message preSend(Message message, MessageChannel channel) { + this.spanHolder.set(TraceContextHolder.getCurrentSpan()); if (TraceContextHolder.isTracing() || message.getHeaders().containsKey(NOT_SAMPLED_NAME)) { return SpanMessageHeaders.addSpanHeaders(message, @@ -55,9 +73,7 @@ public class TraceChannelInterceptor extends ChannelInterceptorAdapter { } String spanId = getHeader(message, SPAN_ID_NAME); String traceId = getHeader(message, TRACE_ID_NAME); - String name = "message/" - + ((channel instanceof IntegrationObjectSupport) ? ((IntegrationObjectSupport) channel) - .getComponentName() : channel.toString()); + String name = "message/" + getChannelName(channel); TraceScope traceScope; if (hasText(spanId) && hasText(traceId)) { @@ -82,9 +98,24 @@ public class TraceChannelInterceptor extends ChannelInterceptorAdapter { else { traceScope = this.trace.startSpan(name); } + this.traceScopeHolder.set(traceScope); return SpanMessageHeaders.addSpanHeaders(message, traceScope.getSpan()); } + private String getChannelName(MessageChannel channel) { + String name = null; + if (channel instanceof IntegrationObjectSupport) { + name = ((IntegrationObjectSupport) channel).getComponentName(); + } + if (name == null && channel instanceof AbstractMessageChannel) { + name = ((AbstractMessageChannel) channel).getFullChannelName(); + } + if (name == null) { + name = channel.toString(); + } + return name; + } + private String getHeader(Message message, String name) { return (String) message.getHeaders().get(name); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java index 15755b8a4..10046e2d9 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java @@ -35,6 +35,7 @@ import javax.servlet.http.HttpServletResponse; import org.springframework.cloud.sleuth.MilliSpan; import org.springframework.cloud.sleuth.MilliSpan.MilliSpanBuilder; import org.springframework.cloud.sleuth.Trace; +import org.springframework.cloud.sleuth.TraceContextHolder; import org.springframework.cloud.sleuth.TraceScope; import org.springframework.core.Ordered; import org.springframework.core.annotation.Order; @@ -145,6 +146,7 @@ public class TraceFilter extends OncePerRequestFilter { addResponseAnnotations(response); traceScope.close(); } + TraceContextHolder.setCurrentSpan(null); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java index aa1bab808..b6caa9702 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java @@ -100,6 +100,7 @@ public class TraceChannelInterceptorTests implements MessageHandler { String traceId = this.message.getHeaders().get(TRACE_ID_NAME, String.class); assertNotNull("traceId was null", traceId); + assertNull(TraceContextHolder.getCurrentSpan()); } @Test @@ -115,6 +116,7 @@ public class TraceChannelInterceptorTests implements MessageHandler { String traceId = this.message.getHeaders().get(TRACE_ID_NAME, String.class); assertNotNull("traceId was null", traceId); + assertNull(TraceContextHolder.getCurrentSpan()); } @Configuration diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterIntegrationTests.java new file mode 100644 index 000000000..36a9ec369 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterIntegrationTests.java @@ -0,0 +1,87 @@ +/* + * Copyright 2013-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; + +import static org.junit.Assert.assertNull; +import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME; +import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME; +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import lombok.SneakyThrows; + +import org.junit.Before; +import org.junit.Test; +import org.springframework.cloud.sleuth.RandomUuidGenerator; +import org.springframework.cloud.sleuth.Trace; +import org.springframework.cloud.sleuth.TraceContextHolder; +import org.springframework.cloud.sleuth.sampler.AlwaysSampler; +import org.springframework.cloud.sleuth.trace.DefaultTrace; +import org.springframework.context.support.StaticApplicationContext; +import org.springframework.http.MediaType; +import org.springframework.mock.web.MockFilterChain; +import org.springframework.mock.web.MockHttpServletRequest; +import org.springframework.mock.web.MockHttpServletResponse; +import org.springframework.mock.web.MockServletContext; +import org.springframework.test.web.servlet.request.MockHttpServletRequestBuilder; + +/** + * @author Spencer Gibb + * @author Dave Syer + */ +public class TraceFilterIntegrationTests { + + private StaticApplicationContext context = new StaticApplicationContext(); + + private Trace trace = new DefaultTrace(new AlwaysSampler(), + new RandomUuidGenerator(), this.context); + + private MockHttpServletRequest request; + private MockHttpServletResponse response; + private MockFilterChain filterChain; + + @Before + @SneakyThrows + public void init() { + TraceContextHolder.setCurrentSpan(null); + this.context.refresh(); + this.request = builder().buildRequest(new MockServletContext()); + this.response = new MockHttpServletResponse(); + this.response.setContentType(MediaType.APPLICATION_JSON_VALUE); + this.filterChain = new MockFilterChain(); + } + + public MockHttpServletRequestBuilder builder() { + return get("/").accept(MediaType.APPLICATION_JSON) + .header("User-Agent", "MockMvc"); + } + + @Test + public void startsNewTrace() throws Exception { + TraceFilter filter = new TraceFilter(this.trace); + filter.doFilter(this.request, this.response, this.filterChain); + assertNull(TraceContextHolder.getCurrentSpan()); + } + + @Test + public void continuesSpanFromHeaders() throws Exception { + this.request = builder().header(SPAN_ID_NAME, "myspan") + .header(TRACE_ID_NAME, "mytrace").buildRequest(new MockServletContext()); + TraceFilter filter = new TraceFilter(this.trace); + filter.doFilter(this.request, this.response, this.filterChain); + assertNull(TraceContextHolder.getCurrentSpan()); + } + +}