Removing entries from MDC (#252)

when iterating over saved spans during span closing we can apply a function that should be executed on each iteration.

fixes #248

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

@@ -144,7 +144,6 @@ public class TraceFilter extends OncePerRequestFilter {
 				} else {
 					spanFromRequest.logEvent(Span.SERVER_SEND);
 				}
-				// Double close to clean up the parent (remote span as well)
 				this.tracer.close(spanFromRequest);
 			}
 		}

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

@@ -67,7 +67,7 @@ public class Slf4jSpanLogger implements SpanLogger {
 	@Override
 	public void logStoppedSpan(Span parent, Span span) {
 		log("Stopped span: {}", span);
-		if (parent != null) {
+		if (span != null && parent != null) {
 			log("With parent: {}", parent);
 			MDC.put(Span.SPAN_ID_NAME, Span.idToHex(parent.getSpanId()));
 			MDC.put(Span.SPAN_EXPORT_NAME, String.valueOf(parent.isExportable()));
@@ -80,7 +80,7 @@ public class Slf4jSpanLogger implements SpanLogger {
 	}
 
 	private void log(String text, Span span) {
-		if (this.nameSkipPattern.matcher(span.getName()).matches()) {
+		if (span != null && this.nameSkipPattern.matcher(span.getName()).matches()) {
 			return;
 		}
 		this.log.trace(text, span);

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

@@ -115,7 +115,7 @@ public class DefaultTracer implements Tracer {
 			return null;
 		}
 		Span cur = SpanContextHolder.getCurrentSpan();
-		Span savedSpan = span.getSavedSpan();
+		final Span savedSpan = span.getSavedSpan();
 		if (!span.equals(cur)) {
 			ExceptionUtils.warn(
 					"Tried to close span but " + "it is not the current span: " + span
@@ -133,7 +133,11 @@ public class DefaultTracer implements Tracer {
 					this.spanLogger.logStoppedSpan(null, span);
 				}
 			}
-			SpanContextHolder.close();
+			SpanContextHolder.close(new SpanContextHolder.SpanFunction() {
+				@Override public void apply(Span span) {
+					DefaultTracer.this.spanLogger.logStoppedSpan(savedSpan, span);
+				}
+			});
 		}
 		return savedSpan;
 	}

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

@@ -68,12 +68,14 @@ class SpanContextHolder {
 
 	/**
 	 * Close the current span and all parents that can be auto closed.
+	 * On every iteration a function will be applied on the closed Span.
 	 */
-	static void close() {
+	static void close(SpanFunction spanFunction) {
 		SpanContext current = CURRENT_SPAN.get();
 		CURRENT_SPAN.remove();
 		while (current != null) {
 			current = current.parent;
+			spanFunction.apply(current != null ? current.span : null);
 			if (current != null) {
 				if (!current.autoClose) {
 					CURRENT_SPAN.set(current);
@@ -83,6 +85,13 @@ class SpanContextHolder {
 		}
 	}
 
+	/**
+	 * Close the current span and all parents that can be auto closed.
+	 */
+	static void close() {
+		close(new NoOpFunction());
+	}
+
 	/**
 	 * Push a span into the thread context, with the option to have it auto close if any
 	 * child spans are themselves closed. Use autoClose=true if you start a new span with
@@ -113,4 +122,12 @@ class SpanContextHolder {
 			this.parent = CURRENT_SPAN.get();
 		}
 	}
+
+	interface SpanFunction {
+		void apply(Span span);
+	}
+
+	private static class NoOpFunction implements SpanFunction {
+		@Override public void apply(Span span) { }
+	}
 }

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

@@ -7,6 +7,7 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.slf4j.MDC;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.actuate.autoconfigure.ManagementServerProperties;
 import org.springframework.boot.test.SpringApplicationConfiguration;
@@ -81,6 +82,16 @@ public class TraceFilterIntegrationTests extends AbstractMvcIntegrationTest {
 		then(tracingHeaderFrom(mvcResult)).isEqualTo(expectedTraceId);
 	}
 
+	@Test
+	public void when_message_is_sent_should_eventually_clear_mdc()
+			throws Exception {
+		Long expectedTraceId = new Random().nextLong();
+
+		whenSentPingWithTraceId(expectedTraceId);
+
+		then(MDC.getCopyOfContextMap()).isEmpty();
+	}
+
 	@Test
 	public void when_traceId_is_sent_to_async_endpoint_span_is_joined()
 			throws Exception {