Merge pull request #30 from spring-cloud/threadlocal-remove

Use ThreadLocal.remove() rather than .set(null)

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceContextHolder.java

@@ -18,6 +18,7 @@ package org.springframework.cloud.sleuth;
 
 import lombok.extern.apachecommons.CommonsLog;
 import org.springframework.core.NamedThreadLocal;
+import org.springframework.util.Assert;
 
 /**
  * @author Spencer Gibb
@@ -31,12 +32,17 @@ public class TraceContextHolder {
 	}
 
 	public static void setCurrentSpan(Span span) {
+		Assert.notNull(span, "span can not be null.  Use removeCurrentSpan() instead.");
 		if (log.isTraceEnabled()) {
 			log.trace("Setting current span " + span);
 		}
 		currentSpan.set(span);
 	}
 
+	public static void removeCurrentSpan() {
+		currentSpan.remove();
+	}
+
 	public static boolean isTracing() {
 		return currentSpan.get() != null;
 	}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java

@@ -58,7 +58,7 @@ public class TraceChannelInterceptor extends ChannelInterceptorAdapter {
 		if (traceScope != null) {
 			traceScope.close();
 		}
-		this.traceScopeHolder.set(null);
+		this.traceScopeHolder.remove();
 		// TODO: Maybe the TraceScope could handle this
 		TraceContextHolder.setCurrentSpan(this.spanHolder.get());
 	}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java

@@ -22,6 +22,7 @@ import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME;
 import static org.springframework.cloud.sleuth.Trace.SPAN_NAME_NAME;
 import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME;
 import static org.springframework.cloud.sleuth.TraceContextHolder.getCurrentSpan;
+import static org.springframework.cloud.sleuth.TraceContextHolder.removeCurrentSpan;
 import static org.springframework.cloud.sleuth.TraceContextHolder.setCurrentSpan;
 
 import java.util.HashMap;
@@ -113,7 +114,7 @@ implements ExecutorChannelInterceptor {
 		Span originalContext = ORIGINAL_CONTEXT.get();
 		try {
 			if (originalContext == null) {
-				setCurrentSpan(null);
+				removeCurrentSpan();
 				ORIGINAL_CONTEXT.remove();
 			}
 			else {
@@ -121,7 +122,7 @@ implements ExecutorChannelInterceptor {
 			}
 		}
 		catch (Throwable t) {// NOSONAR
-			setCurrentSpan(null);
+			removeCurrentSpan();
 		}
 	}
 

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java

@@ -146,7 +146,7 @@ public class TraceFilter extends OncePerRequestFilter {
 				addResponseAnnotations(response);
 				traceScope.close();
 			}
-			TraceContextHolder.setCurrentSpan(null);
+			TraceContextHolder.removeCurrentSpan();
 		}
 	}
 

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/TraceContextHolderTests.java

@@ -0,0 +1,30 @@
+/*
+ * 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;
+
+import org.junit.Test;
+
+/**
+ * @author Spencer Gibb
+ */
+public class TraceContextHolderTests {
+
+	@Test(expected = IllegalArgumentException.class)
+	public void setCurrentSpanNotNull() {
+		TraceContextHolder.setCurrentSpan(null);
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java

@@ -76,7 +76,7 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 
 	@After
 	public void close() {
-		TraceContextHolder.setCurrentSpan(null);
+		TraceContextHolder.removeCurrentSpan();
 		this.channel.unsubscribe(this);
 	}
 

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java

@@ -61,7 +61,7 @@ public class TraceContextPropagationChannelInterceptorTests {
 
 	@After
 	public void close() {
-		TraceContextHolder.setCurrentSpan(null);
+		TraceContextHolder.removeCurrentSpan();
 	}
 
 	@Test

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterIntegrationTests.java

@@ -55,7 +55,7 @@ public class TraceFilterIntegrationTests {
 	@Before
 	@SneakyThrows
 	public void init() {
-		TraceContextHolder.setCurrentSpan(null);
+		TraceContextHolder.removeCurrentSpan();
 		this.context.refresh();
 		this.request = builder().buildRequest(new MockServletContext());
 		this.response = new MockHttpServletResponse();

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptorTests.java

@@ -60,7 +60,7 @@ public class TraceRestTemplateInterceptorTests {
 
 	@After
 	public void clean() {
-		TraceContextHolder.setCurrentSpan(null);
+		TraceContextHolder.removeCurrentSpan();
 	}
 
 	@Test