From 2e45374800a3b8f4b2fb9c8dfc4cd7d139002132 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Mon, 17 Dec 2018 13:20:01 +0100
Subject: [PATCH] Clearing MDC for null context; fixes gh-1161

---
 .../cloud/sleuth/log/Slf4jScopeDecorator.java | 40 ++++++++++++-------
 .../cloud/sleuth/log/Slf4JSpanLoggerTest.java | 20 ++++++++++
 2 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jScopeDecorator.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jScopeDecorator.java
index 26d472710..4419b941b 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jScopeDecorator.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jScopeDecorator.java
@@ -77,8 +77,8 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator {
 		final String legacyPreviousSpanId = MDC.get(LEGACY_SPAN_ID_NAME);
 		final String legacySpanExportable = MDC.get(LEGACY_EXPORTABLE_NAME);
 		final List<AbstractMap.SimpleEntry<String, String>> previousMdc = Stream
-				.concat(whitelistedBaggageKeys(currentSpan),
-						whitelistedPropagationKeys(currentSpan))
+				.concat(whitelistedBaggageKeysWithValue(currentSpan),
+						whitelistedPropagationKeysWithValue(currentSpan))
 				.map((s) -> new AbstractMap.SimpleEntry<>(s, MDC.get(s)))
 				.collect(Collectors.toList());
 
@@ -102,9 +102,9 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator {
 					log.trace("With parent: {}", currentSpan.parentId());
 				}
 			}
-			whitelistedBaggageKeys(currentSpan).forEach(
+			whitelistedBaggageKeysWithValue(currentSpan).forEach(
 					(s) -> MDC.put(s, ExtraFieldPropagation.get(currentSpan, s)));
-			whitelistedPropagationKeys(currentSpan).forEach(
+			whitelistedPropagationKeysWithValue(currentSpan).forEach(
 					(s) -> MDC.put(s, ExtraFieldPropagation.get(currentSpan, s)));
 		}
 		else {
@@ -116,8 +116,8 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator {
 			MDC.remove(LEGACY_PARENT_ID_NAME);
 			MDC.remove(LEGACY_SPAN_ID_NAME);
 			MDC.remove(LEGACY_EXPORTABLE_NAME);
-			whitelistedBaggageKeys(currentSpan).forEach(MDC::remove);
-			whitelistedPropagationKeys(currentSpan).forEach(MDC::remove);
+			whitelistedBaggageKeys().forEach(MDC::remove);
+			whitelistedPropagationKeys().forEach(MDC::remove);
 		}
 
 		/**
@@ -146,18 +146,30 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator {
 		return new ThreadContextCurrentTraceContextScope();
 	}
 
-	private Stream<String> whitelistedBaggageKeys(TraceContext context) {
+	private Stream<String> whitelistedBaggageKeys() {
 		return this.sleuthProperties.getBaggageKeys().stream().filter(
-				(s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s)
-						&& context != null
-						&& StringUtils.hasText(ExtraFieldPropagation.get(context, s)));
+				(s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s));
 	}
 
-	private Stream<String> whitelistedPropagationKeys(TraceContext context) {
+	private Stream<String> whitelistedBaggageKeysWithValue(TraceContext context) {
+		if (context == null) {
+			return Stream.empty();
+		}
+		return whitelistedBaggageKeys().filter(
+				(s) -> StringUtils.hasText(ExtraFieldPropagation.get(context, s)));
+	}
+
+	private Stream<String> whitelistedPropagationKeys() {
 		return this.sleuthProperties.getPropagationKeys().stream().filter(
-				(s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s)
-						&& context != null
-						&& StringUtils.hasText(ExtraFieldPropagation.get(context, s)));
+				(s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s));
+	}
+
+	private Stream<String> whitelistedPropagationKeysWithValue(TraceContext context) {
+		if (context == null) {
+			return Stream.empty();
+		}
+		return whitelistedPropagationKeys().filter(
+				(s) -> StringUtils.hasText(ExtraFieldPropagation.get(context, s)));
 	}
 
 	private void log(String text, TraceContext span) {
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
index f016c8a41..8e818d409 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
@@ -92,6 +92,26 @@ public class Slf4JSpanLoggerTest {
 		assertThat(MDC.get("my-propagation")).isNullOrEmpty();
 	}
 
+	@Test
+	public void should_remove_entries_from_mdc_for_null_span() throws Exception {
+		ExtraFieldPropagation.set(this.span.context(), "my-baggage", "my-value");
+		ExtraFieldPropagation.set(this.span.context(), "my-propagation",
+				"my-propagation-value");
+		this.slf4jScopeDecorator.decorateScope(this.span.context(), () -> {
+		});
+
+		assertThat(MDC.get("my-baggage")).isEqualTo("my-value");
+		assertThat(MDC.get("my-propagation")).isEqualTo("my-propagation-value");
+
+		Scope scope = this.slf4jScopeDecorator.decorateScope(null, () -> {
+		});
+
+		scope.close();
+
+		assertThat(MDC.get("my-baggage")).isNullOrEmpty();
+		assertThat(MDC.get("my-propagation")).isNullOrEmpty();
+	}
+
 	@Test
 	public void should_remove_entries_from_mdc_from_null_span() throws Exception {
 		MDC.put("X-B3-TraceId", "A");