Support Lambda PassThrough trace header propagation

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

@@ -18,6 +18,7 @@
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.AWSXRayRecorder;
 import com.amazonaws.xray.entities.Namespace;
+import com.amazonaws.xray.entities.NoOpSegment;
 import com.amazonaws.xray.entities.Segment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
@@ -104,7 +105,16 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
         Segment parentSegment = subsegment.getParentSegment();
 
         if (subsegment.shouldPropagate()) {
-            request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(subsegment).toString());
+            // If no-op, only propagate root trace ID to not taint sampling decision
+            TraceHeader t = TraceHeader.fromEntity(subsegment);
+            if (parentSegment instanceof NoOpSegment) {
+                request.addHeader(
+                        TraceHeader.HEADER_KEY,
+                        "Root=" + t.getRootTraceId().toString());
+            } else {
+                // This will propagate Parent and Sampled
+                request.addHeader(TraceHeader.HEADER_KEY, t.toString());
+            }
         }
 
         Map<String, Object> requestInformation = new HashMap<>();

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

@@ -19,6 +19,7 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
 import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
 import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.matching;
 import static com.github.tomakehurst.wiremock.client.WireMock.ok;
 import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
@@ -99,7 +100,7 @@ public void unsampledPropagation() throws Exception {
 
         verify(getRequestedFor(urlPathEqualTo("/"))
                    .withHeader(TraceHeader.HEADER_KEY,
-                               equalTo("Root=1-00000000-000000000000000000000000;Sampled=0")));
+                               equalTo("Root=1-00000000-000000000000000000000000")));
     }
 
     @Test
@@ -114,7 +115,7 @@ public void unsampledButForcedPropagation() throws Exception {
 
         verify(getRequestedFor(urlPathEqualTo("/"))
                    .withHeader(TraceHeader.HEADER_KEY,
-                               equalTo("Root=1-67891233-abcdef012345678912345678;Sampled=0")));
+                               matching("Root=1-67891233-abcdef012345678912345678;Parent=[a-z0-9]{16};Sampled=0")));
     }
 
     @Test

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

@@ -21,6 +21,7 @@
 import com.amazonaws.xray.entities.EntityDataKeys;
 import com.amazonaws.xray.entities.EntityHeaderKeys;
 import com.amazonaws.xray.entities.Namespace;
+import com.amazonaws.xray.entities.NoOpSegment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
 import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -296,9 +297,18 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
             return httpRequest;
         }
 
+        // If no-op, only propagate root trace ID to not taint sampling decision
+        TraceHeader t = TraceHeader.fromEntity(subsegment);
+        if (subsegment.getParentSegment() instanceof NoOpSegment) {
+            return httpRequest.toBuilder().putHeader(
+                    TraceHeader.HEADER_KEY,
+                    "Root=" + t.getRootTraceId().toString()).build();
+        }
+
+        // This will propagate Parent and Sampled
         return httpRequest.toBuilder().putHeader(
                 TraceHeader.HEADER_KEY,
-                TraceHeader.fromEntity(subsegment).toString()).build();
+                t.toString()).build();
     }
 
     @Override

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

@@ -31,6 +31,7 @@
 import com.amazonaws.xray.entities.EntityDataKeys;
 import com.amazonaws.xray.entities.EntityHeaderKeys;
 import com.amazonaws.xray.entities.Namespace;
+import com.amazonaws.xray.entities.NoOpSegment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
 import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -192,7 +193,16 @@ public void beforeRequest(Request<?> request) {
         currentSubsegment.setNamespace(Namespace.AWS.toString());
 
         if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
-            request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(currentSubsegment).toString());
+            // If no-op, only propagate root trace ID to not taint sampling decision
+            TraceHeader t = TraceHeader.fromEntity(currentSubsegment);
+            if (currentSubsegment.getParentSegment() instanceof NoOpSegment) {
+                request.addHeader(
+                        TraceHeader.HEADER_KEY,
+                        "Root=" + t.getRootTraceId().toString());
+            } else {
+                // This will propagate Parent and Sampled
+                request.addHeader(TraceHeader.HEADER_KEY, t.toString());
+            }
         }
     }
 

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

@@ -142,7 +142,10 @@ public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boole
                         TraceHeader.SampleDecision.SAMPLED :
                         TraceHeader.SampleDecision.NOT_SAMPLED);
         assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
-        assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
+        assertThat(traceHeader.getParentId()).isEqualTo(
+                subsegment.isSampled() ?
+                        serviceEntityId :
+                        subsegment.getParentSegment().getId());
 
         tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
 

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

@@ -18,11 +18,11 @@
 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;
 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;
@@ -39,36 +39,42 @@ public class LambdaSegmentContext implements SegmentContext {
     // 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() {
+    public static TraceHeader getTraceHeaderFromEnvironment() {
         String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY);
         return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0 
             ? lambdaTraceHeaderKey 
             : System.getProperty(LAMBDA_TRACE_HEADER_PROP));
     }
 
-    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());
-    }
-
+    // SuppressWarnings is needed for passing Root TraceId to noOp segment
+    @SuppressWarnings("nullness")
     @Override
     public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
         if (logger.isDebugEnabled()) {
             logger.debug("Beginning subsegment named: " + name);
         }
 
+        TraceHeader traceHeader = LambdaSegmentContext.getTraceHeaderFromEnvironment();
+        logger.warn("TRACE HEADER IN CODE: " + traceHeader.toString());
         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;
+            // Trace header either takes the structure `Root=...;<extra-data>` or
+            // `Root=...;Parent=...;Sampled=...;<extra-data>`
+            if (traceHeader.getRootTraceId() != null && traceHeader.getParentId() != null && traceHeader.getSampled() != null) {
+                logger.warn("CREATING FACADE SEGMENT");
+                parentSegment = new FacadeSegment(
+                    recorder,
+                    traceHeader.getRootTraceId(),
+                    traceHeader.getParentId(),
+                    traceHeader.getSampled());
+            } else {
+                if (logger.isDebugEnabled()) {
+                    logger.debug("Creating No-Op parent segment");
+                }
+                TraceID t = traceHeader.getRootTraceId() != null ? traceHeader.getRootTraceId() : TraceID.create(recorder);
+                parentSegment = Segment.noOp(t, recorder);
+            }
 
             boolean isRecording = parentSegment.isRecording();
 
@@ -145,6 +151,8 @@ public void endSubsegment(AWSXRayRecorder recorder) {
                     current.getCreator().getEmitter().sendSubsegment((Subsegment) current);
                 }
                 clearTraceEntity();
+            } else if (parentEntity instanceof NoOpSegment) {
+                clearTraceEntity();
             } else {
                 setTraceEntity(current.getParent());
             }

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;

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

@@ -92,9 +92,16 @@ public TraceHeader(@Nullable TraceID rootTraceId, @Nullable String parentId, Sam
     }
 
     public static TraceHeader fromEntity(Entity entity) {
+        String parentId = null;
+        if (entity instanceof Subsegment) {
+            Segment segment = entity.getParentSegment();
+            if (segment != null) {
+                parentId = segment.getId();
+            }
+        }
         return new TraceHeader(
                 entity.getTraceId(),
-                entity.isSampled() ? entity.getId() : null,
+                (entity.getId() == null || entity.getId() == "") ? parentId : entity.getId(),
                 entity.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
     }
 

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,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 ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER =
+        "Root=1-5759e988-bd862e3fe1be46a994272711;Lineage=10:1234abcd:3";
 
     @BeforeEach
     public void setupAWSXRay() {
@@ -63,14 +68,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 +90,18 @@ 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 = ROOT_LAMBDA_PASSTHROUGH_TRACE_HEADER)
+    void testBeginSubsegmentWithRootLambdaPassthroughTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
+        testContextResultsInNoOpSegmentParent();
+    }
+
     @Test
     @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2)
     void testNotSampledSetsParentToSubsegment() {
@@ -149,4 +166,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();
+    }
 }