From 7bb0a731077d26d44ee1e853a30697fe87d61df6 Mon Sep 17 00:00:00 2001
From: Jonathan Lee <107072447+jj22ee@users.noreply.github.com>
Date: Mon, 7 Oct 2024 09:48:51 -0700
Subject: [PATCH] Ensure all XRay Sampler functionality is under ParentBased
 logic (#1488)

---
 .../contrib/awsxray/AwsXrayRemoteSampler.java | 18 ++++++++++++------
 .../contrib/awsxray/SamplingRuleApplier.java  |  4 ++--
 .../awsxray/AwsXrayRemoteSamplerTest.java     | 19 +++++++++++++++++++
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSampler.java b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSampler.java
index 9b5a2e7e6..ad9b72a2c 100644
--- a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSampler.java
+++ b/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,7 +126,7 @@ private void getAndUpdateSampler() {
       GetSamplingRulesResponse response =
           client.getSamplingRules(GetSamplingRulesRequest.create(null));
       if (!response.equals(previousRulesResponse)) {
-        sampler =
+        updateInternalSamplers(
             new XrayRulesSampler(
                 clientId,
                 resource,
@@ -133,7 +134,8 @@ private void getAndUpdateSampler() {
                 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;
diff --git a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/SamplingRuleApplier.java b/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/SamplingRuleApplier.java
index 1ef8abf50..53ab1efaf 100644
--- a/aws-xray/src/main/java/io/opentelemetry/contrib/awsxray/SamplingRuleApplier.java
+++ b/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
diff --git a/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java b/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java
index 654a4d57a..4e5cd13bc 100644
--- a/aws-xray/src/test/java/io/opentelemetry/contrib/awsxray/AwsXrayRemoteSamplerTest.java
+++ b/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() {