Ensure span context is cleared at end of interceptors

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<TraceScope> traceScopeHolder = new ThreadLocal<TraceScope>();
+
+	private ThreadLocal<Span> spanHolder = new ThreadLocal<Span>();
+
 	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);
 	}

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);
 		}
 	}
 

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

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());
+	}
+
+}