[Lambda] Do not propagate trace header if dummy subsegment

aws · May 28, 2024 · 91a0f62 · 91a0f62
1 parent c275035
commit 91a0f62
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 8 deletions.
diff --git a/...dk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java b/...dk-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());
         }
 

diff --git a/...pache-http/src/test/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClientTest.java b/...pache-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

diff --git a/...rder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java b/...rder-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;
         }
 

diff --git a/aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java b/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());
         }
     }

diff --git a/...order-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java b/...order-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)));
 

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSubSegment.java b/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;