diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjector.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjector.java index 7a6856c0a..17db826ab 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjector.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjector.java @@ -33,28 +33,28 @@ public class HeaderBasedMessagingInjector implements MessagingSpanTextMapInjecto } return; } - addHeaders(span, carrier); + addHeaders(map, span, carrier); } private boolean isSampled(Map initialMessage, String sampledHeaderName) { return Span.SPAN_SAMPLED.equals(initialMessage.get(sampledHeaderName)); } - private void addHeaders(Span span, SpanTextMap textMap) { - addHeader(textMap, TraceMessageHeaders.TRACE_ID_NAME, span.traceIdString()); - addHeader(textMap, TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(span.getSpanId())); + private void addHeaders(Map map, Span span, SpanTextMap textMap) { + addHeader(map, textMap, TraceMessageHeaders.TRACE_ID_NAME, span.traceIdString()); + addHeader(map, textMap, TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(span.getSpanId())); if (span.isExportable()) { addAnnotations(this.traceKeys, textMap, span); Long parentId = getFirst(span.getParents()); if (parentId != null) { - addHeader(textMap, TraceMessageHeaders.PARENT_ID_NAME, Span.idToHex(parentId)); + addHeader(map, textMap, TraceMessageHeaders.PARENT_ID_NAME, Span.idToHex(parentId)); } - addHeader(textMap, TraceMessageHeaders.SPAN_NAME_NAME, span.getName()); - addHeader(textMap, TraceMessageHeaders.PROCESS_ID_NAME, span.getProcessId()); - addHeader(textMap, TraceMessageHeaders.SAMPLED_NAME, Span.SPAN_SAMPLED); + addHeader(map, textMap, TraceMessageHeaders.SPAN_NAME_NAME, span.getName()); + addHeader(map, textMap, TraceMessageHeaders.PROCESS_ID_NAME, span.getProcessId()); + addHeader(map, textMap, TraceMessageHeaders.SAMPLED_NAME, Span.SPAN_SAMPLED); } else { - addHeader(textMap, TraceMessageHeaders.SAMPLED_NAME, Span.SPAN_NOT_SAMPLED); + addHeader(map, textMap, TraceMessageHeaders.SAMPLED_NAME, Span.SPAN_NOT_SAMPLED); } for (Map.Entry entry : span.baggageItems()) { textMap.put(prefixedKey(entry.getKey()), entry.getValue()); @@ -92,8 +92,8 @@ public class HeaderBasedMessagingInjector implements MessagingSpanTextMapInjecto } } - private void addHeader(SpanTextMap textMap, String name, String value) { - if (StringUtils.hasText(value)) { + private void addHeader(Map map, SpanTextMap textMap, String name, String value) { + if (StringUtils.hasText(value) && !map.containsKey(name)) { textMap.put(name, value); } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjector.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjector.java index 8788ff2de..6c085be1a 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjector.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjector.java @@ -4,6 +4,7 @@ import java.util.Map; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.SpanTextMap; +import org.springframework.cloud.sleuth.util.TextMapUtil; import org.springframework.util.StringUtils; /** @@ -17,15 +18,16 @@ public class ZipkinHttpSpanInjector implements HttpSpanInjector { private static final String HEADER_DELIMITER = "-"; @Override - public void inject(Span span, SpanTextMap carrier) { - setHeader(carrier, Span.TRACE_ID_NAME, span.traceIdString()); - setIdHeader(carrier, Span.SPAN_ID_NAME, span.getSpanId()); - setHeader(carrier, Span.SAMPLED_NAME, span.isExportable() ? Span.SPAN_SAMPLED : Span.SPAN_NOT_SAMPLED); - setHeader(carrier, Span.SPAN_NAME_NAME, span.getName()); - setIdHeader(carrier, Span.PARENT_ID_NAME, getParentId(span)); - setHeader(carrier, Span.PROCESS_ID_NAME, span.getProcessId()); + public void inject(Span span, SpanTextMap map) { + Map carrier = TextMapUtil.asMap(map); + setHeader(map, carrier, Span.TRACE_ID_NAME, span.traceIdString()); + setIdHeader(map, carrier, Span.SPAN_ID_NAME, span.getSpanId()); + setHeader(map, carrier, Span.SAMPLED_NAME, span.isExportable() ? Span.SPAN_SAMPLED : Span.SPAN_NOT_SAMPLED); + setHeader(map, carrier, Span.SPAN_NAME_NAME, span.getName()); + setIdHeader(map, carrier, Span.PARENT_ID_NAME, getParentId(span)); + setHeader(map, carrier, Span.PROCESS_ID_NAME, span.getProcessId()); for (Map.Entry entry : span.baggageItems()) { - carrier.put(prefixedKey(entry.getKey()), entry.getValue()); + map.put(prefixedKey(entry.getKey()), entry.getValue()); } } @@ -40,15 +42,15 @@ public class ZipkinHttpSpanInjector implements HttpSpanInjector { return !span.getParents().isEmpty() ? span.getParents().get(0) : null; } - private void setIdHeader(SpanTextMap carrier, String name, Long value) { + private void setIdHeader(SpanTextMap map, Map carrier, String name, Long value) { if (value != null) { - setHeader(carrier, name, Span.idToHex(value)); + setHeader(map, carrier, name, Span.idToHex(value)); } } - private void setHeader(SpanTextMap carrier, String name, String value) { - if (StringUtils.hasText(value)) { - carrier.put(name, value); + private void setHeader(SpanTextMap map, Map carrier, String name, String value) { + if (StringUtils.hasText(value) && !carrier.containsKey(name)) { + map.put(name, value); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java new file mode 100644 index 000000000..4ee13d5e4 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java @@ -0,0 +1,62 @@ +package org.springframework.cloud.sleuth.instrument.messaging; + +import java.util.AbstractMap; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; + +import org.junit.Test; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanTextMap; +import org.springframework.cloud.sleuth.TraceKeys; + +import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then; + +/** + * @author Marcin Grzejszczak + */ +public class HeaderBasedMessagingInjectorTests { + + HeaderBasedMessagingInjector injector = new HeaderBasedMessagingInjector(new TraceKeys()); + + @SuppressWarnings("unchecked") + @Test + public void should_not_override_already_existing_headers() throws Exception { + Span span = Span.builder() + .spanId(1L) + .traceId(2L) + .parent(3L) + .baggage("foo", "bar") + .name("span") + .exportable(true) + .build(); + Map holder = new HashMap<>(); + final SpanTextMap map = textMap(holder); + holder.put(TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(10L)); + holder.put(TraceMessageHeaders.TRACE_ID_NAME, Span.idToHex(20L)); + holder.put(TraceMessageHeaders.PARENT_ID_NAME, Span.idToHex(30L)); + holder.put(TraceMessageHeaders.SPAN_NAME_NAME, "anotherSpan"); + + injector.inject(span, map); + + then(map) + .contains(new AbstractMap.SimpleEntry(TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(10L))) + .contains(new AbstractMap.SimpleEntry(TraceMessageHeaders.TRACE_ID_NAME, Span.idToHex(20L))) + .contains(new AbstractMap.SimpleEntry(TraceMessageHeaders.PARENT_ID_NAME, Span.idToHex(30L))) + .contains(new AbstractMap.SimpleEntry(TraceMessageHeaders.SPAN_NAME_NAME, "anotherSpan")) + .contains(new AbstractMap.SimpleEntry("baggage_foo", "bar")); + } + + private SpanTextMap textMap(Map textMap) { + return new SpanTextMap() { + @Override public Iterator> iterator() { + return textMap.entrySet().iterator(); + } + + @Override public void put(String key, String value) { + textMap.put(key, value); + } + }; + } + +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjectorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjectorTests.java new file mode 100644 index 000000000..4110eaeb1 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanInjectorTests.java @@ -0,0 +1,60 @@ +package org.springframework.cloud.sleuth.instrument.web; + +import java.util.AbstractMap; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; + +import org.junit.Test; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanTextMap; + +import static org.assertj.core.api.BDDAssertions.then; + +/** + * @author Marcin Grzejszczak + */ +public class ZipkinHttpSpanInjectorTests { + + ZipkinHttpSpanInjector injector = new ZipkinHttpSpanInjector(); + + @SuppressWarnings("unchecked") + @Test + public void should_not_override_already_existing_headers() throws Exception { + Span span = Span.builder() + .spanId(1L) + .traceId(2L) + .parent(3L) + .baggage("foo", "bar") + .name("span") + .build(); + Map holder = new HashMap<>(); + final SpanTextMap map = textMap(holder); + holder.put(Span.SPAN_ID_NAME, Span.idToHex(10L)); + holder.put(Span.TRACE_ID_NAME, Span.idToHex(20L)); + holder.put(Span.PARENT_ID_NAME, Span.idToHex(30L)); + holder.put(Span.SPAN_NAME_NAME, "anotherSpan"); + + injector.inject(span, map); + + then(map) + .contains(new AbstractMap.SimpleEntry(Span.SPAN_ID_NAME, Span.idToHex(10L))) + .contains(new AbstractMap.SimpleEntry(Span.TRACE_ID_NAME, Span.idToHex(20L))) + .contains(new AbstractMap.SimpleEntry(Span.PARENT_ID_NAME, Span.idToHex(30L))) + .contains(new AbstractMap.SimpleEntry(Span.SPAN_NAME_NAME, "anotherSpan")) + .contains(new AbstractMap.SimpleEntry("baggage-foo", "bar")); + } + + private SpanTextMap textMap(Map textMap) { + return new SpanTextMap() { + @Override public Iterator> iterator() { + return textMap.entrySet().iterator(); + } + + @Override public void put(String key, String value) { + textMap.put(key, value); + } + }; + } + +} \ No newline at end of file