From 4f3558a1651fea6a85118dc0dc20edfbe7cab977 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Wed, 28 Jul 2021 20:06:15 +0200 Subject: [PATCH] Fixed W3C tracestate and baggage setting without this change we're not setting tracestate whatsover. Baggage is also not being propagated properly (via the baggage field) with this change when tracestate is there in the headers it will be treated as baggage. When the context is propagated then tracestate will not be put to the baggage field fixes gh-2001 --- .../cloud/sleuth/BaggageManager.java | 4 + .../brave/bridge/BraveBaggageManager.java | 6 +- .../sleuth/brave/bridge/W3CPropagation.java | 90 +++++++++++++------ .../bridge/W3CBaggagePropagatorTest.java | 4 +- .../brave/bridge/W3CPropagationTest.java | 32 +++++++ 5 files changed, 108 insertions(+), 28 deletions(-) diff --git a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/BaggageManager.java b/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/BaggageManager.java index 861ae6843..541bdda4a 100644 --- a/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/BaggageManager.java +++ b/spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/BaggageManager.java @@ -18,6 +18,8 @@ package org.springframework.cloud.sleuth; import java.util.Map; +import org.springframework.lang.Nullable; + /** * Manages {@link BaggageInScope} entries. Upon retrieval / creation of a baggage entry * puts it in scope. Scope must be closed. @@ -38,6 +40,7 @@ public interface BaggageManager { * @param name baggage name * @return baggage or {@code null} if not present */ + @Nullable BaggageInScope getBaggage(String name); /** @@ -46,6 +49,7 @@ public interface BaggageManager { * @param name baggage name * @return baggage or {@code null} if not present */ + @Nullable BaggageInScope getBaggage(TraceContext traceContext, String name); /** diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveBaggageManager.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveBaggageManager.java index 9970342ff..c94d3b49e 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveBaggageManager.java +++ b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveBaggageManager.java @@ -47,7 +47,11 @@ public class BraveBaggageManager implements Closeable, BaggageManager { @Override public BaggageInScope getBaggage(TraceContext traceContext, String name) { - return new BraveBaggageInScope(BaggageField.getByName(BraveTraceContext.toBrave(traceContext), name)); + BaggageField baggageField = BaggageField.getByName(BraveTraceContext.toBrave(traceContext), name); + if (baggageField == null) { + return null; + } + return new BraveBaggageInScope(baggageField); } @Override diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagation.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagation.java index 0931b52e2..e5e5eb2ba 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagation.java +++ b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagation.java @@ -28,6 +28,8 @@ import java.util.Set; import java.util.stream.Collectors; import brave.baggage.BaggageField; +import brave.baggage.BaggagePropagation; +import brave.baggage.BaggagePropagationConfig; import brave.internal.baggage.BaggageFields; import brave.internal.propagation.StringPropagationAdapter; import brave.propagation.Propagation; @@ -37,6 +39,7 @@ import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.springframework.cloud.sleuth.BaggageInScope; +import org.springframework.util.StringUtils; import static java.util.Collections.singletonList; @@ -117,8 +120,11 @@ class W3CPropagation extends Propagation.Factory implements Propagation private final W3CBaggagePropagator baggagePropagator; + private final BraveBaggageManager braveBaggageManager; + W3CPropagation(BraveBaggageManager braveBaggageManager, List localFields) { this.baggagePropagator = new W3CBaggagePropagator(braveBaggageManager, localFields); + this.braveBaggageManager = braveBaggageManager; } @Override @@ -152,11 +158,25 @@ class W3CPropagation extends Propagation.Factory implements Propagation chars[TRACE_OPTION_OFFSET - 1] = TRACEPARENT_DELIMITER; copyTraceFlagsHexTo(chars, TRACE_OPTION_OFFSET, context); setter.put(carrier, TRACE_PARENT, new String(chars, 0, TRACEPARENT_HEADER_SIZE)); - // Add baggage + addTraceState(setter, context, carrier); this.baggagePropagator.injector(setter).inject(context, carrier); }; } + private void addTraceState(Setter setter, TraceContext context, R carrier) { + if (carrier != null) { + BaggageInScope baggage = this.braveBaggageManager.getBaggage(BraveTraceContext.fromBrave(context), + TRACE_STATE); + if (baggage == null) { + return; + } + String traceState = baggage.get(BraveTraceContext.fromBrave(context)); + if (StringUtils.hasText(traceState)) { + setter.put(carrier, TRACE_STATE, traceState); + } + } + } + private String padLeftWithZeros(String string, int length) { if (string.length() >= length) { return string; @@ -202,14 +222,15 @@ class W3CPropagation extends Propagation.Factory implements Propagation } TraceContextOrSamplingFlags context(TraceContext contextFromParentHeader, String traceStateHeader) { - if (traceStateHeader == null || traceStateHeader.isEmpty()) { + if (!StringUtils.hasText(traceStateHeader)) { return TraceContextOrSamplingFlags.create(contextFromParentHeader); } try { return TraceContextOrSamplingFlags - .create(TraceContext.newBuilder().traceId(contextFromParentHeader.traceId()) + .newBuilder(TraceContext.newBuilder().traceId(contextFromParentHeader.traceId()) .traceIdHigh(contextFromParentHeader.traceIdHigh()).spanId(contextFromParentHeader.spanId()) - .sampled(contextFromParentHeader.sampled()).shared(true).build()); + .sampled(contextFromParentHeader.sampled()).shared(true).build()) + .build(); } catch (IllegalArgumentException e) { logger.info("Unparseable tracestate header. Returning span context without state."); @@ -279,6 +300,10 @@ class W3CBaggagePropagator { private static final Log log = LogFactory.getLog(W3CBaggagePropagator.class); + private static final String TRACE_STATE = "tracestate"; + + private static final BaggageField TRACE_STATE_BAGGAGE = BaggageField.create(TRACE_STATE); + private static final String FIELD = "baggage"; private static final List FIELDS = singletonList(FIELD); @@ -292,6 +317,15 @@ class W3CBaggagePropagator { this.localFields = localFields; } + private BaggagePropagation.FactoryBuilder factory() { + return BaggagePropagation.newFactoryBuilder(new Propagation.Factory() { + @Override + public Propagation create(Propagation.KeyFactory keyFactory) { + return null; + } + }); + } + public List keys() { return FIELDS; } @@ -307,6 +341,9 @@ class W3CBaggagePropagator { String[] strings = this.localFields.toArray(new String[0]); Map filtered = extra.toMapFilteringFieldNames(strings); for (Map.Entry entry : filtered.entrySet()) { + if (TRACE_STATE.equalsIgnoreCase(entry.getKey())) { + continue; + } headerContent.append(entry.getKey()).append("=").append(entry.getValue()); // TODO: [OTEL] No metadata support // String metadataValue = entry.getEntryMetadata().getValue(); @@ -322,28 +359,37 @@ class W3CBaggagePropagator { }; } - TraceContextOrSamplingFlags contextWithBaggage(R carrier, TraceContextOrSamplingFlags context, + TraceContextOrSamplingFlags contextWithBaggage(R carrier, TraceContextOrSamplingFlags flags, Propagation.Getter getter) { + BaggagePropagation.FactoryBuilder factoryBuilder = factory(); + String traceState = getter.get(carrier, TRACE_STATE); + boolean hasTraceState = StringUtils.hasText(traceState); + if (hasTraceState) { + factoryBuilder = factoryBuilder + .add(BaggagePropagationConfig.SingleBaggageField.remote(TRACE_STATE_BAGGAGE)); + } String baggageHeader = getter.get(carrier, FIELD); - if (baggageHeader == null) { - return context; + List> pairs = baggageHeader == null || baggageHeader.isEmpty() + ? Collections.emptyList() : addBaggageToContext(baggageHeader); + Set names = pairs.stream().map(e -> e.getKey().name()).collect(Collectors.toSet()); + for (String name : names) { + factoryBuilder = factoryBuilder.add(BaggagePropagationConfig.SingleBaggageField + .remote(((BraveBaggageInScope) this.braveBaggageManager.createBaggage(name)).unwrap())); } - if (baggageHeader.isEmpty()) { - return context; + TraceContext decoratedContext = factoryBuilder.build().decorate(flags.context()); + if (hasTraceState) { + BaggageInScope baggageInScope = this.braveBaggageManager.createBaggage(TRACE_STATE); + baggageInScope.set(new BraveTraceContext(decoratedContext), traceState); } - TraceContextOrSamplingFlags.Builder builder = context.toBuilder(); - List> pairs = addBaggageToContext(baggageHeader, builder); - TraceContextOrSamplingFlags built = builder.build(); pairs.forEach(e -> { - BaggageField baggage = e.getKey(); - baggage.updateValue(built, e.getValue()); + BaggageField baggage = ((BraveBaggageInScope) e.getKey()).unwrap(); + baggage.updateValue(decoratedContext, e.getValue()); }); - return built; + return TraceContextOrSamplingFlags.create(decoratedContext); } - List> addBaggageToContext(String baggageHeader, - TraceContextOrSamplingFlags.Builder builder) { - List> pairs = new ArrayList<>(); + List> addBaggageToContext(String baggageHeader) { + List> pairs = new ArrayList<>(); String[] entries = baggageHeader.split(","); for (String entry : entries) { int beginningOfMetadata = entry.indexOf(";"); @@ -356,8 +402,7 @@ class W3CBaggagePropagator { String key = keyAndValue[i].trim(); String value = keyAndValue[i + 1].trim(); BaggageInScope baggage = this.braveBaggageManager.createBaggage(key); - BaggageField field = ((BraveBaggageInScope) baggage).unwrap(); - pairs.add(new AbstractMap.SimpleEntry<>(field, value)); + pairs.add(new AbstractMap.SimpleEntry<>(baggage, value)); } catch (Exception e) { if (log.isDebugEnabled()) { @@ -367,11 +412,6 @@ class W3CBaggagePropagator { } } } - // TODO: [OTEL] Magic number for max dynamic entries - builder.addExtra(BaggageFields - .newFactory(pairs.stream().map(AbstractMap.SimpleEntry::getKey).collect(Collectors.toList()), - pairs.size() * 2) - .create()); return pairs; } diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CBaggagePropagatorTest.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CBaggagePropagatorTest.java index ed5887c58..b997013c7 100644 --- a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CBaggagePropagatorTest.java +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CBaggagePropagatorTest.java @@ -51,7 +51,7 @@ class W3CBaggagePropagatorTest { TraceContextOrSamplingFlags contextWithBaggage = propagator.contextWithBaggage(carrier, context, Map::get); - assertThat(contextWithBaggage).isSameAs(context); + assertThat(contextWithBaggage).isEqualTo(context); } @Test @@ -62,7 +62,7 @@ class W3CBaggagePropagatorTest { TraceContextOrSamplingFlags contextWithBaggage = propagator.contextWithBaggage(carrier, context, Map::get); - assertThat(contextWithBaggage).isSameAs(context); + assertThat(contextWithBaggage).isEqualTo(context); } @Test diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagationTest.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagationTest.java index 265e0e76b..58ca2c048 100644 --- a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagationTest.java +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CPropagationTest.java @@ -17,11 +17,14 @@ package org.springframework.cloud.sleuth.brave.bridge; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; +import brave.baggage.BaggageField; +import brave.internal.baggage.BaggageFields; import brave.propagation.Propagation; import brave.propagation.TraceContext; import brave.propagation.TraceContextOrSamplingFlags; @@ -52,6 +55,8 @@ class W3CPropagationTest { private static final String TRACESTATE_NOT_DEFAULT_ENCODING_WITH_SPACES = "bar=baz , foo=bar"; + private static final String TRACESTATE_HEADER = "sappp=CwAAmEnGj0gThK52TCXZ270X8nBhc3Nwb3J0LWFwcABQT1NU"; + private final W3CPropagation w3CPropagation = new W3CPropagation(new BraveBaggageManager(), new ArrayList<>()); @Test @@ -81,6 +86,20 @@ class W3CPropagationTest { assertThat(carrier).containsExactly(entry(TRACE_PARENT, TRACEPARENT_HEADER_SAMPLED)); } + @Test + void inject_tracestate() { + Map carrier = new LinkedHashMap<>(); + BaggageField traceStateField = BaggageField.create(TRACE_STATE); + BaggageField mybaggageField = BaggageField.create("mybaggage"); + TraceContext traceContext = sampledTraceContext() + .addExtra(BaggageFields.newFactory(Arrays.asList(traceStateField, mybaggageField), 2).create()).build(); + traceStateField.updateValue(traceContext, TRACESTATE_HEADER); + mybaggageField.updateValue(traceContext, "mybaggagevalue"); + w3CPropagation.injector((ignored, key, value) -> carrier.put(key, value)).inject(traceContext, carrier); + assertThat(carrier).containsEntry("baggage", "mybaggage=mybaggagevalue").containsEntry("tracestate", + "sappp=CwAAmEnGj0gThK52TCXZ270X8nBhc3Nwb3J0LWFwcABQT1NU"); + } + @Test void inject_NotSampledContext() { Map carrier = new LinkedHashMap<>(); @@ -161,6 +180,19 @@ class W3CPropagationTest { .isEqualTo(sharedTraceContext().sampled(false).build()); } + @Test + void extract_tracestate_shouldNotBePartOfBaggage() { + Map carrier = new LinkedHashMap<>(); + carrier.put(TRACE_PARENT, TRACEPARENT_HEADER_NOT_SAMPLED); + carrier.put(TRACE_STATE, TRACESTATE_HEADER); + carrier.put("baggage", "mybaggage=mybaggagevalue"); + + TraceContext context = w3CPropagation.extractor(getter).extract(carrier).context(); + + assertThat(BaggageField.getByName(context, TRACE_STATE)).isNotNull(); + assertThat(BaggageField.getByName(context, "mybaggage").getValue(context)).isEqualTo("mybaggagevalue"); + } + @Test void extract_EmptyHeader() { Map invalidHeaders = new LinkedHashMap<>();