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