diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessor.java index b94340317..7783edc19 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessor.java @@ -17,9 +17,12 @@ package org.springframework.cloud.sleuth.autoconfig; import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Properties; @@ -64,8 +67,10 @@ class TraceStreamEnvironmentPostProcessor implements EnvironmentPostProcessor { try { for (Resource resource : getAllSpringBinders(resolver)) { for (String binderType : parseBinderConfigurations(resource)) { + List existingHeaders = existingHeaders(environment, binderType); int startIndex = findStartIndex(environment, binderType); - addHeaders(map, environment.getPropertySources(), binderType, startIndex); + startIndex = startIndex + existingHeaders.size(); + addHeaders(map, environment.getPropertySources(), binderType, startIndex, existingHeaders); } } } @@ -81,6 +86,15 @@ class TraceStreamEnvironmentPostProcessor implements EnvironmentPostProcessor { .getResources("classpath*:META-INF/spring.binders"); } + private List existingHeaders(ConfigurableEnvironment environment, String binder) { + String prefix = "spring.cloud.stream." + binder + ".binder.headers"; + String oldHeaders = environment.getProperty(prefix); + if (oldHeaders != null) { + return Arrays.asList(oldHeaders.split(",")); + } + return new ArrayList<>(); + } + private int findStartIndex(ConfigurableEnvironment environment, String binder) { String prefix = "spring.cloud.stream." + binder + ".binder.HEADERS"; int i = 0; @@ -130,16 +144,30 @@ class TraceStreamEnvironmentPostProcessor implements EnvironmentPostProcessor { } private void addHeaders(Map map, MutablePropertySources propertySources, - String binder, int startIndex) { + String binder, int startIndex, List existingHeaders) { String stem = "spring.cloud.stream." + binder + ".binder.HEADERS"; + for (int i = 0; i < existingHeaders.size(); i++) { + String header = existingHeaders.get(i); + if (!hasHeaderKey(propertySources, header)) { + putHeader(map, stem, i, header); + } + } for (int i = 0; i < HEADERS.length; i++) { - if (!hasTracingHeadersValue(propertySources, HEADERS[i])) { - map.put(stem + "[" + (i + startIndex) + "]", HEADERS[i]); + boolean hasHeader = hasHeaderKey(propertySources, HEADERS[i]); + if (!hasHeader) { + putHeader(map, stem, i + startIndex, HEADERS[i]); + } else if (!existingHeaders.isEmpty() && hasHeader) { + removeEntryWithHeader(propertySources, HEADERS[i]); + putHeader(map, stem, i + startIndex, HEADERS[i]); } } } - private boolean hasTracingHeadersValue(MutablePropertySources propertySources, String header) { + private void putHeader(Map map, String stem, int i2, String header2) { + map.put(stem + "[" + (i2) + "]", header2); + } + + private boolean hasHeaderKey(MutablePropertySources propertySources, String header) { PropertySource source = propertySources.get(PROPERTY_SOURCE_NAME); if (source instanceof MapPropertySource) { Collection values = ((MapPropertySource) source).getSource().values(); @@ -148,4 +176,14 @@ class TraceStreamEnvironmentPostProcessor implements EnvironmentPostProcessor { return false; } + private void removeEntryWithHeader(MutablePropertySources propertySources, String header) { + PropertySource source = propertySources.get(PROPERTY_SOURCE_NAME); + if (source instanceof MapPropertySource) { + Collection values = ((MapPropertySource) source).getSource().values(); + if (values.contains(header)) { + values.remove(header); + } + } + } + } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessorTests.java index f971804d3..e99897596 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceStreamEnvironmentPostProcessorTests.java @@ -26,11 +26,10 @@ import org.junit.Test; import org.springframework.boot.SpringApplication; import org.springframework.boot.test.util.EnvironmentTestUtils; import org.springframework.cloud.sleuth.instrument.messaging.TraceMessageHeaders; -import org.springframework.core.env.ConfigurableEnvironment; -import org.springframework.core.env.StandardEnvironment; import org.springframework.core.io.ClassPathResource; import org.springframework.core.io.Resource; import org.springframework.core.io.support.PathMatchingResourcePatternResolver; +import org.springframework.mock.env.MockEnvironment; import static org.assertj.core.api.Assertions.assertThat; @@ -50,14 +49,24 @@ public class TraceStreamEnvironmentPostProcessorTests { return Collections.singleton("test"); } }; - private ConfigurableEnvironment environment = new StandardEnvironment(); + private MockEnvironment environment = new MockEnvironment(); @Test public void should_append_tracing_headers() { + this.environment.setProperty("spring.cloud.stream.test.binder.headers", "foo,bar,baz"); postProcess(); assertThat(this.environment .getProperty("spring.cloud.stream.test.binder.HEADERS[0]")) - .isEqualTo(TraceMessageHeaders.SPAN_ID_NAME); + .isEqualTo("foo"); + assertThat(this.environment + .getProperty("spring.cloud.stream.test.binder.HEADERS[1]")) + .isEqualTo("bar"); + assertThat(this.environment + .getProperty("spring.cloud.stream.test.binder.HEADERS[2]")) + .isEqualTo("baz"); + assertThat(this.environment + .getProperty("spring.cloud.stream.test.binder.HEADERS[3]")) + .isEqualTo(TraceMessageHeaders.SPAN_ID_NAME); } @Test