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<>();