Skip to content

Commit c275035

Browse files
majanjua-amznmajanjua-amzn
authored
[Lambda] Send NoOp segment when trace header is incomplete (#403)
* [Lambda] Send NoOp segment when trace header is incomplete * [Lambda] Add debug log and root lambda passthrough test case * Fix checkstyle
1 parent 0c2cfa7 commit c275035

File tree

3 files changed

+52
-21
lines changed

3 files changed

+52
-21
lines changed

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

Lines changed: 22 additions & 16 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;
@@ -35,40 +35,44 @@ public class LambdaSegmentContext implements SegmentContext {
3535
private static final Log logger = LogFactory.getLog(LambdaSegmentContext.class);
3636

3737
private static final String LAMBDA_TRACE_HEADER_KEY = "_X_AMZN_TRACE_ID";
38-
38+
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

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

4949
private static boolean isInitializing(TraceHeader traceHeader) {
5050
return traceHeader.getRootTraceId() == null || traceHeader.getSampled() == null || traceHeader.getParentId() == null;
5151
}
5252

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-
6353
@Override
6454
public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
6555
if (logger.isDebugEnabled()) {
6656
logger.debug("Beginning subsegment named: " + name);
6757
}
6858

59+
TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment();
6960
Entity entity = getTraceEntity();
70-
if (entity == null) { // First subsgment of a subsegment branch.
71-
Segment parentSegment = newFacadeSegment(recorder, name);
61+
if (entity == null) { // First subsegment of a subsegment branch
62+
Segment parentSegment;
63+
if (isInitializing(traceHeader)) {
64+
if (logger.isDebugEnabled()) {
65+
logger.debug("Creating No-Op parent segment");
66+
}
67+
parentSegment = Segment.noOp(TraceID.create(recorder), recorder);
68+
} else {
69+
parentSegment = new FacadeSegment(
70+
recorder,
71+
traceHeader.getRootTraceId(),
72+
traceHeader.getParentId(),
73+
traceHeader.getSampled()
74+
);
75+
}
7276

7377
boolean isRecording = parentSegment.isRecording();
7478

@@ -145,6 +149,8 @@ public void endSubsegment(AWSXRayRecorder recorder) {
145149
current.getCreator().getEmitter().sendSubsegment((Subsegment) current);
146150
}
147151
clearTraceEntity();
152+
} else if (parentEntity instanceof NoOpSegment) {
153+
clearTraceEntity();
148154
} else {
149155
setTraceEntity(current.getParent());
150156
}

‎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/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)