Fixed the header overriding for stream env post processor

fixes gh-1064
This commit is contained in:
Marcin Grzejszczak
2018-09-03 10:37:34 +02:00
parent e2de25f02d
commit d1dba30086
2 changed files with 56 additions and 9 deletions

View File

@@ -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<String> 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<String> 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<String, Object> map, MutablePropertySources propertySources,
String binder, int startIndex) {
String binder, int startIndex, List<String> 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<String, Object> 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<Object> 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<Object> values = ((MapPropertySource) source).getSource().values();
if (values.contains(header)) {
values.remove(header);
}
}
}
}

View File

@@ -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