Ensure all XRay Sampler functionality is under ParentBased logic (#1488)

Author: jj22ee

aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSampler.java

@@ -56,6 +56,7 @@ public final class AwsXrayRemoteSampler implements Sampler, Closeable {
   @Nullable private volatile ScheduledFuture<?> pollFuture;
   @Nullable private volatile ScheduledFuture<?> fetchTargetsFuture;
   @Nullable private volatile GetSamplingRulesResponse previousRulesResponse;
+  @Nullable private volatile XrayRulesSampler internalXrayRulesSampler;
   private volatile Sampler sampler;
 
   /**
@@ -125,15 +126,16 @@ private void getAndUpdateSampler() {
       GetSamplingRulesResponse response =
           client.getSamplingRules(GetSamplingRulesRequest.create(null));
       if (!response.equals(previousRulesResponse)) {
-        sampler =
+        updateInternalSamplers(
             new XrayRulesSampler(
                 clientId,
                 resource,
                 clock,
                 initialSampler,
                 response.getSamplingRules().stream()
                     .map(SamplingRuleRecord::getRule)
-                    .collect(Collectors.toList()));
+                    .collect(Collectors.toList())));
+
         previousRulesResponse = response;
         ScheduledFuture<?> existingFetchTargetsFuture = fetchTargetsFuture;
         if (existingFetchTargetsFuture != null) {
@@ -170,11 +172,11 @@ Duration getNextSamplerUpdateScheduledDuration() {
   }
 
   private void fetchTargets() {
-    if (!(sampler instanceof XrayRulesSampler)) {
+    if (this.internalXrayRulesSampler == null) {
       throw new IllegalStateException("Programming bug.");
     }
 
-    XrayRulesSampler xrayRulesSampler = (XrayRulesSampler) sampler;
+    XrayRulesSampler xrayRulesSampler = this.internalXrayRulesSampler;
     try {
       Date now = Date.from(Instant.ofEpochSecond(0, clock.now()));
       List<SamplingStatisticsDocument> statistics = xrayRulesSampler.snapshot(now);
@@ -188,8 +190,7 @@ private void fetchTargets() {
       Map<String, SamplingTargetDocument> targets =
           response.getDocuments().stream()
               .collect(Collectors.toMap(SamplingTargetDocument::getRuleName, Function.identity()));
-      sampler =
-          xrayRulesSampler = xrayRulesSampler.withTargets(targets, requestedTargetRuleNames, now);
+      updateInternalSamplers(xrayRulesSampler.withTargets(targets, requestedTargetRuleNames, now));
     } catch (Throwable t) {
       // Might be a transient API failure, try again after a default interval.
       fetchTargetsFuture =
@@ -225,6 +226,11 @@ private static String generateClientId() {
     return new String(clientIdChars);
   }
 
+  private void updateInternalSamplers(XrayRulesSampler xrayRulesSampler) {
+    this.internalXrayRulesSampler = xrayRulesSampler;
+    this.sampler = Sampler.parentBased(internalXrayRulesSampler);
+  }
+
   // Visible for testing
   XraySamplerClient getClient() {
     return client;

aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/SamplingRuleApplier.java

@@ -464,11 +464,11 @@ public String toString() {
   }
 
   private Sampler createRateLimited(int numPerSecond) {
-    return Sampler.parentBased(new RateLimitingSampler(numPerSecond, clock));
+    return new RateLimitingSampler(numPerSecond, clock);
   }
 
   private static Sampler createFixedRate(double rate) {
-    return Sampler.parentBased(Sampler.traceIdRatioBased(rate));
+    return Sampler.traceIdRatioBased(rate);
   }
 
   // We keep track of sampling requests and decisions to report to X-Ray to allow it to allocate

aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java

@@ -168,6 +168,25 @@ void defaultInitialSampler() {
     }
   }
 
+  @Test
+  void parentBasedXraySamplerAfterDefaultSampler() {
+    rulesResponse.set(RULE_RESPONSE_1);
+    try (AwsXrayRemoteSampler samplerWithLongerPollingInterval =
+        AwsXrayRemoteSampler.newBuilder(Resource.empty())
+            .setInitialSampler(Sampler.alwaysOn())
+            .setEndpoint(server.httpUri().toString())
+            .setPollingInterval(Duration.ofMillis(5))
+            .build()) {
+      await()
+          .pollDelay(Duration.ofMillis(10))
+          .untilAsserted(
+              () -> {
+                assertThat(sampler.getDescription())
+                    .startsWith("AwsXrayRemoteSampler{ParentBased{root:XrayRulesSampler{[");
+              });
+    }
+  }
+
   // https://github.com/open-telemetry/opentelemetry-java-contrib/issues/376
   @Test
   void testJitterTruncation() {