improved ForcedSampling support (#232)

* improved ForcedSampling support

* removed setter from all and none strategies

* addressed comments (literally)

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

@@ -415,6 +415,14 @@ public Segment beginSegmentWithSampling(String name) {
                 segment.setRuleName(samplingResponse.getRuleName().get());
             }
 
+            return segment;
+        } else if (this.getSamplingStrategy().isForcedSamplingSupported()) {
+            Segment segment = beginSegment(name);
+            segment.setSampled(false);
+            if (samplingResponse.getRuleName().isPresent()) {
+                segment.setRuleName(samplingResponse.getRuleName().get());
+            }
+
             return segment;
         }
 
@@ -981,7 +989,8 @@ public final IdGenerator getIdGenerator() {
     }
 
     /**
-     * Checks whether the current {@code SamplingStrategy} supports forced sampling.
+     * Checks whether the current {@code SamplingStrategy} supports forced sampling. Use with caution, since segments sampled in
+     * this manner will not count towards your sampling statistic counts.
      *
      * @return true if forced sampling is supported and the current segment was changed from not sampled to sampled.
      */

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/CentralizedSamplingStrategy.java

@@ -51,23 +51,29 @@ public class CentralizedSamplingStrategy implements SamplingStrategy {
     private final LocalizedSamplingStrategy fallback;
     private final RulePoller rulePoller;
     private final TargetPoller targetPoller;
+    private final boolean forcedSamplingSupport;
 
     private boolean isStarted = false;
 
     public CentralizedSamplingStrategy() {
-        this.manifest = new CentralizedManifest();
-        this.fallback = new LocalizedSamplingStrategy();
-        UnsignedXrayClient client = new UnsignedXrayClient();
-        this.rulePoller = new RulePoller(client, manifest, Clock.systemUTC());
-        this.targetPoller = new TargetPoller(client, manifest, Clock.systemUTC());
+        this(LocalizedSamplingStrategy.DEFAULT_RULES, false);
     }
 
-    public CentralizedSamplingStrategy(URL ruleLocation) {
+    public CentralizedSamplingStrategy(@Nullable URL ruleLocation) {
+        this(ruleLocation, false);
+    }
+
+    public CentralizedSamplingStrategy(boolean forcedSamplingSupport) {
+        this(LocalizedSamplingStrategy.DEFAULT_RULES, forcedSamplingSupport);
+    }
+
+    public CentralizedSamplingStrategy(@Nullable URL ruleLocation, boolean forcedSamplingSupport) {
         this.manifest = new CentralizedManifest();
         this.fallback = new LocalizedSamplingStrategy(ruleLocation);
         UnsignedXrayClient client = new UnsignedXrayClient();
         this.rulePoller = new RulePoller(client, manifest, Clock.systemUTC());
         this.targetPoller = new TargetPoller(client, manifest, Clock.systemUTC());
+        this.forcedSamplingSupport = forcedSamplingSupport;
     }
 
     @Nullable
@@ -148,8 +154,6 @@ private synchronized void startPoller() {
 
     @Override
     public boolean isForcedSamplingSupported() {
-        //TODO address this
-        return false;
+        return forcedSamplingSupport;
     }
-
 }

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/LocalizedSamplingStrategy.java

@@ -35,7 +35,10 @@ public class LocalizedSamplingStrategy implements SamplingStrategy {
     private static final Log logger =
         LogFactory.getLog(LocalizedSamplingStrategy.class);
 
-    private static final URL DEFAULT_RULES;
+    private final boolean forcedSamplingSupport;
+
+    // Visible for other sampling strategies
+    static final URL DEFAULT_RULES;
 
     static {
         URL defaultRules =
@@ -64,11 +67,20 @@ public class LocalizedSamplingStrategy implements SamplingStrategy {
     private SamplingRule defaultRule;
 
     public LocalizedSamplingStrategy() {
-        this(DEFAULT_RULES);
+        this(DEFAULT_RULES, false);
     }
 
     public LocalizedSamplingStrategy(@Nullable URL ruleLocation) {
+        this(ruleLocation, false);
+    }
+
+    public LocalizedSamplingStrategy(boolean forcedSamplingSupport) {
+        this(DEFAULT_RULES, forcedSamplingSupport);
+    }
+
+    public LocalizedSamplingStrategy(@Nullable URL ruleLocation, boolean forcedSamplingSupport) {
         this.samplingRulesLocation = ruleLocation;
+        this.forcedSamplingSupport = forcedSamplingSupport;
 
         SamplingRuleManifest manifest = getRuleManifest(ruleLocation);
         if (manifest != null) {
@@ -182,7 +194,6 @@ private boolean shouldTrace(@Nullable SamplingRule samplingRule) {
 
     @Override
     public boolean isForcedSamplingSupported() {
-        //TODO address this
-        return false;
+        return forcedSamplingSupport;
     }
 }

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

@@ -17,6 +17,8 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
 
 import com.amazonaws.xray.contexts.LambdaSegmentContext;
 import com.amazonaws.xray.contexts.LambdaSegmentContextResolver;
@@ -38,7 +40,6 @@
 import com.amazonaws.xray.strategy.LogErrorContextMissingStrategy;
 import com.amazonaws.xray.strategy.RuntimeErrorContextMissingStrategy;
 import com.amazonaws.xray.strategy.sampling.LocalizedSamplingStrategy;
-import com.amazonaws.xray.strategy.sampling.SamplingRequest;
 import com.amazonaws.xray.strategy.sampling.SamplingResponse;
 import com.amazonaws.xray.strategy.sampling.SamplingStrategy;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
@@ -68,7 +69,9 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.MethodSorters;
 import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
 import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PowerMockIgnore;
 import org.powermock.core.classloader.annotations.PrepareForTest;
@@ -86,6 +89,9 @@ public class AWSXRayRecorderTest {
 
     private static ExecutorService threadExecutor;
 
+    @Mock
+    private SamplingStrategy mockSamplingStrategy;
+
     @Rule
     public EnvironmentVariables environmentVariables = new EnvironmentVariables();
     @Rule
@@ -103,6 +109,7 @@ public static void stopExecutor() {
 
     @Before
     public void setupAWSXRay() {
+        MockitoAnnotations.initMocks(this);
         Emitter blankEmitter = Mockito.mock(Emitter.class);
         LocalizedSamplingStrategy defaultSamplingStrategy = new LocalizedSamplingStrategy();
         Mockito.doReturn(true).when(blankEmitter).sendSegment(Mockito.anyObject());
@@ -270,7 +277,7 @@ public void testNotSendingUnsampledSegment() {
         segment.setSampled(false);
         recorder.endSegment();
 
-        Mockito.verify(mockEmitter, Mockito.times(0)).sendSegment(Mockito.any());
+        Mockito.verify(mockEmitter, Mockito.times(0)).sendSegment(any());
     }
 
     @Test
@@ -283,7 +290,7 @@ public void testSegmentEmitted() {
         recorder.endSubsegment();
         recorder.endSegment();
 
-        Mockito.verify(mockEmitter, Mockito.times(1)).sendSegment(Mockito.any());
+        Mockito.verify(mockEmitter, Mockito.times(1)).sendSegment(any());
     }
 
     @Test
@@ -296,7 +303,7 @@ public void testExplicitSubsegmentEmitted() {
         recorder.endSubsegment(subsegment);
         recorder.endSegment();
 
-        Mockito.verify(mockEmitter, Mockito.times(1)).sendSegment(Mockito.any());
+        Mockito.verify(mockEmitter, Mockito.times(1)).sendSegment(any());
     }
 
     @Test
@@ -309,7 +316,7 @@ public void testDummySegmentNotEmitted() {
         recorder.endSubsegment();
         recorder.endSegment();
 
-        Mockito.verify(mockEmitter, Mockito.times(0)).sendSegment(Mockito.any());
+        Mockito.verify(mockEmitter, Mockito.times(0)).sendSegment(any());
     }
 
     @Test
@@ -378,7 +385,7 @@ public void testSubsegmentNotEmittedWithoutExceptionInLambdaInitContext() {
         recorder.createSubsegment("test", () -> {
         });
 
-        Mockito.verify(mockEmitter, Mockito.times(0)).sendSubsegment(Mockito.any());
+        Mockito.verify(mockEmitter, Mockito.times(0)).sendSubsegment(any());
     }
 
     @Test
@@ -851,7 +858,9 @@ public void noOpSubsegmentWithParent() {
 
     @Test
     public void testBeginSegmentWithSamplingDoesSample() {
-        AWSXRay.getGlobalRecorder().setSamplingStrategy(new TestSamplingStrategy(true));
+        SamplingResponse response = new SamplingResponse(true, "rule");
+        when(mockSamplingStrategy.shouldTrace(any())).thenReturn(response);
+        AWSXRay.getGlobalRecorder().setSamplingStrategy(mockSamplingStrategy);
 
         Segment segment = AWSXRay.beginSegmentWithSampling("test");
         assertThat(segment.isSampled()).isTrue();
@@ -861,27 +870,29 @@ public void testBeginSegmentWithSamplingDoesSample() {
 
     @Test
     public void testBeginSegmentWithSamplingDoesNotSample() {
-        AWSXRay.getGlobalRecorder().setSamplingStrategy(new TestSamplingStrategy(false));
+        SamplingResponse response = new SamplingResponse(false, "rule");
+        when(mockSamplingStrategy.shouldTrace(any())).thenReturn(response);
+        AWSXRay.getGlobalRecorder().setSamplingStrategy(mockSamplingStrategy);
 
         Segment segment = AWSXRay.beginSegmentWithSampling("test");
         assertThat(segment.isSampled()).isFalse();
+
+        segment.setUser("user");
+        assertThat(segment.getUser()).isEmpty();  // Loose way to test that segment is a no-op
     }
 
-    private static class TestSamplingStrategy implements SamplingStrategy {
-        boolean sampled;
+    @Test
+    public void testBeginSegmentWithForcedSampling() {
+        SamplingResponse response = new SamplingResponse(false, "rule");
+        when(mockSamplingStrategy.isForcedSamplingSupported()).thenReturn(true);
+        when(mockSamplingStrategy.shouldTrace(any())).thenReturn(response);
+        AWSXRay.getGlobalRecorder().setSamplingStrategy(mockSamplingStrategy);
 
-        TestSamplingStrategy(boolean sampled) {
-            this.sampled = sampled;
-        }
+        Segment segment = AWSXRay.beginSegmentWithSampling("test");
+        assertThat(segment.isSampled()).isFalse();
 
-        @Override
-        public SamplingResponse shouldTrace(SamplingRequest sampleRequest) {
-            return new SamplingResponse(sampled, "rule");
-        }
+        segment.setUser("user");
+        assertThat(segment.getUser()).isEqualTo("user");  // Loose way to test that segment is real
 
-        @Override
-        public boolean isForcedSamplingSupported() {
-            return false;
-        }
     }
 }