Introduce AWS X-Ray Adaptive Sampling support (#1170)

Author: majanjua-amzn

.github/patches/opentelemetry-java-contrib.patch

@@ -1 +1,2 @@
 OTEL_JAVA_INSTRUMENTATION_VERSION=v2.18.1
+OTEL_JAVA_CONTRIB_VERSION=v1.48.0

.github/patches/versions

@@ -38,8 +38,10 @@ dependencies {
   implementation("io.opentelemetry.contrib:opentelemetry-aws-xray")
   // AWS Resource Detectors
   implementation("io.opentelemetry.contrib:opentelemetry-aws-resources")
-  // Json file reader
+  // JSON file reader
   implementation("com.fasterxml.jackson.core:jackson-databind:2.16.1")
+  // YAML file reader
+  implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.16.1")
   // Import AWS SDK v1 core for ARN parsing utilities
   implementation("com.amazonaws:aws-java-sdk-core:1.12.773")
   // Export configuration

awsagentprovider/build.gradle.kts

@@ -112,6 +112,7 @@ public void onStart(Context parentContext, ReadWriteSpan span) {
     if (propagationData != null) {
       span.setAttribute(propagationDataKey, propagationData);
     }
+    span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, span.getSpanContext().isSampled());
   }
 
   private boolean isConsumerKind(ReadableSpan span) {

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AttributePropagatingSpanProcessor.java

@@ -26,7 +26,7 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
         () ->
             new HashMap<String, String>() {
               {
-                put("otel.propagators", "baggage,xray,tracecontext,b3,b3multi");
+                put("otel.propagators", "baggage,xray,tracecontext");
                 put("otel.instrumentation.aws-sdk.experimental-span-attributes", "true");
                 put(
                     "otel.instrumentation.aws-sdk.experimental-record-individual-http-error",

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAgentPropertiesCustomizerProvider.java

@@ -15,9 +15,14 @@
 
 package software.amazon.opentelemetry.javaagent.providers;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
-import io.opentelemetry.contrib.awsxray.AlwaysRecordSampler;
+import io.opentelemetry.contrib.awsxray.AwsXrayAdaptiveSamplingConfig;
+import io.opentelemetry.contrib.awsxray.AwsXrayRemoteSampler;
 import io.opentelemetry.contrib.awsxray.ResourceHolder;
 import io.opentelemetry.exporter.otlp.http.logs.OtlpHttpLogRecordExporter;
 import io.opentelemetry.exporter.otlp.http.metrics.OtlpHttpMetricExporter;
@@ -42,6 +47,11 @@
 import io.opentelemetry.sdk.trace.SpanProcessor;
 import io.opentelemetry.sdk.trace.export.SpanExporter;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -142,11 +152,16 @@ public final class AwsApplicationSignalsCustomizerProvider
   private static final String OTEL_EXPORTER_OTLP_LOGS_COMPRESSION_CONFIG =
       "otel.exporter.otlp.logs.compression";
 
+  private static final String AWS_XRAY_ADAPTIVE_SAMPLING_CONFIG =
+      "aws.xray.adaptive.sampling.config";
+
   // UDP packet can be upto 64KB. To limit the packet size, we limit the exported batch size.
   // This is a bit of a magic number, as there is no simple way to tell how many spans can make a
   // 64KB batch since spans can vary in size.
   private static final int LAMBDA_SPAN_EXPORT_BATCH_SIZE = 10;
 
+  private Sampler sampler;
+
   public void customize(AutoConfigurationCustomizer autoConfiguration) {
     autoConfiguration.addPropertiesCustomizer(this::customizeProperties);
     autoConfiguration.addPropertiesCustomizer(this::customizeLambdaEnvProperties);
@@ -281,6 +296,27 @@ private Resource customizeResource(Resource resource, ConfigProperties configPro
   }
 
   private Sampler customizeSampler(Sampler sampler, ConfigProperties configProps) {
+    if (sampler instanceof AwsXrayRemoteSampler) {
+      String config = configProps.getString(AWS_XRAY_ADAPTIVE_SAMPLING_CONFIG);
+      AwsXrayAdaptiveSamplingConfig parsedConfig = null;
+
+      try {
+        parsedConfig = parseConfigString(config);
+      } catch (Exception e) {
+        logger.log(
+            Level.WARNING, "Failed to parse adaptive sampling configuration: {0}", e.getMessage());
+      }
+
+      if (parsedConfig != null) {
+        try {
+          ((AwsXrayRemoteSampler) sampler).setAdaptiveSamplingConfig(parsedConfig);
+        } catch (Exception e) {
+          logger.log(
+              Level.WARNING, "Error processing adaptive sampling config: {0}", e.getMessage());
+        }
+      }
+      this.sampler = sampler;
+    }
     if (isApplicationSignalsEnabled(configProps)) {
       return AlwaysRecordSampler.create(sampler);
     }
@@ -344,10 +380,13 @@ private SdkTracerProviderBuilder customizeTracerProviderBuilder(
               .build();
 
       // Construct and set application signals metrics processor
-      SpanProcessor spanMetricsProcessor =
+      AwsSpanMetricsProcessorBuilder awsSpanMetricsProcessorBuilder =
           AwsSpanMetricsProcessorBuilder.create(
-                  meterProvider, ResourceHolder.getResource(), meterProvider::forceFlush)
-              .build();
+              meterProvider, ResourceHolder.getResource(), meterProvider::forceFlush);
+      if (this.sampler != null) {
+        awsSpanMetricsProcessorBuilder.setSampler(this.sampler);
+      }
+      SpanProcessor spanMetricsProcessor = awsSpanMetricsProcessorBuilder.build();
       tracerProviderBuilder.addSpanProcessor(spanMetricsProcessor);
     }
     return tracerProviderBuilder;
@@ -423,11 +462,14 @@ SpanExporter customizeSpanExporter(SpanExporter spanExporter, ConfigProperties c
     }
 
     if (isApplicationSignalsEnabled(configProps)) {
-      return AwsMetricAttributesSpanExporterBuilder.create(
-              spanExporter, ResourceHolder.getResource())
-          .build();
+      spanExporter =
+          AwsMetricAttributesSpanExporterBuilder.create(spanExporter, ResourceHolder.getResource())
+              .build();
     }
 
+    if (this.sampler instanceof AwsXrayRemoteSampler) {
+      ((AwsXrayRemoteSampler) this.sampler).setSpanExporter(spanExporter);
+    }
     return spanExporter;
   }
 
@@ -467,6 +509,44 @@ LogRecordExporter customizeLogsExporter(
     return logsExporter;
   }
 
+  static AwsXrayAdaptiveSamplingConfig parseConfigString(String config)
+      throws JsonProcessingException {
+    if (config == null) {
+      return null;
+    }
+
+    // Check if the config is a file path and the file exists
+    Path path = Paths.get(config);
+    if (Files.exists(path)) {
+      try {
+        config = String.join("\n", Files.readAllLines(path, StandardCharsets.UTF_8));
+      } catch (IOException e) {
+        throw new IllegalArgumentException(
+            "Failed to read adaptive sampling configuration file: " + e.getMessage(), e);
+      }
+    }
+
+    ObjectMapper yamlMapper = new ObjectMapper(new YAMLFactory());
+    Map<String, Object> configMap =
+        yamlMapper.readValue(config, new TypeReference<Map<String, Object>>() {});
+
+    Object versionObj = configMap.get("version");
+    if (versionObj == null) {
+      throw new IllegalArgumentException(
+          "Missing required 'version' field in adaptive sampling configuration");
+    }
+
+    double version = ((Number) versionObj).doubleValue();
+    if (version >= 2L) {
+      throw new IllegalArgumentException(
+          "Incompatible adaptive sampling config version: "
+              + version
+              + ". This version of the AWS X-Ray remote sampler only supports versions strictly below 2.0.");
+    }
+
+    return yamlMapper.readValue(config, AwsXrayAdaptiveSamplingConfig.class);
+  }
+
   private enum ApplicationSignalsExporterProvider {
     INSTANCE;
 

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsApplicationSignalsCustomizerProvider.java

@@ -25,12 +25,14 @@
 import io.opentelemetry.api.metrics.LongHistogram;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.contrib.awsxray.AwsXrayRemoteSampler;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
 import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.util.Map;
 import java.util.function.Supplier;
 import javax.annotation.concurrent.Immutable;
@@ -75,16 +77,25 @@ public final class AwsSpanMetricsProcessor implements SpanProcessor {
   private final Resource resource;
   private final Supplier<CompletableResultCode> forceFlushAction;
 
+  private Sampler sampler;
+
   /** Use {@link AwsSpanMetricsProcessorBuilder} to construct this processor. */
   static AwsSpanMetricsProcessor create(
       LongHistogram errorHistogram,
       LongHistogram faultHistogram,
       DoubleHistogram latencyHistogram,
       MetricAttributeGenerator generator,
       Resource resource,
+      Sampler sampler,
       Supplier<CompletableResultCode> forceFlushAction) {
     return new AwsSpanMetricsProcessor(
-        errorHistogram, faultHistogram, latencyHistogram, generator, resource, forceFlushAction);
+        errorHistogram,
+        faultHistogram,
+        latencyHistogram,
+        generator,
+        resource,
+        sampler,
+        forceFlushAction);
   }
 
   private AwsSpanMetricsProcessor(
@@ -93,12 +104,14 @@ private AwsSpanMetricsProcessor(
       DoubleHistogram latencyHistogram,
       MetricAttributeGenerator generator,
       Resource resource,
+      Sampler sampler,
       Supplier<CompletableResultCode> forceFlushAction) {
     this.errorHistogram = errorHistogram;
     this.faultHistogram = faultHistogram;
     this.latencyHistogram = latencyHistogram;
     this.generator = generator;
     this.resource = resource;
+    this.sampler = sampler;
     this.forceFlushAction = forceFlushAction;
   }
 
@@ -125,6 +138,9 @@ public void onEnd(ReadableSpan span) {
     for (Map.Entry<String, Attributes> attribute : attributeMap.entrySet()) {
       recordMetrics(span, spanData, attribute.getValue());
     }
+    if (sampler != null) {
+      ((AwsXrayRemoteSampler) sampler).adaptSampling(span, spanData);
+    }
   }
 
   @Override

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsSpanMetricsProcessor.java

@@ -24,6 +24,7 @@
 import io.opentelemetry.api.metrics.MeterProvider;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.util.function.Supplier;
 
 /** A builder for {@link AwsSpanMetricsProcessor} */
@@ -51,6 +52,7 @@ public final class AwsSpanMetricsProcessorBuilder {
 
   // Optional builder elements
   private MetricAttributeGenerator generator = DEFAULT_GENERATOR;
+  private Sampler sampler;
   private String scopeName = DEFAULT_SCOPE_NAME;
 
   public static AwsSpanMetricsProcessorBuilder create(
@@ -80,6 +82,17 @@ public AwsSpanMetricsProcessorBuilder setGenerator(MetricAttributeGenerator gene
     return this;
   }
 
+  /**
+   * Sets the sampler used to determine if the spans should be sampled This will be used to increase
+   * sampling rate in the case of errors
+   */
+  @CanIgnoreReturnValue
+  public AwsSpanMetricsProcessorBuilder setSampler(Sampler sampler) {
+    requireNonNull(sampler, "sampler");
+    this.sampler = sampler;
+    return this;
+  }
+
   /**
    * Sets the scope name used in the creation of metrics by the span metrics processor. If unset,
    * defaults to {@link #DEFAULT_SCOPE_NAME}. Must not be null.
@@ -99,6 +112,12 @@ public AwsSpanMetricsProcessor build() {
         meter.histogramBuilder(LATENCY).setUnit(LATENCY_UNITS).build();
 
     return AwsSpanMetricsProcessor.create(
-        errorHistogram, faultHistogram, latencyHistogram, generator, resource, forceFlushAction);
+        errorHistogram,
+        faultHistogram,
+        latencyHistogram,
+        generator,
+        resource,
+        sampler,
+        forceFlushAction);
   }
 }