First try rate limited sampler for initialSampler of XRay remote sampler. (#68)

Author: Anuraag Agrawal

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

@@ -11,6 +11,7 @@
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.time.Duration;
 import java.util.concurrent.TimeUnit;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** A builder for {@link AwsXrayRemoteSampler}. */
 public final class AwsXrayRemoteSamplerBuilder {
@@ -22,7 +23,7 @@ public final class AwsXrayRemoteSamplerBuilder {
 
   private Clock clock = Clock.getDefault();
   private String endpoint = DEFAULT_ENDPOINT;
-  private Sampler initialSampler = Sampler.parentBased(Sampler.traceIdRatioBased(0.05));
+  @Nullable private Sampler initialSampler;
   private long pollingIntervalNanos = TimeUnit.SECONDS.toNanos(DEFAULT_POLLING_INTERVAL_SECS);
 
   AwsXrayRemoteSamplerBuilder(Resource resource) {
@@ -84,6 +85,13 @@ public AwsXrayRemoteSamplerBuilder setClock(Clock clock) {
 
   /** Returns a {@link AwsXrayRemoteSampler} with the configuration of this builder. */
   public AwsXrayRemoteSampler build() {
+    Sampler initialSampler = this.initialSampler;
+    if (initialSampler == null) {
+      initialSampler =
+          Sampler.parentBased(
+              new OrElseSampler(
+                  new RateLimitingSampler(1, clock), Sampler.traceIdRatioBased(0.05)));
+    }
     return new AwsXrayRemoteSampler(
         resource, clock, endpoint, initialSampler, pollingIntervalNanos);
   }

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.awsxray;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
+
+class OrElseSampler implements Sampler {
+
+  private final Sampler first;
+  private final Sampler second;
+
+  OrElseSampler(Sampler first, Sampler second) {
+    this.first = first;
+    this.second = second;
+  }
+
+  @Override
+  public SamplingResult shouldSample(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    SamplingResult result =
+        first.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+    if (result.getDecision() != SamplingDecision.DROP) {
+      return result;
+    }
+    return second.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
+  }
+
+  @Override
+  public String getDescription() {
+    return "OrElse{"
+        + "first:"
+        + first.getDescription()
+        + ", second:"
+        + second.getDescription()
+        + "}";
+  }
+}

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

@@ -236,7 +236,12 @@ void getAndUpdate() throws Exception {
 
   @Test
   void defaultInitialSampler() {
-    assertThat(AwsXrayRemoteSampler.newBuilder(Resource.empty()).build().getDescription())
-        .startsWith("AwsXrayRemoteSampler{ParentBased{root:TraceIdRatioBased{0.050000}");
+    try (AwsXrayRemoteSampler sampler = AwsXrayRemoteSampler.newBuilder(Resource.empty()).build()) {
+      assertThat(sampler.getDescription())
+          .startsWith(
+              "AwsXrayRemoteSampler{"
+                  + "ParentBased{root:OrElse{"
+                  + "first:RateLimitingSampler{1}, second:TraceIdRatioBased{0.050000}");
+    }
   }
 }

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package io.opentelemetry.contrib.awsxray;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceId;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.Collections;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@ExtendWith(MockitoExtension.class)
+class OrElseSamplerTest {
+
+  @Test
+  void firstWins() {
+    Sampler sampler = new OrElseSampler(Sampler.alwaysOn(), Sampler.alwaysOff());
+    assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+  }
+
+  @Test
+  void fallsBackToSecond() {
+    Sampler sampler = new OrElseSampler(Sampler.alwaysOff(), Sampler.alwaysOn());
+    assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.RECORD_AND_SAMPLE);
+    sampler = new OrElseSampler(Sampler.alwaysOff(), Sampler.alwaysOff());
+    assertThat(doSample(sampler).getDecision()).isEqualTo(SamplingDecision.DROP);
+  }
+
+  private SamplingResult doSample(Sampler sampler) {
+    return sampler.shouldSample(
+        Context.current(),
+        TraceId.fromLongs(1, 2),
+        "span",
+        SpanKind.CLIENT,
+        Attributes.empty(),
+        Collections.emptyList());
+  }
+}