diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java index db70ad24..792b9423 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java +++ b/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. */ diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/CentralizedSamplingStrategy.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/CentralizedSamplingStrategy.java index f7f67f33..72307f69 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/CentralizedSamplingStrategy.java +++ b/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; } - } diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/LocalizedSamplingStrategy.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/LocalizedSamplingStrategy.java index 888e12d5..92a44e68 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/strategy/sampling/LocalizedSamplingStrategy.java +++ b/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; } } diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java index 01262863..25154758 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java +++ b/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; - } } }