Add AWS X-Ray Adaptive Sampling Support

aws-xray/build.gradle.kts

@@ -11,6 +11,7 @@ dependencies {
   api("io.opentelemetry:opentelemetry-sdk-trace")
 
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
+  implementation("io.opentelemetry.semconv:opentelemetry-semconv:1.32.0-alpha")
 
   implementation("com.squareup.okhttp3:okhttp")
   implementation("io.opentelemetry.semconv:opentelemetry-semconv")

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

@@ -0,0 +1,52 @@
+/*
+ * 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.TraceState;
+import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+
+final class AwsSamplingResult implements SamplingResult {
+
+  public static final String AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY = "awsxraysamplingrule";
+
+  private final SamplingDecision decision;
+  private final Attributes attributes;
+  private final String samplingRuleName;
+
+  private AwsSamplingResult(
+      SamplingDecision decision, Attributes attributes, String samplingRuleName) {
+    this.decision = decision;
+    this.attributes = attributes;
+    this.samplingRuleName = samplingRuleName;
+  }
+
+  static AwsSamplingResult create(
+      SamplingDecision decision, Attributes attributes, String samplingRuleName) {
+    return new AwsSamplingResult(decision, attributes, samplingRuleName);
+  }
+
+  @Override
+  public SamplingDecision getDecision() {
+    return decision;
+  }
+
+  @Override
+  public Attributes getAttributes() {
+    return attributes;
+  }
+
+  @Override
+  public TraceState getUpdatedTraceState(TraceState parentTraceState) {
+    if (parentTraceState.get(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY) == null) {
+      return parentTraceState.toBuilder()
+          .put(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY, samplingRuleName)
+          .build();
+    }
+    return parentTraceState;
+  }
+}

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

@@ -0,0 +1,139 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.awsxray;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import com.google.auto.value.AutoValue;
+import java.util.List;
+import javax.annotation.Nullable;
+
+@AutoValue
+@JsonSerialize(as = AwsXrayAdaptiveSamplingConfig.class)
+@JsonDeserialize(builder = AutoValue_AwsXrayAdaptiveSamplingConfig.Builder.class)
+public abstract class AwsXrayAdaptiveSamplingConfig {
+
+  @JsonProperty("version")
+  public abstract double getVersion();
+
+  @JsonProperty("anomalyConditions")
+  @Nullable
+  public abstract List<AnomalyConditions> getAnomalyConditions();
+
+  @JsonProperty("errorCaptureLimit")
+  @Nullable
+  public abstract ErrorCaptureLimit getErrorCaptureLimit();
+
+  public static Builder builder() {
+    return new AutoValue_AwsXrayAdaptiveSamplingConfig.Builder();
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+    @JsonProperty("version")
+    public abstract Builder setVersion(double value);
+
+    @JsonProperty("anomalyConditions")
+    public abstract Builder setAnomalyConditions(List<AnomalyConditions> value);
+
+    @JsonProperty("errorCaptureLimit")
+    public abstract Builder setErrorCaptureLimit(ErrorCaptureLimit value);
+
+    public abstract AwsXrayAdaptiveSamplingConfig build();
+  }
+
+  @AutoValue
+  @JsonDeserialize(
+      builder = AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyConditions.Builder.class)
+  public abstract static class AnomalyConditions {
+    @JsonProperty("errorCodeRegex")
+    @Nullable
+    public abstract String getErrorCodeRegex();
+
+    @JsonProperty("operations")
+    @Nullable
+    public abstract List<String> getOperations();
+
+    @JsonProperty("highLatencyMs")
+    @Nullable
+    public abstract Long getHighLatencyMs();
+
+    @JsonProperty("usage")
+    @Nullable
+    public abstract UsageType getUsage();
+
+    public static Builder builder() {
+      return new AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyConditions.Builder();
+    }
+
+    @AutoValue.Builder
+    public abstract static class Builder {
+      @JsonProperty("errorCodeRegex")
+      public abstract Builder setErrorCodeRegex(String value);
+
+      @JsonProperty("operations")
+      public abstract Builder setOperations(List<String> value);
+
+      @JsonProperty("highLatencyMs")
+      public abstract Builder setHighLatencyMs(Long value);
+
+      @JsonProperty("usage")
+      public abstract Builder setUsage(UsageType value);
+
+      public abstract AnomalyConditions build();
+    }
+  }
+
+  public enum UsageType {
+    BOTH("both"),
+    SAMPLING_BOOST("sampling-boost"),
+    ERROR_SPAN_CAPTURE("error-span-capture");
+
+    private final String value;
+
+    UsageType(String value) {
+      this.value = value;
+    }
+
+    @JsonValue
+    public String getValue() {
+      return value;
+    }
+
+    @JsonCreator
+    public static UsageType fromValue(String value) {
+      for (UsageType type : values()) {
+        if (type.value.equals(value)) {
+          return type;
+        }
+      }
+      throw new IllegalArgumentException("Invalid usage value: " + value);
+    }
+  }
+
+  @AutoValue
+  @JsonDeserialize(
+      builder = AutoValue_AwsXrayAdaptiveSamplingConfig_ErrorCaptureLimit.Builder.class)
+  public abstract static class ErrorCaptureLimit {
+    @JsonProperty("errorSpansPerSecond")
+    public abstract int getErrorSpansPerSecond();
+
+    public static Builder builder() {
+      return new AutoValue_AwsXrayAdaptiveSamplingConfig_ErrorCaptureLimit.Builder();
+    }
+
+    @AutoValue.Builder
+    public abstract static class Builder {
+      @JsonProperty("errorSpansPerSecond")
+      public abstract Builder setErrorSpansPerSecond(int value);
+
+      public abstract ErrorCaptureLimit build();
+    }
+  }
+}

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

@@ -9,16 +9,22 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.contrib.awsxray.GetSamplingRulesResponse.SamplingRuleRecord;
+import io.opentelemetry.contrib.awsxray.GetSamplingTargetsRequest.SamplingBoostStatisticsDocument;
 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.ReadableSpan;
 import io.opentelemetry.sdk.trace.data.LinkData;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 import java.io.Closeable;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.Date;
 import java.util.Iterator;
 import java.util.List;
@@ -43,6 +49,9 @@ public final class AwsXrayRemoteSampler implements Sampler, Closeable {
 
   private static final Logger logger = Logger.getLogger(AwsXrayRemoteSampler.class.getName());
 
+  // Default batch size to be same as OTel BSP default
+  private static final int maxExportBatchSize = 512;
+
   private final Resource resource;
   private final Clock clock;
   private final Sampler initialSampler;
@@ -59,6 +68,9 @@ public final class AwsXrayRemoteSampler implements Sampler, Closeable {
   @Nullable private volatile XrayRulesSampler internalXrayRulesSampler;
   private volatile Sampler sampler;
 
+  @Nullable private AwsXrayAdaptiveSamplingConfig adaptiveSamplingConfig;
+  @Nullable private BatchSpanProcessor bsp;
+
   /**
    * Returns a {@link AwsXrayRemoteSamplerBuilder} with the given {@link Resource}. This {@link
    * Resource} should be the same as what the OpenTelemetry SDK is configured with.
@@ -120,6 +132,40 @@ public String getDescription() {
     return "AwsXrayRemoteSampler{" + sampler.getDescription() + "}";
   }
 
+  public void setAdaptiveSamplingConfig(AwsXrayAdaptiveSamplingConfig config) {
+    if (this.adaptiveSamplingConfig != null) {
+      throw new IllegalStateException("Programming bug - Adaptive sampling config is already set");
+    } else if (config != null && this.adaptiveSamplingConfig == null) {
+      // Save here and also pass to XrayRulesSampler directly as it already exists
+      this.adaptiveSamplingConfig = config;
+      if (sampler instanceof XrayRulesSampler) {
+        ((XrayRulesSampler) sampler).setAdaptiveSamplingConfig(config);
+      }
+    }
+  }
+
+  public void setSpanExporter(SpanExporter spanExporter) {
+    if (this.bsp != null) {
+      throw new IllegalStateException("Programming bug - BatchSpanProcessor is already set");
+    } else if (spanExporter != null && this.bsp == null) {
+      this.bsp =
+          BatchSpanProcessor.builder(spanExporter)
+              .setExportUnsampledSpans(true) // Required to capture the unsampled anomaly spans
+              .setMaxExportBatchSize(maxExportBatchSize)
+              .build();
+    }
+  }
+
+  public void adaptSampling(ReadableSpan span, SpanData spanData) {
+    if (this.bsp == null) {
+      throw new IllegalStateException(
+          "Programming bug - BatchSpanProcessor is null while trying to adapt sampling");
+    }
+    if (sampler instanceof XrayRulesSampler) {
+      ((XrayRulesSampler) sampler).adaptSampling(span, spanData, this.bsp::onEnd);
+    }
+  }
+
   private void getAndUpdateSampler() {
     try {
       // No pagination support yet, or possibly ever.
@@ -134,8 +180,8 @@ private void getAndUpdateSampler() {
                 initialSampler,
                 response.getSamplingRules().stream()
                     .map(SamplingRuleRecord::getRule)
-                    .collect(Collectors.toList())));
-
+                    .collect(Collectors.toList()),
+                adaptiveSamplingConfig));
         previousRulesResponse = response;
         ScheduledFuture<?> existingFetchTargetsFuture = fetchTargetsFuture;
         if (existingFetchTargetsFuture != null) {
@@ -179,14 +225,29 @@ private void fetchTargets() {
     XrayRulesSampler xrayRulesSampler = this.internalXrayRulesSampler;
     try {
       Date now = Date.from(Instant.ofEpochSecond(0, clock.now()));
-      List<SamplingStatisticsDocument> statistics = xrayRulesSampler.snapshot(now);
+      List<SamplingRuleApplier.SamplingRuleStatisticsSnapshot> statisticsSnapshot =
+          xrayRulesSampler.snapshot(now);
+      List<SamplingStatisticsDocument> statistics = new ArrayList<SamplingStatisticsDocument>();
+      List<SamplingBoostStatisticsDocument> boostStatistics =
+          new ArrayList<SamplingBoostStatisticsDocument>();
+      statisticsSnapshot.stream()
+          .forEach(
+              snapshot -> {
+                if (snapshot.getStatisticsDocument() != null) {
+                  statistics.add(snapshot.getStatisticsDocument());
+                }
+                if (snapshot.getBoostStatisticsDocument() != null
+                    && snapshot.getBoostStatisticsDocument().getTotalCount() > 0) {
+                  boostStatistics.add(snapshot.getBoostStatisticsDocument());
+                }
+              });
       Set<String> requestedTargetRuleNames =
           statistics.stream()
               .map(SamplingStatisticsDocument::getRuleName)
               .collect(Collectors.toSet());
 
-      GetSamplingTargetsResponse response =
-          client.getSamplingTargets(GetSamplingTargetsRequest.create(statistics));
+      GetSamplingTargetsRequest req = GetSamplingTargetsRequest.create(statistics, boostStatistics);
+      GetSamplingTargetsResponse response = client.getSamplingTargets(req);
       Map<String, SamplingTargetDocument> targets =
           response.getDocuments().stream()
               .collect(Collectors.toMap(SamplingTargetDocument::getRuleName, Function.identity()));

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

@@ -62,7 +62,8 @@ static SamplingRule create(
         @JsonProperty("ServiceName") String serviceName,
         @JsonProperty("ServiceType") String serviceType,
         @JsonProperty("URLPath") String urlPath,
-        @JsonProperty("Version") int version) {
+        @JsonProperty("Version") int version,
+        @JsonProperty("SamplingRateBoost") @Nullable SamplingRateBoost samplingRateBoost) {
       return new AutoValue_GetSamplingRulesResponse_SamplingRule(
           attributes,
           fixedRate,
@@ -76,7 +77,8 @@ static SamplingRule create(
           serviceName,
           serviceType,
           urlPath,
-          version);
+          version,
+          samplingRateBoost);
     }
 
     abstract Map<String, String> getAttributes();
@@ -106,5 +108,23 @@ static SamplingRule create(
     abstract String getUrlPath();
 
     abstract int getVersion();
+
+    @Nullable
+    abstract SamplingRateBoost getSamplingRateBoost();
+  }
+
+  @AutoValue
+  abstract static class SamplingRateBoost {
+    @JsonCreator
+    static SamplingRateBoost create(
+        @JsonProperty("MaxRate") double maxRate,
+        @JsonProperty("CooldownWindowMinutes") long cooldownWindowMinutes) {
+      return new AutoValue_GetSamplingRulesResponse_SamplingRateBoost(
+          maxRate, cooldownWindowMinutes);
+    }
+
+    abstract double getMaxRate();
+
+    abstract long getCooldownWindowMinutes();
   }
 }