From 9950b38e4a34ba0c7fb4559e88dca31350267dec Mon Sep 17 00:00:00 2001 From: Mahad Janjua Date: Mon, 6 May 2024 13:07:06 -0700 Subject: [PATCH 1/3] [Lambda] Send NoOp segment when trace header is incomplete --- .../xray/contexts/LambdaSegmentContext.java | 21 +++++++--------- .../amazonaws/xray/entities/NoOpSegment.java | 2 +- .../contexts/LambdaSegmentContextTest.java | 25 ++++++++++++++++--- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index 25c99506..e1f2aac4 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java @@ -18,6 +18,7 @@ import com.amazonaws.xray.AWSXRayRecorder; import com.amazonaws.xray.entities.Entity; import com.amazonaws.xray.entities.FacadeSegment; +import com.amazonaws.xray.entities.NoOpSegment; import com.amazonaws.xray.entities.Segment; import com.amazonaws.xray.entities.Subsegment; import com.amazonaws.xray.entities.SubsegmentImpl; @@ -50,25 +51,18 @@ private static boolean isInitializing(TraceHeader traceHeader) { return traceHeader.getRootTraceId() == null || traceHeader.getSampled() == null || traceHeader.getParentId() == null; } - private static FacadeSegment newFacadeSegment(AWSXRayRecorder recorder, String name) { - TraceHeader traceHeader = getTraceHeaderFromEnvironment(); - if (isInitializing(traceHeader)) { - logger.warn(LAMBDA_TRACE_HEADER_KEY + " is missing a trace ID, parent ID, or sampling decision. Subsegment " - + name + " discarded."); - return new FacadeSegment(recorder, TraceID.create(recorder), "", SampleDecision.NOT_SAMPLED); - } - return new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled()); - } - @Override public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) { if (logger.isDebugEnabled()) { logger.debug("Beginning subsegment named: " + name); } + TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment(); Entity entity = getTraceEntity(); - if (entity == null) { // First subsgment of a subsegment branch. - Segment parentSegment = newFacadeSegment(recorder, name); + if (entity == null) { // First subsegment of a subsegment branch + Segment parentSegment = isInitializing(traceHeader) + ? Segment.noOp(TraceID.create(recorder), recorder) + : new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled()); boolean isRecording = parentSegment.isRecording(); @@ -145,6 +139,9 @@ public void endSubsegment(AWSXRayRecorder recorder) { current.getCreator().getEmitter().sendSubsegment((Subsegment) current); } clearTraceEntity(); + } + else if (parentEntity instanceof NoOpSegment) { + clearTraceEntity(); } else { setTraceEntity(current.getParent()); } diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java index e406e9ac..d470d2df 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSegment.java @@ -22,7 +22,7 @@ import java.util.concurrent.locks.ReentrantLock; import org.checkerframework.checker.nullness.qual.Nullable; -class NoOpSegment implements Segment { +public class NoOpSegment implements Segment { private final TraceID traceId; private final AWSXRayRecorder creator; diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java index 5f62747e..4d9628b7 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java +++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java @@ -22,6 +22,7 @@ import com.amazonaws.xray.AWSXRayRecorderBuilder; import com.amazonaws.xray.emitters.Emitter; import com.amazonaws.xray.entities.FacadeSegment; +import com.amazonaws.xray.entities.NoOpSegment; import com.amazonaws.xray.entities.Subsegment; import com.amazonaws.xray.exceptions.SubsegmentNotFoundException; import com.amazonaws.xray.strategy.LogErrorContextMissingStrategy; @@ -49,6 +50,8 @@ class LambdaSegmentContextTest { private static final String MALFORMED_TRACE_HEADER = ";;Root=1-57ff426a-80c11c39b0c928905eb0828d;;Parent=1234abcd1234abcd;;;Sampled=1;;;"; + private static final String MALFORMED_TRACE_HEADER_2 = + ";;root-missing;;Parent=1234abcd1234abcd;;;Sampled=1;;;"; @BeforeEach public void setupAWSXRay() { @@ -63,14 +66,14 @@ public void setupAWSXRay() { } @Test - void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() { - testContextResultsInFacadeSegmentParent(); + void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() { + testContextResultsInNoOpSegmentParent(); } @Test @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = "a") - void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() { - testContextResultsInFacadeSegmentParent(); + void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() { + testContextResultsInNoOpSegmentParent(); } @Test @@ -85,6 +88,12 @@ void testBeginSubsegmentWithCompleteButMalformedTraceHeaderEnvironmentVariableRe testContextResultsInFacadeSegmentParent(); } + @Test + @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = MALFORMED_TRACE_HEADER_2) + void testBeginSubsegmentWithIncompleteAndMalformedTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() { + testContextResultsInNoOpSegmentParent(); + } + @Test @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2) void testNotSampledSetsParentToSubsegment() { @@ -149,4 +158,12 @@ private static void testContextResultsInFacadeSegmentParent() { mockContext.endSubsegment(AWSXRay.getGlobalRecorder()); assertThat(AWSXRay.getTraceEntity()).isNull(); } + + private static void testContextResultsInNoOpSegmentParent() { + LambdaSegmentContext mockContext = new LambdaSegmentContext(); + assertThat(mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test").getParent()) + .isInstanceOf(NoOpSegment.class); + mockContext.endSubsegment(AWSXRay.getGlobalRecorder()); + assertThat(AWSXRay.getTraceEntity()).isNull(); + } } From ed061becf4208878e0f9244541ae5766826f2217 Mon Sep 17 00:00:00 2001 From: Mahad Janjua Date: Mon, 6 May 2024 14:34:42 -0700 Subject: [PATCH 2/3] [Lambda] Add debug log and root lambda passthrough test case --- .../xray/contexts/LambdaSegmentContext.java | 12 +++++++++--- .../xray/contexts/LambdaSegmentContextTest.java | 12 ++++++++++-- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index e1f2aac4..85d690fe 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java @@ -60,9 +60,15 @@ public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) { TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment(); Entity entity = getTraceEntity(); if (entity == null) { // First subsegment of a subsegment branch - Segment parentSegment = isInitializing(traceHeader) - ? Segment.noOp(TraceID.create(recorder), recorder) - : new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled()); + Segment parentSegment; + if (isInitializing(traceHeader)) { + if (logger.isDebugEnabled()) { + logger.debug("Creating No-Op parent segment"); + } + parentSegment = Segment.noOp(TraceID.create(recorder), recorder); + } else { + parentSegment = new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled()); + } boolean isRecording = parentSegment.isRecording(); diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java index 4d9628b7..2f7a728b 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java +++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java @@ -50,8 +50,10 @@ class LambdaSegmentContextTest { private static final String MALFORMED_TRACE_HEADER = ";;Root=1-57ff426a-80c11c39b0c928905eb0828d;;Parent=1234abcd1234abcd;;;Sampled=1;;;"; - private static final String MALFORMED_TRACE_HEADER_2 = - ";;root-missing;;Parent=1234abcd1234abcd;;;Sampled=1;;;"; + private static final String MALFORMED_TRACE_HEADER_2 = ";;root-missing;;Parent=1234abcd1234abcd;;;Sampled=1;;;"; + + private static final String ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER = + "Root=1-5759e988-bd862e3fe1be46a994272711;Lineage=10:1234abcd:3"; @BeforeEach public void setupAWSXRay() { @@ -94,6 +96,12 @@ void testBeginSubsegmentWithIncompleteAndMalformedTraceHeaderEnvironmentVariable testContextResultsInNoOpSegmentParent(); } + @Test + @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER) + void testBeginSubsegmentWithRootLambdaPassthroughTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() { + testContextResultsInNoOpSegmentParent(); + } + @Test @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2) void testNotSampledSetsParentToSubsegment() { From 8ba9ff837311414ef095fc172871c8074ec60b7e Mon Sep 17 00:00:00 2001 From: Mahad Janjua Date: Mon, 6 May 2024 15:05:31 -0700 Subject: [PATCH 3/3] Fix checkstyle --- .../xray/contexts/LambdaSegmentContext.java | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index 85d690fe..bb92a166 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java @@ -23,7 +23,6 @@ import com.amazonaws.xray.entities.Subsegment; import com.amazonaws.xray.entities.SubsegmentImpl; import com.amazonaws.xray.entities.TraceHeader; -import com.amazonaws.xray.entities.TraceHeader.SampleDecision; import com.amazonaws.xray.entities.TraceID; import com.amazonaws.xray.exceptions.SubsegmentNotFoundException; import com.amazonaws.xray.listeners.SegmentListener; @@ -36,14 +35,14 @@ public class LambdaSegmentContext implements SegmentContext { private static final Log logger = LogFactory.getLog(LambdaSegmentContext.class); private static final String LAMBDA_TRACE_HEADER_KEY = "_X_AMZN_TRACE_ID"; - + // See: https://github.com/aws/aws-xray-sdk-java/issues/251 private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader"; private static TraceHeader getTraceHeaderFromEnvironment() { String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY); - return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0 - ? lambdaTraceHeaderKey + return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0 + ? lambdaTraceHeaderKey : System.getProperty(LAMBDA_TRACE_HEADER_PROP)); } @@ -67,7 +66,12 @@ public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) { } parentSegment = Segment.noOp(TraceID.create(recorder), recorder); } else { - parentSegment = new FacadeSegment(recorder, traceHeader.getRootTraceId(), traceHeader.getParentId(), traceHeader.getSampled()); + parentSegment = new FacadeSegment( + recorder, + traceHeader.getRootTraceId(), + traceHeader.getParentId(), + traceHeader.getSampled() + ); } boolean isRecording = parentSegment.isRecording(); @@ -145,8 +149,7 @@ public void endSubsegment(AWSXRayRecorder recorder) { current.getCreator().getEmitter().sendSubsegment((Subsegment) current); } clearTraceEntity(); - } - else if (parentEntity instanceof NoOpSegment) { + } else if (parentEntity instanceof NoOpSegment) { clearTraceEntity(); } else { setTraceEntity(current.getParent());