Simplifying Oversampling Mitigation

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

@@ -32,7 +32,6 @@
 import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
 import com.amazonaws.xray.internal.FastIdGenerator;
 import com.amazonaws.xray.internal.IdGenerator;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import com.amazonaws.xray.internal.SecureIdGenerator;
 import com.amazonaws.xray.listeners.SegmentListener;
 import com.amazonaws.xray.strategy.ContextMissingStrategy;
@@ -636,10 +635,7 @@ public Subsegment beginSubsegmentWithoutSampling(String name) {
             // context to be propagated downstream
             return Subsegment.noOp(this, false);
         }
-        return context.beginSubsegmentWithSamplingOverride(
-                this,
-                name,
-                SamplingStrategyOverride.FALSE);
+        return context.beginSubsegmentWithoutSampling(this, name);
     }
 
     /**

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

@@ -25,7 +25,6 @@
 import com.amazonaws.xray.entities.TraceHeader.SampleDecision;
 import com.amazonaws.xray.entities.TraceID;
 import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import com.amazonaws.xray.listeners.SegmentListener;
 import java.util.List;
 import java.util.Objects;
@@ -62,11 +61,17 @@ private static FacadeSegment newFacadeSegment(AWSXRayRecorder recorder, String n
     }
 
     @Override
-    public Subsegment beginSubsegmentWithSamplingOverride(
+    public Subsegment beginSubsegmentWithoutSampling(
             AWSXRayRecorder recorder,
-            String name,
-            SamplingStrategyOverride samplingStrategyOverride) {
+            String name) {
 
+        Subsegment subsegment = beginSubsegment(recorder, name);
+        subsegment.setSampledFalse();
+        return subsegment;
+    }
+
+    @Override
+    public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
         if (logger.isDebugEnabled()) {
             logger.debug("Beginning subsegment named: " + name);
         }
@@ -75,12 +80,11 @@ public Subsegment beginSubsegmentWithSamplingOverride(
         if (entity == null) { // First subsgment of a subsegment branch.
             Segment parentSegment = newFacadeSegment(recorder, name);
 
-            boolean isRecording = parentSegment.isRecording() &&
-                    samplingStrategyOverride == SamplingStrategyOverride.DISABLED;
+            boolean isRecording = parentSegment.isRecording();
 
             Subsegment subsegment = isRecording
-                    ? new SubsegmentImpl(recorder, name, parentSegment, samplingStrategyOverride)
-                    : Subsegment.noOp(parentSegment, recorder, samplingStrategyOverride);
+                    ? new SubsegmentImpl(recorder, name, parentSegment)
+                    : Subsegment.noOp(parentSegment, recorder);
             subsegment.setParent(parentSegment);
             // Enable FacadeSegment to keep track of its subsegments for subtree streaming
             parentSegment.addSubsegment(subsegment);
@@ -97,12 +101,11 @@ public Subsegment beginSubsegmentWithSamplingOverride(
             }
             Segment parentSegment = parentSubsegment.getParentSegment();
 
-            boolean isRecording = parentSegment.isRecording() &&
-                    samplingStrategyOverride == SamplingStrategyOverride.DISABLED;
+            boolean isRecording = parentSegment.isRecording();
 
             Subsegment subsegment = isRecording
-                    ? new SubsegmentImpl(recorder, name, parentSegment, samplingStrategyOverride)
-                    : Subsegment.noOp(parentSegment, recorder, samplingStrategyOverride);
+                    ? new SubsegmentImpl(recorder, name, parentSegment)
+                    : Subsegment.noOp(parentSegment, recorder);
             subsegment.setParent(parentSubsegment);
             parentSubsegment.addSubsegment(subsegment);
             setTraceEntity(subsegment);
@@ -116,11 +119,6 @@ public Subsegment beginSubsegmentWithSamplingOverride(
         }
     }
 
-    @Override
-    public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
-        return beginSubsegmentWithSamplingOverride(recorder, name, SamplingStrategyOverride.DISABLED);
-    }
-
     @Override
     public void endSubsegment(AWSXRayRecorder recorder) {
         Entity current = getTraceEntity();

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

@@ -20,7 +20,6 @@
 import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Segment;
 import com.amazonaws.xray.entities.Subsegment;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import java.util.Objects;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -66,10 +65,9 @@ default void clearTraceEntity() {
 
     Subsegment beginSubsegment(AWSXRayRecorder recorder, String name);
 
-    Subsegment beginSubsegmentWithSamplingOverride(
+    Subsegment beginSubsegmentWithoutSampling(
             AWSXRayRecorder recorder,
-            String name,
-            SamplingStrategyOverride samplingStrategyOverride);
+            String name);
 
     void endSubsegment(AWSXRayRecorder recorder);
 }

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

@@ -22,7 +22,6 @@
 import com.amazonaws.xray.entities.SubsegmentImpl;
 import com.amazonaws.xray.exceptions.SegmentNotFoundException;
 import com.amazonaws.xray.exceptions.SubsegmentNotFoundException;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import com.amazonaws.xray.listeners.SegmentListener;
 import java.util.List;
 import java.util.Objects;
@@ -34,23 +33,29 @@ public class ThreadLocalSegmentContext implements SegmentContext {
         LogFactory.getLog(ThreadLocalSegmentContext.class);
 
     @Override
-    public Subsegment beginSubsegmentWithSamplingOverride(
+    public Subsegment beginSubsegmentWithoutSampling(
             AWSXRayRecorder recorder,
-            String name,
-            SamplingStrategyOverride samplingStrategyOverride) {
+            String name) {
+        Subsegment subsegment = beginSubsegment(recorder, name);
+        subsegment.setSampledFalse();
+        return subsegment;
+    }
+
+    @Override
+    public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
         Entity current = getTraceEntity();
         if (current == null) {
             recorder.getContextMissingStrategy().contextMissing("Failed to begin subsegment named '" + name
                     + "': segment cannot be found.", SegmentNotFoundException.class);
-            return Subsegment.noOp(recorder, false, samplingStrategyOverride);
+            return Subsegment.noOp(recorder, false);
         }
         if (logger.isDebugEnabled()) {
             logger.debug("Beginning subsegment named: " + name);
         }
         Segment parentSegment = current.getParentSegment();
-        Subsegment subsegment = parentSegment.isRecording() && (samplingStrategyOverride == SamplingStrategyOverride.DISABLED)
-                ? new SubsegmentImpl(recorder, name, parentSegment, samplingStrategyOverride)
-                : Subsegment.noOp(parentSegment, recorder, samplingStrategyOverride);
+        Subsegment subsegment = parentSegment.isRecording()
+                ? new SubsegmentImpl(recorder, name, parentSegment)
+                : Subsegment.noOp(parentSegment, recorder);
         subsegment.setParent(current);
         current.addSubsegment(subsegment);
         setTraceEntity(subsegment);
@@ -63,11 +68,6 @@ public Subsegment beginSubsegmentWithSamplingOverride(
         return subsegment;
     }
 
-    @Override
-    public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
-        return beginSubsegmentWithSamplingOverride(recorder, name, SamplingStrategyOverride.DISABLED);
-    }
-
     @Override
     public void endSubsegment(AWSXRayRecorder recorder) {
         Entity current = getTraceEntity();
@@ -87,8 +87,7 @@ public void endSubsegment(AWSXRayRecorder recorder) {
                     .filter(Objects::nonNull)
                     .forEach(listener -> listener.beforeEndSubsegment(currentSubsegment));
 
-            if ((currentSubsegment.end() &&
-                    currentSubsegment.getSamplingStrategyOverride() == SamplingStrategyOverride.DISABLED)) {
+            if (currentSubsegment.end() && currentSubsegment.isSampled()) {
                 recorder.sendSegment(currentSubsegment.getParentSegment());
             } else {
                 if (recorder.getStreamingStrategy().requiresStreaming(currentSubsegment.getParentSegment())) {

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

@@ -16,7 +16,7 @@
 package com.amazonaws.xray.entities;
 
 import com.amazonaws.xray.AWSXRayRecorder;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -44,21 +44,20 @@ public class DummySubsegment implements Subsegment {
     private TraceID traceId;
     private Segment parentSegment;
 
-    private SamplingStrategyOverride samplingStrategyOverride;
+    @JsonIgnore
+    private boolean isSampled;
+
+    @JsonIgnore
+    private boolean isRecording;
 
     public DummySubsegment(AWSXRayRecorder creator) {
         this(creator, TraceID.create(creator));
     }
 
-    public DummySubsegment(AWSXRayRecorder creator, TraceID traceId, SamplingStrategyOverride samplingStrategyOverride) {
+    public DummySubsegment(AWSXRayRecorder creator, TraceID traceId) {
         this.creator = creator;
         this.traceId = traceId;
         this.parentSegment = new DummySegment(creator);
-        this.samplingStrategyOverride = samplingStrategyOverride;
-    }
-
-    public DummySubsegment(AWSXRayRecorder creator, TraceID traceId) {
-        this(creator, traceId, SamplingStrategyOverride.DISABLED);
     }
 
     @Override
@@ -396,7 +395,21 @@ public void removeSubsegment(Subsegment subsegment) {
     }
 
     @Override
-    public SamplingStrategyOverride getSamplingStrategyOverride() {
-        return samplingStrategyOverride;
+    @JsonIgnore
+    public boolean isSampled() {
+        return isSampled;
+    }
+
+    @Override
+    @JsonIgnore
+    public boolean isRecording() {
+        return isRecording;
+    }
+
+    @Override
+    @JsonIgnore
+    public void setSampledFalse() {
+        isSampled = false;
+        isRecording = false;
     }
 }

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

@@ -16,7 +16,6 @@
 package com.amazonaws.xray.entities;
 
 import com.amazonaws.xray.AWSXRayRecorder;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import java.util.List;
 import java.util.Map;
@@ -34,30 +33,24 @@ class NoOpSubSegment implements Subsegment {
     private volatile Entity parent;
 
     @JsonIgnore
-    private SamplingStrategyOverride samplingStrategyOverride;
+    private boolean isSampled;
+
+    @JsonIgnore
+    private boolean isRecording;
 
     NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator) {
         this(parentSegment, creator, true);
     }
 
-    NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator, boolean shouldPropagate) {
-        this(parentSegment, creator, shouldPropagate, SamplingStrategyOverride.DISABLED);
-    }
-
-    NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator, SamplingStrategyOverride samplingStrategyOverride) {
-        this(parentSegment, creator, true, samplingStrategyOverride);
-    }
 
     NoOpSubSegment(
             Segment parentSegment,
             AWSXRayRecorder creator,
-            boolean shouldPropagate,
-            SamplingStrategyOverride samplingStrategyOverride) {
+            boolean shouldPropagate) {
         this.parentSegment = parentSegment;
         this.creator = creator;
         this.shouldPropagate = shouldPropagate;
         parent = parentSegment;
-        this.samplingStrategyOverride = samplingStrategyOverride;
     }
 
     @Override
@@ -391,7 +384,21 @@ public void close() {
     }
 
     @Override
-    public SamplingStrategyOverride getSamplingStrategyOverride() {
-        return samplingStrategyOverride;
+    @JsonIgnore
+    public boolean isSampled() {
+        return isSampled;
+    }
+
+    @Override
+    @JsonIgnore
+    public boolean isRecording() {
+        return isRecording;
+    }
+
+    @Override
+    @JsonIgnore
+    public void setSampledFalse() {
+        isSampled = false;
+        isRecording = false;
     }
 }

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

@@ -16,21 +16,11 @@
 package com.amazonaws.xray.entities;
 
 import com.amazonaws.xray.AWSXRayRecorder;
-import com.amazonaws.xray.internal.SamplingStrategyOverride;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import java.util.Set;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 public interface Subsegment extends Entity {
-
-    static Subsegment noOp(AWSXRayRecorder recorder, boolean shouldPropagate, SamplingStrategyOverride samplingStrategyOverride) {
-        return new NoOpSubSegment(Segment.noOp(TraceID.invalid(), recorder), recorder, shouldPropagate, samplingStrategyOverride);
-    }
-
-    static Subsegment noOp(Segment parent, AWSXRayRecorder recorder, SamplingStrategyOverride samplingStrategyOverride) {
-        return new NoOpSubSegment(parent, recorder, samplingStrategyOverride);
-    }
-
     static Subsegment noOp(AWSXRayRecorder recorder, boolean shouldPropagate) {
         return new NoOpSubSegment(Segment.noOp(TraceID.invalid(), recorder), recorder, shouldPropagate);
     }
@@ -119,15 +109,12 @@ static Subsegment noOp(Segment parent, AWSXRayRecorder recorder) {
     @Override
     void close();
 
-    SamplingStrategyOverride getSamplingStrategyOverride();
+    @JsonIgnore
+    boolean isSampled();
 
     @JsonIgnore
-    default boolean isSampled() {
-        return getParentSegment().isSampled() && getSamplingStrategyOverride() == SamplingStrategyOverride.DISABLED;
-    }
+    boolean isRecording();
 
     @JsonIgnore
-    default boolean isRecording() {
-        return getParentSegment().isRecording() && getSamplingStrategyOverride() == SamplingStrategyOverride.DISABLED;
-    }
+    void setSampledFalse();
 }