Skip to content

Commit e0cdc56

Browse files
liustveliustve
authored
Support for Lineage in XRay Trace Header and removing additional Baggage from being added (#1671)
1 parent 6e411e6 commit e0cdc56

File tree

3 files changed

+261
-92
lines changed

3 files changed

+261
-92
lines changed

‎aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java

+85-39
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,10 @@
55

66
package io.opentelemetry.contrib.awsxray.propagator;
77

8+
import static io.opentelemetry.api.internal.OtelEncodingUtils.isValidBase16String;
9+
810
import io.opentelemetry.api.baggage.Baggage;
911
import io.opentelemetry.api.baggage.BaggageBuilder;
10-
import io.opentelemetry.api.baggage.BaggageEntry;
1112
import io.opentelemetry.api.internal.StringUtils;
1213
import io.opentelemetry.api.trace.Span;
1314
import io.opentelemetry.api.trace.SpanContext;
@@ -21,7 +22,7 @@
2122
import io.opentelemetry.context.propagation.TextMapSetter;
2223
import java.util.Collections;
2324
import java.util.List;
24-
import java.util.function.BiConsumer;
25+
import java.util.Set;
2526
import java.util.logging.Logger;
2627
import javax.annotation.Nullable;
2728

@@ -68,6 +69,17 @@ public final class AwsXrayPropagator implements TextMapPropagator {
6869
private static final char IS_SAMPLED = '1';
6970
private static final char NOT_SAMPLED = '0';
7071

72+
private static final String LINEAGE_KEY = "Lineage";
73+
private static final char LINEAGE_DELIMITER = ':';
74+
private static final int LINEAGE_MAX_LENGTH = 18;
75+
private static final int LINEAGE_MIN_LENGTH = 12;
76+
private static final int LINEAGE_HASH_LENGTH = 8;
77+
private static final int LINEAGE_MAX_COUNTER1 = 32767;
78+
private static final int LINEAGE_MAX_COUNTER2 = 255;
79+
private static final int LINEAGE_MIN_COUNTER = 0;
80+
private static final String INVALID_LINEAGE = "-1:11111111:0";
81+
private static final int NUM_OF_LINEAGE_DELIMITERS = 2;
82+
7183
private static final List<String> FIELDS = Collections.singletonList(TRACE_HEADER_KEY);
7284

7385
private static final AwsXrayPropagator INSTANCE = new AwsXrayPropagator();
@@ -127,34 +139,19 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
127139
.append(samplingFlag);
128140

129141
Baggage baggage = Baggage.fromContext(context);
130-
// Truncate baggage to 256 chars per X-Ray spec.
131-
baggage.forEach(
132-
new BiConsumer<String, BaggageEntry>() {
133-
134-
private int baggageWrittenBytes;
135-
136-
@Override
137-
public void accept(String key, BaggageEntry entry) {
138-
if (key.equals(TRACE_ID_KEY)
139-
|| key.equals(PARENT_ID_KEY)
140-
|| key.equals(SAMPLED_FLAG_KEY)) {
141-
return;
142-
}
143-
// Size is key/value pair, excludes delimiter.
144-
int size = key.length() + entry.getValue().length() + 1;
145-
if (baggageWrittenBytes + size > 256) {
146-
return;
147-
}
148-
traceHeader
149-
.append(TRACE_HEADER_DELIMITER)
150-
.append(key)
151-
.append(KV_DELIMITER)
152-
.append(entry.getValue());
153-
baggageWrittenBytes += size;
154-
}
155-
});
156-
157-
setter.set(carrier, TRACE_HEADER_KEY, traceHeader.toString());
142+
String lineageHeader = baggage.getEntryValue(LINEAGE_KEY);
143+
144+
if (lineageHeader != null) {
145+
traceHeader
146+
.append(TRACE_HEADER_DELIMITER)
147+
.append(LINEAGE_KEY)
148+
.append(KV_DELIMITER)
149+
.append(lineageHeader);
150+
}
151+
152+
// add 256 character truncation
153+
String truncatedTraceHeader = traceHeader.substring(0, Math.min(traceHeader.length(), 256));
154+
setter.set(carrier, TRACE_HEADER_KEY, truncatedTraceHeader);
158155
}
159156

160157
@Override
@@ -183,10 +180,20 @@ private static <C> Context getContextFromHeader(
183180

184181
String traceId = TraceId.getInvalid();
185182
String spanId = SpanId.getInvalid();
183+
String lineageHeader;
186184
Boolean isSampled = false;
187185

188-
BaggageBuilder baggage = null;
189-
int baggageReadBytes = 0;
186+
Baggage contextBaggage = Baggage.fromContext(context);
187+
BaggageBuilder baggageBuilder = Baggage.builder();
188+
Set<String> baggageMap = contextBaggage.asMap().keySet();
189+
190+
// Copying baggage over to new Baggage object to add Lineage key
191+
for (String baggageKey : baggageMap) {
192+
String baggageValue = contextBaggage.getEntryValue(baggageKey);
193+
if (baggageValue != null) {
194+
baggageBuilder.put(baggageKey, baggageValue);
195+
}
196+
}
190197

191198
int pos = 0;
192199
while (pos < traceHeader.length()) {
@@ -215,12 +222,13 @@ private static <C> Context getContextFromHeader(
215222
spanId = parseSpanId(value);
216223
} else if (trimmedPart.startsWith(SAMPLED_FLAG_KEY)) {
217224
isSampled = parseTraceFlag(value);
218-
} else if (baggageReadBytes + trimmedPart.length() <= 256) {
219-
if (baggage == null) {
220-
baggage = Baggage.builder();
225+
} else if (trimmedPart.startsWith(LINEAGE_KEY)) {
226+
lineageHeader = parseLineageHeader(value);
227+
if (isValidLineage(lineageHeader)) {
228+
baggageBuilder.put(LINEAGE_KEY, lineageHeader);
229+
} else {
230+
logger.fine("Invalid Lineage header: " + value);
221231
}
222-
baggage.put(trimmedPart.substring(0, equalsIndex), value);
223-
baggageReadBytes += trimmedPart.length();
224232
}
225233
}
226234
if (isSampled == null) {
@@ -243,12 +251,17 @@ private static <C> Context getContextFromHeader(
243251
spanId,
244252
isSampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
245253
TraceState.getDefault());
254+
246255
if (spanContext.isValid()) {
247256
context = context.with(Span.wrap(spanContext));
248257
}
249-
if (baggage != null) {
250-
context = context.with(baggage.build());
258+
259+
Baggage baggage = baggageBuilder.build();
260+
261+
if (!baggage.isEmpty()) {
262+
context = context.with(baggage);
251263
}
264+
252265
return context;
253266
}
254267

@@ -316,6 +329,31 @@ private static String parseSpanId(String xrayParentId) {
316329
return xrayParentId;
317330
}
318331

332+
private static String parseLineageHeader(String xrayLineageHeader) {
333+
long numOfDelimiters = xrayLineageHeader.chars().filter(ch -> ch == LINEAGE_DELIMITER).count();
334+
335+
if (xrayLineageHeader.length() < LINEAGE_MIN_LENGTH
336+
|| xrayLineageHeader.length() > LINEAGE_MAX_LENGTH
337+
|| numOfDelimiters != NUM_OF_LINEAGE_DELIMITERS) {
338+
return INVALID_LINEAGE;
339+
}
340+
341+
return xrayLineageHeader;
342+
}
343+
344+
private static boolean isValidLineage(String key) {
345+
String[] split = key.split(String.valueOf(LINEAGE_DELIMITER));
346+
String hash = split[1];
347+
int counter1 = parseIntOrReturnNegative(split[0]);
348+
int counter2 = parseIntOrReturnNegative(split[2]);
349+
350+
boolean isHashValid = hash.length() == LINEAGE_HASH_LENGTH && isValidBase16String(hash);
351+
boolean isValidCounter2 = counter2 <= LINEAGE_MAX_COUNTER2 && counter2 >= LINEAGE_MIN_COUNTER;
352+
boolean isValidCounter1 = counter1 <= LINEAGE_MAX_COUNTER1 && counter1 >= LINEAGE_MIN_COUNTER;
353+
354+
return isHashValid && isValidCounter2 && isValidCounter1;
355+
}
356+
319357
@Nullable
320358
private static Boolean parseTraceFlag(String xraySampledFlag) {
321359
if (xraySampledFlag.length() != SAMPLED_FLAG_LENGTH) {
@@ -332,4 +370,12 @@ private static Boolean parseTraceFlag(String xraySampledFlag) {
332370
return null;
333371
}
334372
}
373+
374+
private static int parseIntOrReturnNegative(String num) {
375+
try {
376+
return Integer.parseInt(num);
377+
} catch (NumberFormatException e) {
378+
return -1;
379+
}
380+
}
335381
}

‎aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayCompositePropagatorTest.java

+78
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.contrib.awsxray.propagator;
7+
8+
import static org.assertj.core.api.Assertions.assertThat;
9+
10+
import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
11+
import io.opentelemetry.api.trace.SpanContext;
12+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
13+
import io.opentelemetry.context.Context;
14+
import io.opentelemetry.context.propagation.TextMapPropagator;
15+
import java.util.LinkedHashMap;
16+
import org.junit.jupiter.api.Test;
17+
18+
public class AwsXrayCompositePropagatorTest extends AwsXrayPropagatorTest {
19+
20+
@Override
21+
TextMapPropagator propagator() {
22+
return TextMapPropagator.composite(
23+
W3CBaggagePropagator.getInstance(),
24+
AwsXrayPropagator.getInstance(),
25+
W3CTraceContextPropagator.getInstance());
26+
}
27+
28+
@Test
29+
void extract_traceContextOverridesXray() {
30+
LinkedHashMap<String, String> carrier = new LinkedHashMap<>();
31+
String w3cTraceContextTraceId = "4bf92f3577b34da6a3ce929d0e0e4736";
32+
String w3cTraceContextSpanId = "00f067aa0ba902b7";
33+
String traceParent =
34+
String.format("00-%s-%s-01", w3cTraceContextTraceId, w3cTraceContextSpanId);
35+
String traceState = "rojo=00f067aa0ba902b7";
36+
String xrayTrace = String.format("Root=1-%s;Parent=%s;Sampled=0", TRACE_ID, SPAN_ID);
37+
38+
carrier.put("traceparent", traceParent);
39+
carrier.put("tracestate", traceState);
40+
carrier.put("X-Amzn-Trace-Id", xrayTrace);
41+
42+
SpanContext actualContext = getSpanContext(subject.extract(Context.current(), carrier, GETTER));
43+
44+
assertThat(actualContext.getTraceId()).isEqualTo(w3cTraceContextTraceId);
45+
assertThat(actualContext.getSpanId()).isEqualTo(w3cTraceContextSpanId);
46+
assertThat(actualContext.isSampled()).isEqualTo(true);
47+
}
48+
49+
@Test
50+
void extract_xrayOverridesTraceContext() {
51+
TextMapPropagator propagator =
52+
TextMapPropagator.composite(
53+
W3CBaggagePropagator.getInstance(),
54+
W3CTraceContextPropagator.getInstance(),
55+
AwsXrayPropagator.getInstance());
56+
57+
LinkedHashMap<String, String> carrier = new LinkedHashMap<>();
58+
String w3cTraceContextTraceId = "4bf92f3577b34da6a3ce929d0e0e4736";
59+
String w3cTraceContextSpanId = "00f067aa0ba902b7";
60+
String traceParent =
61+
String.format("00-%s-%s-01", w3cTraceContextTraceId, w3cTraceContextSpanId);
62+
String traceState = "rojo=00f067aa0ba902b7";
63+
String xrayTrace =
64+
String.format(
65+
"Root=1-%s;Parent=%s;Sampled=0", "8a3c60f7-d188f8fa79d48a391a778fa6", SPAN_ID);
66+
67+
carrier.put("traceparent", traceParent);
68+
carrier.put("tracestate", traceState);
69+
carrier.put("X-Amzn-Trace-Id", xrayTrace);
70+
71+
SpanContext actualContext =
72+
getSpanContext(propagator.extract(Context.current(), carrier, GETTER));
73+
74+
assertThat(actualContext.getTraceId()).isEqualTo(TRACE_ID);
75+
assertThat(actualContext.getSpanId()).isEqualTo(SPAN_ID);
76+
assertThat(actualContext.isSampled()).isEqualTo(false);
77+
}
78+
}

0 commit comments

Comments
 (0)