From 8f3ed392bc28ef36d027822d10b1f66e9fc14ea1 Mon Sep 17 00:00:00 2001 From: Vastin <3690049+vastin@users.noreply.github.com> Date: Wed, 8 May 2024 05:34:45 +0000 Subject: [PATCH] Store additional fields in context instead of baggage. --- .../awsxray/propagator/AwsXrayPropagator.java | 73 ++++++++++--------- .../propagator/AwsXrayPropagatorTest.java | 20 +++-- 2 files changed, 53 insertions(+), 40 deletions(-) diff --git a/aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java b/aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java index ef7a4cf8a..0d1ab7d06 100644 --- a/aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java +++ b/aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java @@ -16,6 +16,7 @@ import io.opentelemetry.api.trace.TraceId; import io.opentelemetry.api.trace.TraceState; import io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.propagation.TextMapGetter; import io.opentelemetry.context.propagation.TextMapPropagator; import io.opentelemetry.context.propagation.TextMapSetter; @@ -68,6 +69,10 @@ public final class AwsXrayPropagator implements TextMapPropagator { private static final char IS_SAMPLED = '1'; private static final char NOT_SAMPLED = '0'; + private static final String XRAY_HEADER_ADDITIONAL_FIELDS_KEY_NAME = "XrayHeaderAdditionalFields"; + static final ContextKey XRAY_HEADER_ADDITIONAL_FIELDS_KEY = + ContextKey.named(XRAY_HEADER_ADDITIONAL_FIELDS_KEY_NAME); + private static final List FIELDS = Collections.singletonList(TRACE_HEADER_KEY); private static final AwsXrayPropagator INSTANCE = new AwsXrayPropagator(); @@ -85,6 +90,7 @@ public List fields() { return FIELDS; } + @SuppressWarnings("null") @Override public void inject(Context context, @Nullable C carrier, TextMapSetter setter) { if (context == null) { @@ -126,34 +132,35 @@ public void inject(Context context, @Nullable C carrier, TextMapSetter se .append(KV_DELIMITER) .append(samplingFlag); - Baggage baggage = Baggage.fromContext(context); - // Truncate baggage to 256 chars per X-Ray spec. - baggage.forEach( - new BiConsumer() { - - private int baggageWrittenBytes; - - @Override - public void accept(String key, BaggageEntry entry) { - if (key.equals(TRACE_ID_KEY) - || key.equals(PARENT_ID_KEY) - || key.equals(SAMPLED_FLAG_KEY)) { - return; + Baggage fields = context.get(XRAY_HEADER_ADDITIONAL_FIELDS_KEY); + // Truncate fields to 256 chars per X-Ray spec. + if (fields != null) { + fields.forEach( + new BiConsumer() { + + private int baggageWrittenBytes; + + @Override + public void accept(String key, BaggageEntry entry) { + if (key.equals(TRACE_ID_KEY) + || key.equals(PARENT_ID_KEY) + || key.equals(SAMPLED_FLAG_KEY)) { + return; + } + // Size is key/value pair, excludes delimiter. + int size = key.length() + entry.getValue().length() + 1; + if (baggageWrittenBytes + size > 256) { + return; + } + traceHeader + .append(TRACE_HEADER_DELIMITER) + .append(key) + .append(KV_DELIMITER) + .append(entry.getValue()); + baggageWrittenBytes += size; } - // Size is key/value pair, excludes delimiter. - int size = key.length() + entry.getValue().length() + 1; - if (baggageWrittenBytes + size > 256) { - return; - } - traceHeader - .append(TRACE_HEADER_DELIMITER) - .append(key) - .append(KV_DELIMITER) - .append(entry.getValue()); - baggageWrittenBytes += size; - } - }); - + }); + } setter.set(carrier, TRACE_HEADER_KEY, traceHeader.toString()); } @@ -180,7 +187,7 @@ private static Context getContextFromHeader( String spanId = SpanId.getInvalid(); Boolean isSampled = false; - BaggageBuilder baggage = null; + BaggageBuilder fields = null; int baggageReadBytes = 0; int pos = 0; @@ -211,10 +218,10 @@ private static Context getContextFromHeader( } else if (trimmedPart.startsWith(SAMPLED_FLAG_KEY)) { isSampled = parseTraceFlag(value); } else if (baggageReadBytes + trimmedPart.length() <= 256) { - if (baggage == null) { - baggage = Baggage.builder(); + if (fields == null) { + fields = Baggage.builder(); } - baggage.put(trimmedPart.substring(0, equalsIndex), value); + fields.put(trimmedPart.substring(0, equalsIndex), value); baggageReadBytes += trimmedPart.length(); } } @@ -241,8 +248,8 @@ private static Context getContextFromHeader( if (spanContext.isValid()) { context = context.with(Span.wrap(spanContext)); } - if (baggage != null) { - context = context.with(baggage.build()); + if (fields != null) { + context = context.with(XRAY_HEADER_ADDITIONAL_FIELDS_KEY, fields.build()); } return context; } diff --git a/aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagatorTest.java b/aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagatorTest.java index 2637e78e5..13f57841f 100644 --- a/aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagatorTest.java +++ b/aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagatorTest.java @@ -6,6 +6,7 @@ package io.opentelemetry.contrib.awsxray.propagator; import static io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator.TRACE_HEADER_KEY; +import static io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator.XRAY_HEADER_ADDITIONAL_FIELDS_KEY; import static org.assertj.core.api.Assertions.assertThat; import io.opentelemetry.api.baggage.Baggage; @@ -90,7 +91,7 @@ void inject_NotSampledContext() { } @Test - void inject_WithBaggage() { + void inject_WithAdditionalFields() { Map carrier = new LinkedHashMap<>(); subject.inject( withSpanContext( @@ -98,6 +99,7 @@ void inject_WithBaggage() { TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()), Context.current()) .with( + XRAY_HEADER_ADDITIONAL_FIELDS_KEY, Baggage.builder() .put("cat", "meow") .put("dog", "bark") @@ -116,7 +118,7 @@ void inject_WithBaggage() { } @Test - void inject_WithBaggage_LimitTruncates() { + void inject_WithAdditionalFields_LimitTruncates() { Map carrier = new LinkedHashMap<>(); // Limit is 256 characters for all baggage. We add a 254-character key/value pair and a // 3 character key value pair. @@ -133,7 +135,7 @@ void inject_WithBaggage_LimitTruncates() { SpanContext.create( TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()), Context.current()) - .with(baggage), + .with(XRAY_HEADER_ADDITIONAL_FIELDS_KEY, baggage), carrier, SETTER); @@ -244,11 +246,13 @@ void extract_AdditionalFields() { .isEqualTo( SpanContext.createFromRemoteParent( TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault())); - assertThat(Baggage.fromContext(context).getEntryValue("Foo")).isEqualTo("Bar"); + Baggage additionalFields = context.get(XRAY_HEADER_ADDITIONAL_FIELDS_KEY); + assertThat(additionalFields.getEntryValue("Foo")).isEqualTo("Bar"); + assertThat(Baggage.fromContext(context)).isEqualTo(Baggage.empty()); } @Test - void extract_Baggage_LimitTruncates() { + void extract_AdditionalFields_LimitTruncates() { // Limit is 256 characters for all baggage. We add a 254-character key/value pair and a // 3 character key value pair. String key1 = Stream.generate(() -> "a").limit(252).collect(Collectors.joining()); @@ -274,8 +278,10 @@ void extract_Baggage_LimitTruncates() { .isEqualTo( SpanContext.createFromRemoteParent( TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault())); - assertThat(Baggage.fromContext(context).getEntryValue(key1)).isEqualTo(value1); - assertThat(Baggage.fromContext(context).getEntryValue(key2)).isNull(); + Baggage additionalFields = context.get(XRAY_HEADER_ADDITIONAL_FIELDS_KEY); + assertThat(additionalFields.getEntryValue(key1)).isEqualTo(value1); + assertThat(additionalFields.getEntryValue(key2)).isNull(); + assertThat(Baggage.fromContext(context)).isEqualTo(Baggage.empty()); } @Test