Improvements to no-op subsegment behavior (#294)

* various improvements to no-op subsegment behavior

* removed setter

Co-authored-by: Anuraag Agrawal <aanuraag@amazon.co.jp>

Author: willarmiros

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

@@ -110,10 +110,12 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
         subsegment.setNamespace(Namespace.REMOTE.toString());
         Segment parentSegment = subsegment.getParentSegment();
 
-        TraceHeader header = new TraceHeader(parentSegment.getTraceId(),
-                                parentSegment.isSampled() ? subsegment.getId() : null,
-                                parentSegment.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
-        request.addHeader(TraceHeader.HEADER_KEY, header.toString());
+        if (subsegment.shouldPropagate()) {
+            TraceHeader header = new TraceHeader(parentSegment.getTraceId(),
+                parentSegment.isSampled() ? subsegment.getId() : null,
+                parentSegment.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
+            request.addHeader(TraceHeader.HEADER_KEY, header.toString());
+        }
 
         Map<String, Object> requestInformation = new HashMap<>();
 

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

@@ -245,7 +245,7 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
         SdkHttpRequest httpRequest = context.httpRequest();
 
         Subsegment subsegment = executionAttributes.getAttribute(entityKey);
-        if (subsegment == null) {
+        if (!subsegment.shouldPropagate()) {
             return httpRequest;
         }
 

aws-xray-recorder-sdk-aws-sdk-v2/src/test/java/com/amazonaws/xray/interceptors/TracingInterceptorTest.java

@@ -15,6 +15,11 @@
 
 package com.amazonaws.xray.interceptors;
 
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.AWSXRayRecorderBuilder;
 import com.amazonaws.xray.emitters.Emitter;
@@ -40,10 +45,13 @@
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.core.async.EmptyPublisher;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.http.AbortableInputStream;
 import software.amazon.awssdk.http.ExecutableHttpRequest;
 import software.amazon.awssdk.http.HttpExecuteResponse;
 import software.amazon.awssdk.http.SdkHttpClient;
+import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.SdkHttpResponse;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
@@ -86,8 +94,8 @@ private SdkHttpClient mockSdkHttpClient(SdkHttpResponse response, String body) t
         ExecutableHttpRequest abortableCallable = Mockito.mock(ExecutableHttpRequest.class);
         SdkHttpClient mockClient = Mockito.mock(SdkHttpClient.class);
 
-        Mockito.when(mockClient.prepareRequest(Mockito.any())).thenReturn(abortableCallable);
-        Mockito.when(abortableCallable.call()).thenReturn(HttpExecuteResponse.builder()
+        when(mockClient.prepareRequest(Mockito.any())).thenReturn(abortableCallable);
+        when(abortableCallable.call()).thenReturn(HttpExecuteResponse.builder()
                 .response(response)
                 .responseBody(AbortableInputStream.create(
                         new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8))
@@ -99,7 +107,7 @@ private SdkHttpClient mockSdkHttpClient(SdkHttpResponse response, String body) t
 
     private SdkAsyncHttpClient mockSdkAsyncHttpClient(SdkHttpResponse response) {
         SdkAsyncHttpClient mockClient = Mockito.mock(SdkAsyncHttpClient.class);
-        Mockito.when(mockClient.execute(Mockito.any(AsyncExecuteRequest.class)))
+        when(mockClient.execute(Mockito.any(AsyncExecuteRequest.class)))
                .thenAnswer((Answer<CompletableFuture<Void>>) invocationOnMock -> {
                    AsyncExecuteRequest request = invocationOnMock.getArgument(0);
                    SdkAsyncHttpResponseHandler handler = request.responseHandler();
@@ -562,5 +570,22 @@ public void testAsync500Exception() {
             Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
         }
     }
+
+    @Test
+    public void testNoHeaderAddedWhenPropagationOff() {
+        Subsegment subsegment = Subsegment.noOp(AWSXRay.getGlobalRecorder(), false);
+        TracingInterceptor interceptor = new TracingInterceptor();
+        Context.ModifyHttpRequest context = Mockito.mock(Context.ModifyHttpRequest.class);
+        SdkHttpRequest mockRequest = Mockito.mock(SdkHttpRequest.class);
+        SdkHttpRequest.Builder mockRequestBuilder = Mockito.mock(SdkHttpRequest.Builder.class);
+        when(context.httpRequest()).thenReturn(mockRequest);
+        Mockito.lenient().when(context.httpRequest().toBuilder()).thenReturn(mockRequestBuilder);
+        ExecutionAttributes attributes = new ExecutionAttributes();
+        attributes.putAttribute(TracingInterceptor.entityKey, subsegment);
+
+        interceptor.modifyHttpRequest(context, attributes);
+
+        verify(mockRequest.toBuilder(), never()).appendHeader(anyString(), anyString());
+    }
 }
 

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

@@ -191,7 +191,7 @@ public void beforeRequest(Request<?> request) {
         }
         currentSubsegment.setNamespace(Namespace.AWS.toString());
 
-        if (recorder.getCurrentSegment() != null) {
+        if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
             TraceHeader header =
                 new TraceHeader(recorder.getCurrentSegment().getTraceId(),
                                 recorder.getCurrentSegment().isSampled() ? currentSubsegment.getId() : null,

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

@@ -609,7 +609,9 @@ public Subsegment beginSubsegment(String name) {
         if (context == null) {
             // No context available, we return a no-op subsegment so user code does not have to work around this. Based on
             // ContextMissingStrategy they will still know about the issue unless they explicitly opt-ed out.
-            return Subsegment.noOp(this);
+            // This no-op subsegment is different from unsampled no-op subsegments only in that it should not cause trace
+            // context to be propagated downstream
+            return Subsegment.noOp(this, false);
         }
         return context.beginSubsegment(this, name);
     }

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

@@ -107,7 +107,11 @@ public void endSubsegment(AWSXRayRecorder recorder) {
         Entity current = getTraceEntity();
         if (current instanceof Subsegment) {
             if (logger.isDebugEnabled()) {
-                logger.debug("Ending subsegment named: " + current.getName());
+                if (current.getName().isEmpty() && !current.getParentSegment().isSampled()) {
+                    logger.debug("Ending no-op subsegment");
+                } else {
+                    logger.debug("Ending subsegment named: " + current.getName());
+                }
             }
             Subsegment currentSubsegment = (Subsegment) current;
 
@@ -139,7 +143,8 @@ public void endSubsegment(AWSXRayRecorder recorder) {
             }
 
         } else {
-            throw new SubsegmentNotFoundException("Failed to end a subsegment: subsegment cannot be found.");
+            recorder.getContextMissingStrategy().contextMissing("Failed to end subsegment: subsegment cannot be found.",
+                SubsegmentNotFoundException.class);
         }
     }
 }

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

@@ -39,7 +39,7 @@ public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
         if (current == null) {
             recorder.getContextMissingStrategy().contextMissing("Failed to begin subsegment named '" + name
                                                                 + "': segment cannot be found.", SegmentNotFoundException.class);
-            return Subsegment.noOp(recorder);
+            return Subsegment.noOp(recorder, false);
         }
         if (logger.isDebugEnabled()) {
             logger.debug("Beginning subsegment named: " + name);
@@ -65,7 +65,11 @@ public void endSubsegment(AWSXRayRecorder recorder) {
         Entity current = getTraceEntity();
         if (current instanceof Subsegment) {
             if (logger.isDebugEnabled()) {
-                logger.debug("Ending subsegment named: " + current.getName());
+                if (current.getName().isEmpty() && !current.getParentSegment().isSampled()) {
+                    logger.debug("Ending no-op subsegment");
+                } else {
+                    logger.debug("Ending subsegment named: " + current.getName());
+                }
             }
             Subsegment currentSubsegment = (Subsegment) current;
 

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

@@ -27,7 +27,7 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
- * @deprecated Use {@link Subsegment#noOp(AWSXRayRecorder)}.
+ * @deprecated Use {@link Subsegment#noOp(AWSXRayRecorder, boolean)}.
  */
 @Deprecated
 public class DummySubsegment implements Subsegment {
@@ -340,6 +340,11 @@ public void setPrecursorIds(Set<String> precursorIds) {
     public void addPrecursorId(String precursorId) {
     }
 
+    @Override
+    public boolean shouldPropagate() {
+        return false;
+    }
+
     @Override
     public String streamSerialize() {
         return "";

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

@@ -27,12 +27,18 @@ class NoOpSubSegment implements Subsegment {
 
     private final Segment parentSegment;
     private final AWSXRayRecorder creator;
+    private final boolean shouldPropagate;
 
     private volatile Entity parent;
 
     NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator) {
+        this(parentSegment, creator, true);
+    }
+
+    NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator, boolean shouldPropagate) {
         this.parentSegment = parentSegment;
         this.creator = creator;
+        this.shouldPropagate = shouldPropagate;
         parent = parentSegment;
     }
 
@@ -347,6 +353,11 @@ public void setPrecursorIds(Set<String> precursorIds) {
     public void addPrecursorId(String precursorId) {
     }
 
+    @Override
+    public boolean shouldPropagate() {
+        return this.shouldPropagate;
+    }
+
     @Override
     public String streamSerialize() {
         return "";

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

@@ -21,8 +21,8 @@
 
 public interface Subsegment extends Entity {
 
-    static Subsegment noOp(AWSXRayRecorder recorder) {
-        return new NoOpSubSegment(Segment.noOp(TraceID.invalid(), recorder), recorder);
+    static Subsegment noOp(AWSXRayRecorder recorder, boolean shouldPropagate) {
+        return new NoOpSubSegment(Segment.noOp(TraceID.invalid(), recorder), recorder, shouldPropagate);
     }
 
     static Subsegment noOp(Segment parent, AWSXRayRecorder recorder) {
@@ -78,6 +78,12 @@ static Subsegment noOp(Segment parent, AWSXRayRecorder recorder) {
      */
     void addPrecursorId(String precursorId);
 
+    /**
+     * Determines if this subsegment should propagate its trace context downstream
+     * @return true if its trace context should be propagated downstream, false otherwise
+     */
+    boolean shouldPropagate();
+
     /**
      * Serializes the subsegment as a standalone String with enough information for the subsegment to be streamed on its own.
      * @return

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

@@ -38,6 +38,9 @@ public class SubsegmentImpl extends EntityImpl implements Subsegment {
     @GuardedBy("lock")
     private Set<String> precursorIds;
 
+    @GuardedBy("lock")
+    private boolean shouldPropagate;
+
     @SuppressWarnings("nullness")
     private SubsegmentImpl() {
         super();
@@ -48,6 +51,7 @@ public SubsegmentImpl(AWSXRayRecorder creator, String name, Segment parentSegmen
         this.parentSegment = parentSegment;
         parentSegment.incrementReferenceCount();
         this.precursorIds = new HashSet<>();
+        this.shouldPropagate = true;
     }
 
     @Override
@@ -125,6 +129,13 @@ public void setPrecursorIds(Set<String> precursorIds) {
         }
     }
 
+    @Override
+    public boolean shouldPropagate() {
+        synchronized (lock) {
+            return shouldPropagate;
+        }
+    }
+
     private ObjectNode getStreamSerializeObjectNode() {
         synchronized (lock) {
             ObjectNode obj = (ObjectNode) mapper.valueToTree(this);

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/javax/servlet/AWSXRayServletAsyncListener.java

@@ -15,7 +15,6 @@
 
 package com.amazonaws.xray.javax.servlet;
 
-import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.AWSXRayRecorder;
 import com.amazonaws.xray.entities.Entity;
 import java.io.IOException;
@@ -39,15 +38,7 @@ class AWSXRayServletAsyncListener implements AsyncListener {
         this.recorder = recorder;
     }
 
-    private AWSXRayRecorder getRecorder() {
-        if (recorder == null) {
-            recorder = AWSXRay.getGlobalRecorder();
-        }
-        return recorder;
-    }
-
     private void processEvent(AsyncEvent event) throws IOException {
-        AWSXRayRecorder recorder = getRecorder();
         Entity entity = (Entity) event.getSuppliedRequest().getAttribute(ENTITY_ATTRIBUTE_KEY);
         entity.run(() -> {
             if (event.getThrowable() != null) {

aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java

@@ -505,6 +505,7 @@ public void testBeginSubsegmentWhenMissingContext() {
         // No-op
         subsegment.setNamespace("foo");
         assertThat(subsegment.getNamespace()).isEmpty();
+        assertThat(subsegment.shouldPropagate()).isFalse();
     }
 
     @Test
@@ -736,8 +737,8 @@ public void testNoOpSegment() throws Exception {
         assertThat(segment.getCreator()).isEqualTo(AWSXRay.getGlobalRecorder());
         segment.setRuleName("foo");
         assertThat(segment.getParentSegment()).isSameAs(segment);
-        segment.getSubsegments().add(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
-        segment.addSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
+        segment.getSubsegments().add(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
+        segment.addSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
         assertThat(segment.getSubsegments()).isEmpty();
         segment.addException(new IllegalStateException());
         assertThat(segment.getReferenceCount()).isZero();
@@ -747,7 +748,7 @@ public void testNoOpSegment() throws Exception {
         assertThat(segment.getReferenceCount()).isZero();
         segment.decrementReferenceCount();
         assertThat(segment.getReferenceCount()).isZero();
-        segment.removeSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
+        segment.removeSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
         segment.setEmitted(true);
         segment.compareAndSetEmitted(false, true);
         assertThat(segment.isEmitted()).isFalse();
@@ -788,7 +789,7 @@ public void testAlwaysCreateTraceId() {
 
     @Test
     public void testNoOpSubsegment() throws Exception {
-        Subsegment subsegment = Subsegment.noOp(AWSXRay.getGlobalRecorder());
+        Subsegment subsegment = Subsegment.noOp(AWSXRay.getGlobalRecorder(), true);
 
         Map<String, Object> map = new HashMap<>();
         map.put("dog", "bark");
@@ -854,8 +855,8 @@ public void testNoOpSubsegment() throws Exception {
         subsegment.setCreator(AWSXRayRecorderBuilder.standard().build());
         assertThat(subsegment.getCreator()).isEqualTo(AWSXRay.getGlobalRecorder());
         assertThat(subsegment.getParentSegment().getTraceId()).isEqualTo(TraceID.invalid());
-        subsegment.getSubsegments().add(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
-        subsegment.addSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
+        subsegment.getSubsegments().add(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
+        subsegment.addSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
         assertThat(subsegment.getSubsegments()).isEmpty();
         subsegment.addException(new IllegalStateException());
         assertThat(subsegment.getReferenceCount()).isZero();
@@ -865,7 +866,7 @@ public void testNoOpSubsegment() throws Exception {
         assertThat(subsegment.getReferenceCount()).isZero();
         subsegment.decrementReferenceCount();
         assertThat(subsegment.getReferenceCount()).isZero();
-        subsegment.removeSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder()));
+        subsegment.removeSubsegment(Subsegment.noOp(AWSXRay.getGlobalRecorder(), true));
         subsegment.setEmitted(true);
         subsegment.compareAndSetEmitted(false, true);
         assertThat(subsegment.isEmitted()).isFalse();
@@ -935,4 +936,17 @@ public void testLogGroupFromEnvironment() {
         Set<AWSLogReference> logReferences = (Set<AWSLogReference>) segment.getAws().get("cloudwatch_logs");
         assertThat(logReferences).containsOnly(expected);
     }
+
+    @Test
+    public void testUnsampledSubsegmentPropagation() {
+        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
+            .withSamplingStrategy(new NoSamplingStrategy())
+            .build();
+
+        Segment segment = recorder.beginSegmentWithSampling("test");
+        Subsegment subsegment = recorder.beginSubsegment("test");
+
+        assertThat(segment.isSampled()).isFalse();
+        assertThat(subsegment.shouldPropagate()).isTrue();
+    }
 }