[Lambda] Send NoOp segment when trace header is incomplete

aws · May 6, 2024 · 9950b38 · 9950b38
1 parent 0c2cfa7
commit 9950b38
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 17 deletions.
diff --git a/...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/...ray-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
@@ -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/...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/...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();
+    }
 }