[Lambda] Do not propagate trace header if dummy subsegment (#404)

Author: majanjua-amzn

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.NoOpSubSegment;
 import com.amazonaws.xray.entities.Segment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
@@ -103,7 +104,7 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
         subsegment.setNamespace(Namespace.REMOTE.toString());
         Segment parentSegment = subsegment.getParentSegment();
 
-        if (subsegment.shouldPropagate()) {
+        if (!(subsegment instanceof NoOpSubSegment) && subsegment.shouldPropagate()) {
             request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(subsegment).toString());
         }
 

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

@@ -99,7 +99,7 @@ public void unsampledPropagation() throws Exception {
 
         verify(getRequestedFor(urlPathEqualTo("/"))
                    .withHeader(TraceHeader.HEADER_KEY,
-                               equalTo("Root=1-00000000-000000000000000000000000;Sampled=0")));
+                               equalTo("Root=1-67891233-abcdef012345678912345678;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.NoOpSubSegment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
 import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -292,7 +293,7 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
         }
 
         Subsegment subsegment = executionAttributes.getAttribute(entityKey);
-        if (!subsegment.shouldPropagate()) {
+        if (!subsegment.shouldPropagate() || subsegment instanceof NoOpSubSegment) {
             return httpRequest;
         }
 

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.NoOpSubSegment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
 import com.amazonaws.xray.handlers.config.AWSOperationHandler;
@@ -191,7 +192,9 @@ public void beforeRequest(Request<?> request) {
         }
         currentSubsegment.setNamespace(Namespace.AWS.toString());
 
-        if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
+        if (recorder.getCurrentSegment() != null &&
+            !(currentSubsegment instanceof NoOpSubSegment) &&
+            recorder.getCurrentSubsegment().shouldPropagate()) {
             request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(currentSubsegment).toString());
         }
     }

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

@@ -30,6 +30,7 @@
 import com.amazonaws.xray.emitters.DefaultEmitter;
 import com.amazonaws.xray.emitters.Emitter;
 import com.amazonaws.xray.entities.Subsegment;
+import com.amazonaws.xray.entities.SubsegmentImpl;
 import com.amazonaws.xray.entities.TraceHeader;
 import com.amazonaws.xray.strategy.sampling.NoSamplingStrategy;
 import org.junit.FixMethodOrder;
@@ -137,12 +138,17 @@ public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boole
         TraceHeader traceHeader = TraceHeader.fromString(request.getHeaders().get(TraceHeader.HEADER_KEY));
         String serviceEntityId = recorder.getCurrentSubsegment().getId();
 
-        assertThat(traceHeader.getSampled()).isEqualTo(
+        if (!sampled) {
+            assertThat(subsegment).isInstanceOf(SubsegmentImpl.class);
+        } else {
+            assertThat(subsegment).isInstanceOf(SubsegmentImpl.class);
+            assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
+            assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
+            assertThat(traceHeader.getSampled()).isEqualTo(
                 subsegment.isSampled() ?
                         TraceHeader.SampleDecision.SAMPLED :
                         TraceHeader.SampleDecision.NOT_SAMPLED);
-        assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
-        assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
+        }
 
         tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
 

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

@@ -25,7 +25,7 @@
 import java.util.concurrent.locks.ReentrantLock;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
-class NoOpSubSegment implements Subsegment {
+public class NoOpSubSegment implements Subsegment {
     private final Segment parentSegment;
     private final AWSXRayRecorder creator;
     private final boolean shouldPropagate;