Revert aws#403 and aws#404

majanjua-amzn · Jul 12, 2024 · 35be8fd · 35be8fd
1 parent 63ace4a
commit 35be8fd
Show file tree

Hide file tree

Showing 9 changed files with 29 additions and 71 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,7 +18,6 @@
 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;
@@ -104,7 +103,7 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
         subsegment.setNamespace(Namespace.REMOTE.toString());
         Segment parentSegment = subsegment.getParentSegment();
 
-        if (!(subsegment instanceof NoOpSubSegment) && subsegment.shouldPropagate()) {
+        if (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-67891233-abcdef012345678912345678;Sampled=0")));
+                               equalTo("Root=1-00000000-000000000000000000000000;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,7 +21,6 @@
 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;
@@ -293,7 +292,7 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
         }
 
         Subsegment subsegment = executionAttributes.getAttribute(entityKey);
-        if (!subsegment.shouldPropagate() || subsegment instanceof NoOpSubSegment) {
+        if (!subsegment.shouldPropagate()) {
             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,7 +31,6 @@
 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;
@@ -192,9 +191,7 @@ public void beforeRequest(Request<?> request) {
         }
         currentSubsegment.setNamespace(Namespace.AWS.toString());
 
-        if (recorder.getCurrentSegment() != null &&
-            !(currentSubsegment instanceof NoOpSubSegment) &&
-            recorder.getCurrentSubsegment().shouldPropagate()) {
+        if (recorder.getCurrentSegment() != null && 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,7 +30,6 @@
 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;
@@ -138,17 +137,12 @@ public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boole
         TraceHeader traceHeader = TraceHeader.fromString(request.getHeaders().get(TraceHeader.HEADER_KEY));
         String serviceEntityId = recorder.getCurrentSubsegment().getId();
 
-        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(
+        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/...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,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;
@@ -35,44 +35,40 @@ 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));
     }
 
     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 subsegment of a subsegment branch
-            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()
-                );
-            }
+        if (entity == null) { // First subsgment of a subsegment branch.
+            Segment parentSegment = newFacadeSegment(recorder, name);
 
             boolean isRecording = parentSegment.isRecording();
 
@@ -149,8 +145,6 @@ 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;
 
-public class NoOpSegment implements Segment {
+class NoOpSegment implements Segment {
 
     private final TraceID traceId;
     private final AWSXRayRecorder creator;

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;
 
-public class NoOpSubSegment implements Subsegment {
+class NoOpSubSegment implements Subsegment {
     private final Segment parentSegment;
     private final AWSXRayRecorder creator;
     private final boolean shouldPropagate;

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,7 +22,6 @@
 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;
@@ -50,10 +49,6 @@ 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() {
@@ -68,14 +63,14 @@ public void setupAWSXRay() {
     }
 
     @Test
-    void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
-        testContextResultsInNoOpSegmentParent();
+    void testBeginSubsegmentWithNullTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
+        testContextResultsInFacadeSegmentParent();
     }
 
     @Test
     @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = "a")
-    void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInANoOpSegmentParent() {
-        testContextResultsInNoOpSegmentParent();
+    void testBeginSubsegmentWithIncompleteTraceHeaderEnvironmentVariableResultsInAFacadeSegmentParent() {
+        testContextResultsInFacadeSegmentParent();
     }
 
     @Test
@@ -90,18 +85,6 @@ 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() {
@@ -166,12 +149,4 @@ 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();
-    }
 }