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 60ead6443..87eb87b4f 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 @@ -17,9 +17,9 @@ package org.springframework.cloud.sleuth.log; import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; -import java.util.stream.Collectors; -import java.util.stream.Stream; import brave.internal.HexCodec; import brave.internal.Nullable; @@ -86,11 +86,7 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator { final String legacyPreviousParentId = MDC.get(LEGACY_PARENT_ID_NAME); final String legacyPreviousSpanId = MDC.get(LEGACY_SPAN_ID_NAME); final String legacySpanExportable = MDC.get(LEGACY_EXPORTABLE_NAME); - final List> previousMdc = Stream - .concat(whitelistedBaggageKeysWithValue(currentSpan), - whitelistedPropagationKeysWithValue(currentSpan)) - .map((s) -> new AbstractMap.SimpleEntry<>(s, MDC.get(s))) - .collect(Collectors.toList()); + final List> previousMdc = previousMdc(); if (currentSpan != null) { String traceIdString = currentSpan.traceIdString(); @@ -112,10 +108,12 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator { log.trace("With parent: {}", currentSpan.parentId()); } } - whitelistedBaggageKeysWithValue(currentSpan).forEach( - (s) -> MDC.put(s, ExtraFieldPropagation.get(currentSpan, s))); - whitelistedPropagationKeysWithValue(currentSpan).forEach( - (s) -> MDC.put(s, ExtraFieldPropagation.get(currentSpan, s))); + for (String key : whitelistedBaggageKeysWithValue(currentSpan)) { + MDC.put(key, ExtraFieldPropagation.get(currentSpan, key)); + } + for (String key : whitelistedPropagationKeysWithValue(currentSpan)) { + MDC.put(key, ExtraFieldPropagation.get(currentSpan, key)); + } } else { MDC.remove("traceId"); @@ -126,8 +124,13 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator { MDC.remove(LEGACY_PARENT_ID_NAME); MDC.remove(LEGACY_SPAN_ID_NAME); MDC.remove(LEGACY_EXPORTABLE_NAME); - whitelistedBaggageKeys().forEach(MDC::remove); - whitelistedPropagationKeys().forEach(MDC::remove); + for (String s : whitelistedBaggageKeys()) { + MDC.remove(s); + } + for (String s : whitelistedPropagationKeys()) { + MDC.remove(s); + } + previousMdc.clear(); } /** @@ -158,30 +161,54 @@ final class Slf4jScopeDecorator implements CurrentTraceContext.ScopeDecorator { return new ThreadContextCurrentTraceContextScope(); } - private Stream whitelistedBaggageKeys() { - return this.sleuthProperties.getBaggageKeys().stream().filter( - (s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s)); - } - - private Stream whitelistedBaggageKeysWithValue(TraceContext context) { - if (context == null) { - return Stream.empty(); + private List> previousMdc() { + List> previousMdc = new ArrayList<>(); + List keys = new ArrayList<>(whitelistedBaggageKeys()); + keys.addAll(whitelistedPropagationKeys()); + for (String key : keys) { + previousMdc.add(new AbstractMap.SimpleEntry<>(key, MDC.get(key))); } - return whitelistedBaggageKeys().filter( - (s) -> StringUtils.hasText(ExtraFieldPropagation.get(context, s))); + return previousMdc; } - private Stream whitelistedPropagationKeys() { - return this.sleuthProperties.getPropagationKeys().stream().filter( - (s) -> this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(s)); - } - - private Stream whitelistedPropagationKeysWithValue(TraceContext context) { - if (context == null) { - return Stream.empty(); + private List whitelistedKeys(List keysToFilter) { + List keys = new ArrayList<>(); + for (String baggageKey : keysToFilter) { + if (this.sleuthSlf4jProperties.getWhitelistedMdcKeys().contains(baggageKey)) { + keys.add(baggageKey); + } } - return whitelistedPropagationKeys().filter( - (s) -> StringUtils.hasText(ExtraFieldPropagation.get(context, s))); + return keys; + } + + private List whitelistedBaggageKeys() { + return whitelistedKeys(this.sleuthProperties.getBaggageKeys()); + } + + private List whitelistedKeysWithValue(TraceContext context, + List keys) { + if (context == null) { + return Collections.EMPTY_LIST; + } + List nonEmpty = new ArrayList<>(); + for (String key : keys) { + if (StringUtils.hasText(ExtraFieldPropagation.get(context, key))) { + nonEmpty.add(key); + } + } + return nonEmpty; + } + + private List whitelistedBaggageKeysWithValue(TraceContext context) { + return whitelistedKeysWithValue(context, whitelistedBaggageKeys()); + } + + private List whitelistedPropagationKeys() { + return whitelistedKeys(this.sleuthProperties.getPropagationKeys()); + } + + private List whitelistedPropagationKeysWithValue(TraceContext context) { + return whitelistedKeysWithValue(context, whitelistedPropagationKeys()); } 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 3a8c0ddb3..5bae5e6b6 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 @@ -113,6 +113,27 @@ public class Slf4JSpanLoggerTest { assertThat(MDC.get("my-propagation")).isNullOrEmpty(); } + @Test + public void should_remove_entries_from_mdc_for_null_span_and_mdc_fields_set_directly() + throws Exception { + MDC.put("my-baggage", "my-value"); + MDC.put("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"); @@ -130,4 +151,46 @@ public class Slf4JSpanLoggerTest { assertThat(MDC.get("traceId")).isEqualTo("A"); } + // #1416 + @Test + public void should_clear_any_mdc_entries_when_their_keys_are_whitelisted() + throws Exception { + + Scope scope = this.slf4jScopeDecorator.decorateScope(this.span.context(), () -> { + }); + + MDC.put("my-baggage", "A"); + MDC.put("my-propagation", "B"); + + assertThat(MDC.get("my-baggage")).isEqualTo("A"); + assertThat(MDC.get("my-propagation")).isEqualTo("B"); + + scope.close(); + + assertThat(MDC.get("my-baggage")).isNullOrEmpty(); + assertThat(MDC.get("my-propagation")).isNullOrEmpty(); + } + + @Test + public void should_pick_previous_mdc_entries_when_their_keys_are_whitelisted() + throws Exception { + + MDC.put("my-baggage", "A1"); + MDC.put("my-propagation", "B1"); + + Scope scope = this.slf4jScopeDecorator.decorateScope(this.span.context(), () -> { + }); + + MDC.put("my-baggage", "A2"); + MDC.put("my-propagation", "B2"); + + assertThat(MDC.get("my-baggage")).isEqualTo("A2"); + assertThat(MDC.get("my-propagation")).isEqualTo("B2"); + + scope.close(); + + assertThat(MDC.get("my-baggage")).isEqualTo("A1"); + assertThat(MDC.get("my-propagation")).isEqualTo("B1"); + } + }