Fetch and update rules with sampling targets. (#53)

* Fetch and update rules with sampling targets.

Author: Anuraag Agrawal

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

@@ -7,42 +7,52 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.contrib.awsxray.GetSamplingRulesResponse.SamplingRuleRecord;
+import io.opentelemetry.contrib.awsxray.GetSamplingTargetsRequest.SamplingStatisticsDocument;
+import io.opentelemetry.contrib.awsxray.GetSamplingTargetsResponse.SamplingTargetDocument;
+import io.opentelemetry.sdk.common.Clock;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 import java.io.Closeable;
+import java.time.Instant;
+import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
+import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Remote sampler that gets sampling configuration from AWS X-Ray. */
 public final class AwsXrayRemoteSampler implements Sampler, Closeable {
 
-  private static final Random RANDOM = new Random();
+  static final long DEFAULT_TARGET_INTERVAL_NANOS = TimeUnit.SECONDS.toNanos(10);
 
+  private static final Random RANDOM = new Random();
   private static final Logger logger = Logger.getLogger(AwsXrayRemoteSampler.class.getName());
 
-  private static final String WORKER_THREAD_NAME =
-      AwsXrayRemoteSampler.class.getSimpleName() + "_WorkerThread";
-
   // Unique per-process client ID, generated as a random string.
   private static final String CLIENT_ID = generateClientId();
 
   private final Resource resource;
+  private final Clock clock;
   private final Sampler initialSampler;
   private final XraySamplerClient client;
   private final ScheduledExecutorService executor;
   private final long pollingIntervalNanos;
   private final int jitterNanos;
 
   @Nullable private volatile ScheduledFuture<?> pollFuture;
+  @Nullable private volatile ScheduledFuture<?> fetchTargetsFuture;
   @Nullable private volatile GetSamplingRulesResponse previousRulesResponse;
   private volatile Sampler sampler;
 
@@ -57,8 +67,13 @@ public static AwsXrayRemoteSamplerBuilder newBuilder(Resource resource) {
   }
 
   AwsXrayRemoteSampler(
-      Resource resource, String endpoint, Sampler initialSampler, long pollingIntervalNanos) {
+      Resource resource,
+      Clock clock,
+      String endpoint,
+      Sampler initialSampler,
+      long pollingIntervalNanos) {
     this.resource = resource;
+    this.clock = clock;
     this.initialSampler = initialSampler;
     client = new XraySamplerClient(endpoint);
     executor =
@@ -107,8 +122,22 @@ private void getAndUpdateSampler() {
           client.getSamplingRules(GetSamplingRulesRequest.create(null));
       if (!response.equals(previousRulesResponse)) {
         sampler =
-            new XrayRulesSampler(CLIENT_ID, resource, initialSampler, response.getSamplingRules());
+            new XrayRulesSampler(
+                CLIENT_ID,
+                resource,
+                clock,
+                initialSampler,
+                response.getSamplingRules().stream()
+                    .map(SamplingRuleRecord::getRule)
+                    .collect(Collectors.toList()));
         previousRulesResponse = response;
+        ScheduledFuture<?> existingFetchTargetsFuture = fetchTargetsFuture;
+        if (existingFetchTargetsFuture != null) {
+          existingFetchTargetsFuture.cancel(false);
+        }
+        fetchTargetsFuture =
+            executor.schedule(
+                this::fetchTargets, DEFAULT_TARGET_INTERVAL_NANOS, TimeUnit.NANOSECONDS);
       }
     } catch (Throwable t) {
       logger.log(Level.FINE, "Failed to update sampler", t);
@@ -121,6 +150,39 @@ private void scheduleSamplerUpdate() {
     pollFuture = executor.schedule(this::getAndUpdateSampler, delay, TimeUnit.NANOSECONDS);
   }
 
+  private void fetchTargets() {
+    if (!(sampler instanceof XrayRulesSampler)) {
+      throw new IllegalStateException("Programming bug.");
+    }
+
+    XrayRulesSampler xrayRulesSampler = (XrayRulesSampler) sampler;
+    try {
+      Date now = Date.from(Instant.ofEpochSecond(0, clock.now()));
+      List<SamplingStatisticsDocument> statistics = xrayRulesSampler.snapshot(now);
+      Set<String> requestedTargetRuleNames =
+          statistics.stream()
+              .map(SamplingStatisticsDocument::getRuleName)
+              .collect(Collectors.toSet());
+
+      GetSamplingTargetsResponse response =
+          client.getSamplingTargets(GetSamplingTargetsRequest.create(statistics));
+      Map<String, SamplingTargetDocument> targets =
+          response.getDocuments().stream()
+              .collect(Collectors.toMap(SamplingTargetDocument::getRuleName, Function.identity()));
+      sampler =
+          xrayRulesSampler = xrayRulesSampler.withTargets(targets, requestedTargetRuleNames, now);
+    } catch (Throwable t) {
+      // Might be a transient API failure, try again after a default interval.
+      executor.schedule(this::fetchTargets, DEFAULT_TARGET_INTERVAL_NANOS, TimeUnit.NANOSECONDS);
+      return;
+    }
+
+    long nextTargetFetchIntervalNanos =
+        xrayRulesSampler.nextTargetFetchTimeNanos() - clock.nanoTime();
+    fetchTargetsFuture =
+        executor.schedule(this::fetchTargets, nextTargetFetchIntervalNanos, TimeUnit.NANOSECONDS);
+  }
+
   @Override
   public void close() {
     ScheduledFuture<?> pollFuture = this.pollFuture;

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

@@ -6,6 +6,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import io.opentelemetry.sdk.common.Clock;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.time.Duration;
@@ -19,6 +20,7 @@ public final class AwsXrayRemoteSamplerBuilder {
 
   private final Resource resource;
 
+  private Clock clock = Clock.getDefault();
   private String endpoint = DEFAULT_ENDPOINT;
   private Sampler initialSampler = Sampler.parentBased(Sampler.traceIdRatioBased(0.05));
   private long pollingIntervalNanos = TimeUnit.SECONDS.toNanos(DEFAULT_POLLING_INTERVAL_SECS);
@@ -70,8 +72,19 @@ public AwsXrayRemoteSamplerBuilder setInitialSampler(Sampler initialSampler) {
     return this;
   }
 
+  /**
+   * Sets the {@link Clock} used for time measurements for sampling, such as rate limiting or quota
+   * expiry.
+   */
+  public AwsXrayRemoteSamplerBuilder setClock(Clock clock) {
+    requireNonNull(clock, "clock");
+    this.clock = clock;
+    return this;
+  }
+
   /** Returns a {@link AwsXrayRemoteSampler} with the configuration of this builder. */
   public AwsXrayRemoteSampler build() {
-    return new AwsXrayRemoteSampler(resource, endpoint, initialSampler, pollingIntervalNanos);
+    return new AwsXrayRemoteSampler(
+        resource, clock, endpoint, initialSampler, pollingIntervalNanos);
   }
 }

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

@@ -35,8 +35,8 @@ abstract static class SamplingTargetDocument {
     @JsonCreator
     static SamplingTargetDocument create(
         @JsonProperty("FixedRate") double fixedRate,
-        @JsonProperty("Interval") int intervalSecs,
-        @JsonProperty("ReservoirQuota") int reservoirQuota,
+        @JsonProperty("Interval") @Nullable Integer intervalSecs,
+        @JsonProperty("ReservoirQuota") @Nullable Integer reservoirQuota,
         @JsonProperty("ReservoirQuotaTTL") @Nullable Date reservoirQuotaTtl,
         @JsonProperty("RuleName") String ruleName) {
       return new AutoValue_GetSamplingTargetsResponse_SamplingTargetDocument(
@@ -45,10 +45,14 @@ static SamplingTargetDocument create(
 
     abstract double getFixedRate();
 
-    abstract int getIntervalSecs();
+    @Nullable
+    abstract Integer getIntervalSecs();
 
-    abstract int getReservoirQuota();
+    @Nullable
+    abstract Integer getReservoirQuota();
 
+    // Careful that this is a timestamp when the quota expires, not a duration as we'd normally
+    // expect for a Time to live.
     @Nullable
     abstract Date getReservoirQuotaTtl();
 

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

@@ -19,11 +19,6 @@ final class RateLimitingSampler implements Sampler {
   private final RateLimiter limiter;
   private final int numPerSecond;
 
-  RateLimitingSampler(int numPerSecond) {
-    this(numPerSecond, Clock.getDefault());
-  }
-
-  // Visible for testing
   RateLimitingSampler(int numPerSecond, Clock clock) {
     limiter = new RateLimiter(numPerSecond, numPerSecond, clock);
     this.numPerSecond = numPerSecond;