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

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

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);
 
 	<V> Callable<V> wrap(Callable<V> callable);

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

@@ -46,7 +46,7 @@ public class TraceCallable<V> extends TraceDelegate<Callable<V>>implements Calla
 			return this.getDelegate().call();
 		}
 		finally {
-			close(trace);
+			closeAll(trace);
 		}
 	}
 

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

@@ -48,6 +48,13 @@ public abstract class TraceDelegate<T> {
 		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);
 	}

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

@@ -44,7 +44,7 @@ public class TraceRunnable extends TraceDelegate<Runnable>implements Runnable {
 			this.getDelegate().run();
 		}
 		finally {
-			close(trace);
+			closeAll(trace);
 		}
 	}
 }

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

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<Trace> thatRetrievesTraceFromThreadLocal() {
 		return new Callable<Trace>() {
 			@Override