Skip to content

Commit 42fa0f7

Browse files
liustveliustve
committed
added support for Lineage in Xray Trace Header
1 parent 782b540 commit 42fa0f7

File tree

3 files changed

+239
-96
lines changed

3 files changed

+239
-96
lines changed

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

Lines changed: 87 additions & 39 deletions
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,15 @@ 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_LOOP_COUNTER = 32767;
78+
private static final int LINEAGE_MAX_REQUEST_COUNTER = 255;
79+
private static final int LINEAGE_MIN_COUNTER = 0;
80+
7181
private static final List<String> FIELDS = Collections.singletonList(TRACE_HEADER_KEY);
7282

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

129139
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());
140+
String lineageV2Header = baggage.getEntryValue(LINEAGE_KEY);
141+
142+
if (lineageV2Header != null) {
143+
traceHeader
144+
.append(TRACE_HEADER_DELIMITER)
145+
.append(LINEAGE_KEY)
146+
.append(KV_DELIMITER)
147+
.append(lineageV2Header);
148+
}
149+
150+
// add 256 character truncation
151+
String truncatedTraceHeader = traceHeader.substring(0, Math.min(traceHeader.length(), 256));
152+
setter.set(carrier, TRACE_HEADER_KEY, truncatedTraceHeader);
158153
}
159154

160155
@Override
@@ -183,10 +178,20 @@ private static <C> Context getContextFromHeader(
183178

184179
String traceId = TraceId.getInvalid();
185180
String spanId = SpanId.getInvalid();
181+
String lineageV2Header;
186182
Boolean isSampled = false;
187183

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

191196
int pos = 0;
192197
while (pos < traceHeader.length()) {
@@ -215,12 +220,11 @@ private static <C> Context getContextFromHeader(
215220
spanId = parseSpanId(value);
216221
} else if (trimmedPart.startsWith(SAMPLED_FLAG_KEY)) {
217222
isSampled = parseTraceFlag(value);
218-
} else if (baggageReadBytes + trimmedPart.length() <= 256) {
219-
if (baggage == null) {
220-
baggage = Baggage.builder();
223+
} else if (trimmedPart.startsWith(LINEAGE_KEY)) {
224+
lineageV2Header = parseLineageV2Header(value);
225+
if (isValidLineage(lineageV2Header)) {
226+
baggageBuilder.put(LINEAGE_KEY, lineageV2Header);
221227
}
222-
baggage.put(trimmedPart.substring(0, equalsIndex), value);
223-
baggageReadBytes += trimmedPart.length();
224228
}
225229
}
226230
if (isSampled == null) {
@@ -243,12 +247,17 @@ private static <C> Context getContextFromHeader(
243247
spanId,
244248
isSampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
245249
TraceState.getDefault());
250+
246251
if (spanContext.isValid()) {
247252
context = context.with(Span.wrap(spanContext));
248253
}
249-
if (baggage != null) {
250-
context = context.with(baggage.build());
254+
255+
Baggage baggage = baggageBuilder.build();
256+
257+
if (!baggage.isEmpty()) {
258+
context = context.with(baggage);
251259
}
260+
252261
return context;
253262
}
254263

@@ -316,6 +325,37 @@ private static String parseSpanId(String xrayParentId) {
316325
return xrayParentId;
317326
}
318327

328+
private static String parseLineageV2Header(String xrayLineageHeader) {
329+
long numOfDelimiters = xrayLineageHeader.chars().filter(ch -> ch == LINEAGE_DELIMITER).count();
330+
331+
if (xrayLineageHeader.length() < LINEAGE_MIN_LENGTH
332+
|| xrayLineageHeader.length() > LINEAGE_MAX_LENGTH
333+
|| numOfDelimiters != 2) {
334+
return AwsXrayPropagator.getInvalidLineageV2Header();
335+
}
336+
337+
return xrayLineageHeader;
338+
}
339+
340+
private static boolean isValidLineage(String key) {
341+
String[] split = key.split(":");
342+
String hash = split[1];
343+
int loopCounter = parseIntOrReturnNegative(split[0]);
344+
int requestCounter = parseIntOrReturnNegative(split[2]);
345+
346+
boolean isHashValid = hash.length() == LINEAGE_HASH_LENGTH && isValidBase16String(hash);
347+
boolean isValidRequestCounter =
348+
requestCounter <= LINEAGE_MAX_REQUEST_COUNTER && requestCounter >= LINEAGE_MIN_COUNTER;
349+
boolean isValidLoopCounter =
350+
loopCounter <= LINEAGE_MAX_LOOP_COUNTER && loopCounter >= LINEAGE_MIN_COUNTER;
351+
352+
return isHashValid && isValidRequestCounter && isValidLoopCounter;
353+
}
354+
355+
private static String getInvalidLineageV2Header() {
356+
return "-1:11111111:0";
357+
}
358+
319359
@Nullable
320360
private static Boolean parseTraceFlag(String xraySampledFlag) {
321361
if (xraySampledFlag.length() != SAMPLED_FLAG_LENGTH) {
@@ -332,4 +372,12 @@ private static Boolean parseTraceFlag(String xraySampledFlag) {
332372
return null;
333373
}
334374
}
375+
376+
private static int parseIntOrReturnNegative(String num) {
377+
try {
378+
return Integer.parseInt(num);
379+
} catch (NumberFormatException e) {
380+
return -1;
381+
}
382+
}
335383
}

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

Lines changed: 78 additions & 0 deletions
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)