From 70708e22a10c67f942f52228116f21da9136b004 Mon Sep 17 00:00:00 2001 From: Dave Syer Date: Wed, 2 Dec 2015 07:58:46 +0000 Subject: [PATCH] Remove trace from current thread in thread pool The TraceManager.detach() method should be called by wrapper tasks submitted to background threads. This change fixes it in TraceCallable and TraceRunnable. Possibly still some work to do for Hystrix. The "detached" flag didn't seem to be doing anything, so I removed it in the process. Fixes gh-64, fixes gh-61 --- .../springframework/cloud/sleuth/Trace.java | 9 +----- .../cloud/sleuth/TraceManager.java | 8 +++++- .../sleuth/instrument/TraceCallable.java | 2 +- .../sleuth/instrument/TraceDelegate.java | 7 +++++ .../sleuth/instrument/TraceRunnable.java | 2 +- .../sleuth/trace/DefaultTraceManager.java | 15 ++++++---- .../sleuth/instrument/TraceCallableTest.java | 28 +++++++++++++++++-- 7 files changed, 52 insertions(+), 19 deletions(-) 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