From b1537f38e0b7a3e129eb5b61c937a3b4910b20e9 Mon Sep 17 00:00:00 2001 From: Spencer Gibb Date: Tue, 18 Aug 2015 12:51:04 -0600 Subject: [PATCH] Use ThreadLocal.remove() rather than .set(null) fixes gh-27 --- .../cloud/sleuth/TraceContextHolder.java | 6 ++++ .../integration/TraceChannelInterceptor.java | 2 +- ...eContextPropagationChannelInterceptor.java | 5 ++-- .../sleuth/instrument/web/TraceFilter.java | 2 +- .../cloud/sleuth/TraceContextHolderTests.java | 30 +++++++++++++++++++ .../TraceChannelInterceptorTests.java | 2 +- ...extPropagationChannelInterceptorTests.java | 2 +- .../web/TraceFilterIntegrationTests.java | 2 +- .../TraceRestTemplateInterceptorTests.java | 2 +- 9 files changed, 45 insertions(+), 8 deletions(-) create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/TraceContextHolderTests.java diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceContextHolder.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceContextHolder.java index 283404642..79c55c90d 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceContextHolder.java +++ b/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; } 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 46f70540e..260c78e6d 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 @@ -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()); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java index 3cc2ebf02..26f538c07 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java +++ b/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(); } } 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 10046e2d9..d5410e69a 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 @@ -146,7 +146,7 @@ public class TraceFilter extends OncePerRequestFilter { addResponseAnnotations(response); traceScope.close(); } - TraceContextHolder.setCurrentSpan(null); + TraceContextHolder.removeCurrentSpan(); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/TraceContextHolderTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/TraceContextHolderTests.java new file mode 100644 index 000000000..34b5c3d50 --- /dev/null +++ b/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); + } +} 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 b6caa9702..2100d798e 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 @@ -76,7 +76,7 @@ public class TraceChannelInterceptorTests implements MessageHandler { @After public void close() { - TraceContextHolder.setCurrentSpan(null); + TraceContextHolder.removeCurrentSpan(); this.channel.unsubscribe(this); } 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 1222a5d32..3449aff40 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 @@ -61,7 +61,7 @@ public class TraceContextPropagationChannelInterceptorTests { @After public void close() { - TraceContextHolder.setCurrentSpan(null); + TraceContextHolder.removeCurrentSpan(); } @Test 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 index 36a9ec369..ea92d8227 100644 --- 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 @@ -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(); 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 index 290594456..52dbc3519 100644 --- 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 @@ -60,7 +60,7 @@ public class TraceRestTemplateInterceptorTests { @After public void clean() { - TraceContextHolder.setCurrentSpan(null); + TraceContextHolder.removeCurrentSpan(); } @Test