diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/Trace.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/Trace.java index 58840e7ac..6aa5429a3 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/Trace.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/Trace.java @@ -54,9 +54,6 @@ public class Trace { */ private final Trace savedTrace; - @NonFinal - private boolean detached = false; - public Trace(Trace saved, Span span) { this.savedTrace = saved; this.span = span; @@ -67,13 +64,9 @@ public class Trace { } public void addAnnotation(String key, String value) { - if (this.span != null && !this.detached) { + if (this.span != null) { this.span.addAnnotation(key, value); } } - public void detach() { - this.detached = true; - } - } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceManager.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceManager.java index c474a925c..d3e1821ef 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceManager.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceManager.java @@ -95,7 +95,7 @@ public interface TraceManager extends TraceAccessor { void addAnnotation(String key, String value); /** - * Remove this span as the current thread, but don't stop it yet or send it for + * Remove this span from the current thread, but don't stop it yet or send it for * collection. This is useful if the span object is then passed to another thread for * use with Trace.continueTrace(). * @@ -103,6 +103,12 @@ public interface TraceManager extends TraceAccessor { */ Trace detach(Trace trace); + /** + * Remove this span from the current thread, stop it and send it for collection. + * + * @param trace the trace to close + * @return the saved trace if there was one before the trace started (null otherwise) + */ Trace close(Trace trace); Callable wrap(Callable callable); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java index 88f03a176..902a2376f 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java @@ -46,7 +46,7 @@ public class TraceCallable extends TraceDelegate>implements Calla return this.getDelegate().call(); } finally { - close(trace); + closeAll(trace); } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java index 14242e574..c5fcd2e3b 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java @@ -48,6 +48,13 @@ public abstract class TraceDelegate { this.traceManager.close(scope); } + protected void closeAll(Trace scope) { + scope = this.traceManager.close(scope); + while (scope != null) { + scope = this.traceManager.detach(scope); + } + } + protected Trace startSpan() { return this.traceManager.startSpan(getSpanName(), this.parent); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java index ab367f47b..844f32ca3 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java @@ -44,7 +44,7 @@ public class TraceRunnable extends TraceDelegateimplements Runnable { this.getDelegate().run(); } finally { - close(trace); + closeAll(trace); } } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/DefaultTraceManager.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/DefaultTraceManager.java index 0648b3081..5f2349191 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/DefaultTraceManager.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/DefaultTraceManager.java @@ -85,7 +85,6 @@ public class DefaultTraceManager implements TraceManager { if (trace == null) { return null; } - trace.detach(); Span cur = TraceContextHolder.getCurrentSpan(); Span span = trace.getSpan(); if (cur != span) { @@ -98,6 +97,9 @@ public class DefaultTraceManager implements TraceManager { if (span != NullTrace.INSTANCE) { TraceContextHolder.setCurrentTrace(trace.getSavedTrace()); } + else { + TraceContextHolder.removeCurrentTrace(); + } } return trace.getSavedTrace(); } @@ -117,7 +119,7 @@ public class DefaultTraceManager implements TraceManager { + ". You have " + "probably forgotten to close or detach " + cur); } else { - if (span != NullTrace.INSTANCE && span!=null) { + if (span != NullTrace.INSTANCE && span != null) { span.stop(); if (savedTrace != null && span.getParents().contains(savedTrace.getSpan().getSpanId())) { @@ -130,6 +132,9 @@ public class DefaultTraceManager implements TraceManager { TraceContextHolder.removeCurrentTrace(); } } + else { + TraceContextHolder.removeCurrentTrace(); + } } return savedTrace; } @@ -137,8 +142,7 @@ public class DefaultTraceManager implements TraceManager { protected Span createChild(Span parent, String name) { if (parent == null) { MilliSpan span = MilliSpan.builder().begin(System.currentTimeMillis()) - .name(name).traceId(createId()) - .spanId(createId()).build(); + .name(name).traceId(createId()).spanId(createId()).build(); this.publisher.publishEvent(new SpanAcquiredEvent(this, span)); return span; } @@ -149,8 +153,7 @@ public class DefaultTraceManager implements TraceManager { } MilliSpan span = MilliSpan.builder().begin(System.currentTimeMillis()) .name(name).traceId(parent.getTraceId()).parent(parent.getSpanId()) - .spanId(createId()).processId(parent.getProcessId()) - .build(); + .spanId(createId()).processId(parent.getProcessId()).build(); this.publisher.publishEvent(new SpanAcquiredEvent(this, parent, span)); return span; } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/TraceCallableTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/TraceCallableTest.java index 0d749a45c..cf9c39979 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/TraceCallableTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/TraceCallableTest.java @@ -6,6 +6,7 @@ import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import org.junit.After; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.Mockito; @@ -25,8 +26,13 @@ public class TraceCallableTest { TraceManager traceManager = new DefaultTraceManager(new AlwaysSampler(), new JdkIdGenerator(), Mockito.mock(ApplicationEventPublisher.class)); + @After + public void clean() { + TraceContextHolder.removeCurrentTrace(); + } + @Test - public void should_remove_span_from_thread_local_after_finishing_work() + public void should_not_see_same_trace_id_in_successive_tasks() throws Exception { Trace firstTrace = givenCallableGetsSubmitted( thatRetrievesTraceFromThreadLocal()); @@ -40,7 +46,7 @@ public class TraceCallableTest { } @Test - public void should_not_find_thread_local_in_non_traceable_callback() + public void should_remove_span_from_thread_local_after_finishing_work() throws Exception { givenCallableGetsSubmitted(thatRetrievesTraceFromThreadLocal()); @@ -50,6 +56,24 @@ public class TraceCallableTest { then(secondTrace).isNull(); } + @Test + public void should_remove_parent_span_from_thread_local_after_finishing_work() + throws Exception { + Trace parent = givenSpanIsAlreadyActive(); + Trace child = givenCallableGetsSubmitted(thatRetrievesTraceFromThreadLocal()); + then(parent).as("parent").isNotNull(); + then(child.getSavedTrace()).isEqualTo(parent); + + Trace secondTrace = whenNonTraceableCallableGetsSubmitted( + thatRetrievesTraceFromThreadLocal()); + + then(secondTrace).isNull(); + } + + private Trace givenSpanIsAlreadyActive() { + return this.traceManager.startSpan("parent"); + } + private Callable thatRetrievesTraceFromThreadLocal() { return new Callable() { @Override