Skip to content

Commit 3c15d40

Browse files
majanjua-amznmajanjua-amzn
authored
Support Lambda PassThrough trace header propagation (#409)
* Support Lambda PassThrough trace header propagation * Account for Sampled=0 trace header in Active mode * Handle Sampled=0 Active Tracing case * Ensure Parent is always copied from entity into trace header
1 parent 101a1e6 commit 3c15d40

File tree

9 files changed

+104
-30
lines changed

9 files changed

+104
-30
lines changed

‎aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.amazonaws.xray.AWSXRay;
1919
import com.amazonaws.xray.AWSXRayRecorder;
2020
import com.amazonaws.xray.entities.Namespace;
21+
import com.amazonaws.xray.entities.NoOpSegment;
2122
import com.amazonaws.xray.entities.Segment;
2223
import com.amazonaws.xray.entities.Subsegment;
2324
import com.amazonaws.xray.entities.TraceHeader;
@@ -104,7 +105,16 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
104105
Segment parentSegment = subsegment.getParentSegment();
105106

106107
if (subsegment.shouldPropagate()) {
107-
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(subsegment).toString());
108+
// If no-op, only propagate root trace ID to not taint sampling decision
109+
TraceHeader t = TraceHeader.fromEntity(subsegment);
110+
if (parentSegment instanceof NoOpSegment) {
111+
request.addHeader(
112+
TraceHeader.HEADER_KEY,
113+
"Root=" + t.getRootTraceId().toString());
114+
} else {
115+
// This will propagate Parent and Sampled
116+
request.addHeader(TraceHeader.HEADER_KEY, t.toString());
117+
}
108118
}
109119

110120
Map<String, Object> requestInformation = new HashMap<>();

‎aws-xray-recorder-sdk-apache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
2020
import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
2121
import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
22+
import static com.github.tomakehurst.wiremock.client.WireMock.matching;
2223
import static com.github.tomakehurst.wiremock.client.WireMock.ok;
2324
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
2425
import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
@@ -99,7 +100,7 @@ public void unsampledPropagation() throws Exception {
99100

100101
verify(getRequestedFor(urlPathEqualTo("/"))
101102
.withHeader(TraceHeader.HEADER_KEY,
102-
equalTo("Root=1-00000000-000000000000000000000000;Sampled=0")));
103+
equalTo("Root=1-00000000-000000000000000000000000")));
103104
}
104105

105106
@Test
@@ -114,7 +115,7 @@ public void unsampledButForcedPropagation() throws Exception {
114115

115116
verify(getRequestedFor(urlPathEqualTo("/"))
116117
.withHeader(TraceHeader.HEADER_KEY,
117-
equalTo("Root=1-67891233-abcdef012345678912345678;Sampled=0")));
118+
matching("Root=1-67891233-abcdef012345678912345678;Parent=[a-z0-9]{16};Sampled=0")));
118119
}
119120

120121
@Test

‎aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.amazonaws.xray.entities.EntityDataKeys;
2222
import com.amazonaws.xray.entities.EntityHeaderKeys;
2323
import com.amazonaws.xray.entities.Namespace;
24+
import com.amazonaws.xray.entities.NoOpSegment;
2425
import com.amazonaws.xray.entities.Subsegment;
2526
import com.amazonaws.xray.entities.TraceHeader;
2627
import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -296,9 +297,18 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
296297
return httpRequest;
297298
}
298299

300+
// If no-op, only propagate root trace ID to not taint sampling decision
301+
TraceHeader t = TraceHeader.fromEntity(subsegment);
302+
if (subsegment.getParentSegment() instanceof NoOpSegment) {
303+
return httpRequest.toBuilder().putHeader(
304+
TraceHeader.HEADER_KEY,
305+
"Root=" + t.getRootTraceId().toString()).build();
306+
}
307+
308+
// This will propagate Parent and Sampled
299309
return httpRequest.toBuilder().putHeader(
300310
TraceHeader.HEADER_KEY,
301-
TraceHeader.fromEntity(subsegment).toString()).build();
311+
t.toString()).build();
302312
}
303313

304314
@Override

‎aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.amazonaws.xray.entities.EntityDataKeys;
3232
import com.amazonaws.xray.entities.EntityHeaderKeys;
3333
import com.amazonaws.xray.entities.Namespace;
34+
import com.amazonaws.xray.entities.NoOpSegment;
3435
import com.amazonaws.xray.entities.Subsegment;
3536
import com.amazonaws.xray.entities.TraceHeader;
3637
import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -192,7 +193,16 @@ public void beforeRequest(Request<?> request) {
192193
currentSubsegment.setNamespace(Namespace.AWS.toString());
193194

194195
if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
195-
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(currentSubsegment).toString());
196+
// If no-op, only propagate root trace ID to not taint sampling decision
197+
TraceHeader t = TraceHeader.fromEntity(currentSubsegment);
198+
if (currentSubsegment.getParentSegment() instanceof NoOpSegment) {
199+
request.addHeader(
200+
TraceHeader.HEADER_KEY,
201+
"Root=" + t.getRootTraceId().toString());
202+
} else {
203+
// This will propagate Parent and Sampled
204+
request.addHeader(TraceHeader.HEADER_KEY, t.toString());
205+
}
196206
}
197207
}
198208

‎aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,10 @@ public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boole
142142
TraceHeader.SampleDecision.SAMPLED :
143143
TraceHeader.SampleDecision.NOT_SAMPLED);
144144
assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
145-
assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
145+
assertThat(traceHeader.getParentId()).isEqualTo(
146+
subsegment.isSampled() ?
147+
serviceEntityId :
148+
subsegment.getParentSegment().getId());
146149

147150
tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
148151

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java‎

Lines changed: 26 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,11 @@
1818
import com.amazonaws.xray.AWSXRayRecorder;
1919
import com.amazonaws.xray.entities.Entity;
2020
import com.amazonaws.xray.entities.FacadeSegment;
21+
import com.amazonaws.xray.entities.NoOpSegment;
2122
import com.amazonaws.xray.entities.Segment;
2223
import com.amazonaws.xray.entities.Subsegment;
2324
import com.amazonaws.xray.entities.SubsegmentImpl;
2425
import com.amazonaws.xray.entities.TraceHeader;
25-
import com.amazonaws.xray.entities.TraceHeader.SampleDecision;
2626
import com.amazonaws.xray.entities.TraceID;
2727
import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
2828
import com.amazonaws.xray.listeners.SegmentListener;
@@ -39,36 +39,42 @@ public class LambdaSegmentContext implements SegmentContext {
3939
// See: https://github.com/aws/aws-xray-sdk-java/issues/251
4040
private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
4141

42-
private static TraceHeader getTraceHeaderFromEnvironment() {
42+
public static TraceHeader getTraceHeaderFromEnvironment() {
4343
String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY);
4444
return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0
4545
? lambdaTraceHeaderKey
4646
: System.getProperty(LAMBDA_TRACE_HEADER_PROP));
4747
}
4848

49-
private static boolean isInitializing(TraceHeader traceHeader) {
50-
return traceHeader.getRootTraceId() == null || traceHeader.getSampled() == null || traceHeader.getParentId() == null;
51-
}
52-
53-
private static FacadeSegment newFacadeSegment(AWSXRayRecorder recorder, String name) {
54-
TraceHeader traceHeader = getTraceHeaderFromEnvironment();
55-
if (isInitializing(traceHeader)) {
56-
logger.warn(LAMBDA_TRACE_HEADER_KEY + " is missing a trace ID, parent ID, or sampling decision. Subsegment "
57-
+ name + " discarded.");
58-
return new FacadeSegment(recorder, TraceID.create(recorder), "", SampleDecision.NOT_SAMPLED);
59-
}
60-
return new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled());
61-
}
62-
49+
// SuppressWarnings is needed for passing Root TraceId to noOp segment
50+
@SuppressWarnings("nullness")
6351
@Override
6452
public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
6553
if (logger.isDebugEnabled()) {
6654
logger.debug("Beginning subsegment named: " + name);
6755
}
6856

57+
TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment();
58+
logger.warn("TRACE HEADER IN CODE: " + traceHeader.toString());
6959
Entity entity = getTraceEntity();
70-
if (entity == null) { // First subsgment of a subsegment branch.
71-
Segment parentSegment = newFacadeSegment(recorder, name);
60+
if (entity == null) { // First subsegment of a subsegment branch
61+
Segment parentSegment;
62+
// Trace header either takes the structure `Root=...;<extra-data>` or
63+
// `Root=...;Parent=...;Sampled=...;<extra-data>`
64+
if (traceHeader.getRootTraceId() != null && traceHeader.getParentId() != null && traceHeader.getSampled() != null) {
65+
logger.warn("CREATING FACADE SEGMENT");
66+
parentSegment = new FacadeSegment(
67+
recorder,
68+
traceHeader.getRootTraceId(),
69+
traceHeader.getParentId(),
70+
traceHeader.getSampled());
71+
} else {
72+
if (logger.isDebugEnabled()) {
73+
logger.debug("Creating No-Op parent segment");
74+
}
75+
TraceID t = traceHeader.getRootTraceId() != null ? traceHeader.getRootTraceId() : TraceID.create(recorder);
76+
parentSegment = Segment.noOp(t, recorder);
77+
}
7278

7379
boolean isRecording = parentSegment.isRecording();
7480

@@ -145,6 +151,8 @@ public void endSubsegment(AWSXRayRecorder recorder) {
145151
current.getCreator().getEmitter().sendSubsegment((Subsegment) current);
146152
}
147153
clearTraceEntity();
154+
} else if (parentEntity instanceof NoOpSegment) {
155+
clearTraceEntity();
148156
} else {
149157
setTraceEntity(current.getParent());
150158
}

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import java.util.concurrent.locks.ReentrantLock;
2323
import org.checkerframework.checker.nullness.qual.Nullable;
2424

25-
class NoOpSegment implements Segment {
25+
public class NoOpSegment implements Segment {
2626

2727
private final TraceID traceId;
2828
private final AWSXRayRecorder creator;

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,16 @@ public TraceHeader(@Nullable TraceID rootTraceId, @Nullable String parentId, Sam
9292
}
9393

9494
public static TraceHeader fromEntity(Entity entity) {
95+
String parentId = null;
96+
if (entity instanceof Subsegment) {
97+
Segment segment = entity.getParentSegment();
98+
if (segment != null) {
99+
parentId = segment.getId();
100+
}
101+
}
95102
return new TraceHeader(
96103
entity.getTraceId(),
97-
entity.isSampled() ? entity.getId() : null,
104+
(entity.getId() == null || entity.getId() == "") ? parentId : entity.getId(),
98105
entity.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
99106
}
100107

‎aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.amazonaws.xray.AWSXRayRecorderBuilder;
2323
import com.amazonaws.xray.emitters.Emitter;
2424
import com.amazonaws.xray.entities.FacadeSegment;
25+
import com.amazonaws.xray.entities.NoOpSegment;
2526
import com.amazonaws.xray.entities.Subsegment;
2627
import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
2728
import com.amazonaws.xray.strategy.LogErrorContextMissingStrategy;
@@ -49,6 +50,10 @@ class LambdaSegmentContextTest {
4950

5051
private static final String MALFORMED_TRACE_HEADER =
5152
";;Root=1-57ff426a-80c11c39b0c928905eb0828d;;Parent=1234abcd1234abcd;;;Sampled=1;;;";
53+
private static final String MALFORMED_TRACE_HEADER_2 = ";;root-missing;;Parent=1234abcd1234abcd;;;Sampled=1;;;";
54+
55+
private static final String ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER =
56+
"Root=1-5759e988-bd862e3fe1be46a994272711;Lineage=10:1234abcd:3";
5257

5358
@BeforeEach
5459
public void setupAWSXRay() {
@@ -63,14 +68,14 @@ public void setupAWSXRay() {
6368
}
6469

6570
@Test
66-
void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
67-
testContextResultsInFacadeSegmentParent();
71+
void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
72+
testContextResultsInNoOpSegmentParent();
6873
}
6974

7075
@Test
7176
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = "a")
72-
void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
73-
testContextResultsInFacadeSegmentParent();
77+
void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
78+
testContextResultsInNoOpSegmentParent();
7479
}
7580

7681
@Test
@@ -85,6 +90,18 @@ void testBeginSubsegmentWithCompleteButMalformedTraceHeaderEnvironmentVariableRe
8590
testContextResultsInFacadeSegmentParent();
8691
}
8792

93+
@Test
94+
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = MALFORMED_TRACE_HEADER_2)
95+
void testBeginSubsegmentWithIncompleteAndMalformedTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
96+
testContextResultsInNoOpSegmentParent();
97+
}
98+
99+
@Test
100+
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER)
101+
void testBeginSubsegmentWithRootLambdaPassthroughTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
102+
testContextResultsInNoOpSegmentParent();
103+
}
104+
88105
@Test
89106
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2)
90107
void testNotSampledSetsParentToSubsegment() {
@@ -149,4 +166,12 @@ private static void testContextResultsInFacadeSegmentParent() {
149166
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
150167
assertThat(AWSXRay.getTraceEntity()).isNull();
151168
}
169+
170+
private static void testContextResultsInNoOpSegmentParent() {
171+
LambdaSegmentContext mockContext = new LambdaSegmentContext();
172+
assertThat(mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test").getParent())
173+
.isInstanceOf(NoOpSegment.class);
174+
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
175+
assertThat(AWSXRay.getTraceEntity()).isNull();
176+
}
152177
}

0 commit comments

Comments
 (0)