Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

improved ForcedSampling support #232

Merged
merged 4 commits into from
Nov 12, 2020
Merged

improved ForcedSampling support #232

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9c1ad16
improved ForcedSampling support
Nov 9, 2020
5a5eb91
removed setter from all and none strategies
Nov 9, 2020
05cb1db
addressed comments (literally)
Nov 10, 2020
dc7ba94
Merge branch 'master' into forced-sample
willarmiros Nov 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Copy link
Contributor

@srprash srprash Nov 9, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From what I understand of ForcedSampling is that it can be used to force sample a segment which is un-sampled when created. Or is it only to keep the segment in context even though it is un-sampled?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Forced sampling forces the creation of a regular instead of no-op segment when not sampled, which also places it in context. Then forceSamplingOfCurrentSegment can be used in conjunction to actually switch an unsampled segment to sampled.

if (samplingResponse.getRuleName().isPresent()) {
segment.setRuleName(samplingResponse.getRuleName().get());
}

return segment;
}

Expand Down Expand Up @@ -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.
*/
Expand Down
22 changes: 13 additions & 9 deletions ...-core/src/main/java/com/amazonaws/xray/strategy/sampling/CentralizedSamplingStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -148,8 +154,6 @@ private synchronized void startPoller() {

@Override
public boolean isForcedSamplingSupported() {
//TODO address this
return false;
return forcedSamplingSupport;
}

}
19 changes: 15 additions & 4 deletions ...dk-core/src/main/java/com/amazonaws/xray/strategy/sampling/LocalizedSamplingStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -182,7 +194,6 @@ private boolean shouldTrace(@Nullable SamplingRule samplingRule) {

@Override
public boolean isForcedSamplingSupported() {
//TODO address this
return false;
return forcedSamplingSupport;
}
}
53 changes: 32 additions & 21 deletions aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -86,6 +89,9 @@ public class AWSXRayRecorderTest {

private static ExecutorService threadExecutor;

@Mock
private SamplingStrategy mockSamplingStrategy;

@Rule
public EnvironmentVariables environmentVariables = new EnvironmentVariables();
@Rule
Expand All @@ -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());
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
Expand All @@ -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;
}
}
}