Add AWS X-Ray Adaptive Sampling Support

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

@@ -10,6 +10,7 @@
 
 import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.baggage.BaggageBuilder;
+import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
 import io.opentelemetry.api.internal.StringUtils;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
@@ -80,6 +81,9 @@ public final class AwsXrayPropagator implements TextMapPropagator {
   private static final String INVALID_LINEAGE = "-1:11111111:0";
   private static final int NUM_OF_LINEAGE_DELIMITERS = 2;
 
+  // Copied from AwsSamplingResult in aws-xray extension
+  private static final String AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY = "xrsr";
+
   private static final List<String> FIELDS = singletonList(TRACE_HEADER_KEY);
 
   private static final AwsXrayPropagator INSTANCE = new AwsXrayPropagator();
@@ -140,6 +144,16 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
 
     Baggage baggage = Baggage.fromContext(context);
     String lineageHeader = baggage.getEntryValue(LINEAGE_KEY);
+    // Get sampling rule from trace state and inject into baggage
+    // This is a back up in case the next service does not have trace state propagation
+    String ruleFromTraceState =
+        spanContext.getTraceState().get(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY);
+    if (ruleFromTraceState != null) {
+      baggage =
+          baggage.toBuilder()
+              .put(AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY, ruleFromTraceState)
+              .build();
+    }
 
     if (lineageHeader != null) {
       traceHeader
@@ -152,6 +166,9 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
     // add 256 character truncation
     String truncatedTraceHeader = traceHeader.substring(0, Math.min(traceHeader.length(), 256));
     setter.set(carrier, TRACE_HEADER_KEY, truncatedTraceHeader);
+
+    // Ensure baggage is propagated with any modifications
+    W3CBaggagePropagator.getInstance().inject(context.with(baggage), carrier, setter);
   }
 
   @Override
@@ -245,12 +262,15 @@ private static <C> Context getContextFromHeader(
       logger.finest("Both traceId and spanId are required to extract a valid span context. ");
     }
 
+    SpanContext upstreamSpanContext = Span.fromContext(context).getSpanContext();
     SpanContext spanContext =
         SpanContext.createFromRemoteParent(
             StringUtils.padLeft(traceId, TraceId.getLength()),
             spanId,
             isSampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
-            TraceState.getDefault());
+            upstreamSpanContext.isValid()
+                ? upstreamSpanContext.getTraceState()
+                : TraceState.getDefault());
 
     if (spanContext.isValid()) {
       context = context.with(Span.wrap(spanContext));

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")
@@ -25,6 +26,7 @@ dependencies {
 
   implementation("com.fasterxml.jackson.core:jackson-core")
   implementation("com.fasterxml.jackson.core:jackson-databind")
+  implementation("com.github.ben-manes.caffeine:caffeine:2.9.3")
 
   testImplementation("com.linecorp.armeria:armeria-junit5")
   testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")

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

@@ -0,0 +1,56 @@
+/*
+ * 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;
+import javax.annotation.Nullable;
+
+final class AwsSamplingResult implements SamplingResult {
+
+  // OTel trace state is a space shared with other vendors with a 256 character limit
+  // We keep the key and values as short as possible while still identifiable
+  public static final String AWS_XRAY_SAMPLING_RULE_TRACE_STATE_KEY = "xrsr";
+
+  private final SamplingDecision decision;
+  private final Attributes attributes;
+  @Nullable private final String samplingRuleName;
+
+  private AwsSamplingResult(
+      SamplingDecision decision, Attributes attributes, @Nullable String samplingRuleName) {
+    this.decision = decision;
+    this.attributes = attributes;
+    this.samplingRuleName = samplingRuleName;
+  }
+
+  static AwsSamplingResult create(
+      SamplingDecision decision, Attributes attributes, @Nullable 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
+        && this.samplingRuleName != 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,148 @@
+/*
+ * 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("anomalyCaptureLimit")
+  @Nullable
+  public abstract AnomalyCaptureLimit getAnomalyCaptureLimit();
+
+  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("anomalyCaptureLimit")
+    public abstract Builder setAnomalyCaptureLimit(AnomalyCaptureLimit 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"),
+    ANOMALY_TRACE_CAPTURE("anomaly-trace-capture"),
+    NEITHER("neither"); // Not meant to be used by customers
+
+    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);
+    }
+
+    public static boolean isUsedForBoost(UsageType usage) {
+      return BOTH.equals(usage) || SAMPLING_BOOST.equals(usage);
+    }
+
+    public static boolean isUsedForAnomalyTraceCapture(UsageType usage) {
+      return BOTH.equals(usage) || ANOMALY_TRACE_CAPTURE.equals(usage);
+    }
+  }
+
+  @AutoValue
+  @JsonDeserialize(
+      builder = AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyCaptureLimit.Builder.class)
+  public abstract static class AnomalyCaptureLimit {
+    @JsonProperty("anomalyTracesPerSecond")
+    public abstract int getAnomalyTracesPerSecond();
+
+    public static Builder builder() {
+      return new AutoValue_AwsXrayAdaptiveSamplingConfig_AnomalyCaptureLimit.Builder();
+    }
+
+    @AutoValue.Builder
+    public abstract static class Builder {
+      @JsonProperty("anomalyTracesPerSecond")
+      public abstract Builder setAnomalyTracesPerSecond(int value);
+
+      public abstract AnomalyCaptureLimit build();
+    }
+  }
+}