Wire up xray sampler. (#72)

* Wire up xray sampler.

* Comment

Author: anuraaga

awsagentprovider/build.gradle.kts

@@ -28,6 +28,7 @@ dependencies {
   compileOnly("org.slf4j:slf4j-api")
 
   implementation("io.opentelemetry:opentelemetry-sdk-extension-aws")
+  implementation("io.opentelemetry.contrib:opentelemetry-aws-xray")
 
   testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
 

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

@@ -15,8 +15,8 @@
 
 package software.amazon.opentelemetry.javaagent.providers;
 
+import io.opentelemetry.contrib.awsxray.AwsXrayIdGenerator;
 import io.opentelemetry.sdk.autoconfigure.spi.SdkTracerProviderConfigurer;
-import io.opentelemetry.sdk.extension.aws.trace.AwsXrayIdGenerator;
 import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
 
 public class AwsTracerConfigurer implements SdkTracerProviderConfigurer {

dependencyManagement/build.gradle.kts

@@ -66,6 +66,7 @@ val DEPENDENCIES = listOf(
   "commons-logging:commons-logging:1.2",
   "com.sparkjava:spark-core:2.9.3",
   "com.squareup.okhttp3:okhttp:4.9.1",
+  "io.opentelemetry.contrib:opentelemetry-aws-xray:1.4.0",
   "io.opentelemetry.javaagent:opentelemetry-javaagent:${if (!TEST_SNAPSHOTS) "1.4.0" else "1.5.0-SNAPSHOT"}",
   "net.bytebuddy:byte-buddy:1.11.0"
 )

smoke-tests/runner/build.gradle.kts

@@ -36,7 +36,7 @@ project.evaluationDependsOn(":otelagent")
 
 val otelAgentJarTask = project(":otelagent").tasks.named<Jar>("shadowJar")
 tasks {
-  named<Test>("test") {
+  withType<Test>().configureEach {
     dependsOn(otelAgentJarTask)
 
     jvmArgs(

smoke-tests/runner/src/test/java/io/awsobservability/instrumentation/smoketests/runner/SpringBootSmokeTest.java

@@ -119,6 +119,21 @@ public ExportTraceServiceRequest deserialize(
           .withEnv("OTEL_BSP_SCHEDULE_DELAY", "10")
           .withEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://backend:8080");
 
+  @Container
+  private static final GenericContainer<?> applicationXraySampler =
+      new GenericContainer<>("public.ecr.aws/u0d6r4y4/aws-otel-java-smoketests-springboot:latest")
+          .dependsOn(backend)
+          .withExposedPorts(8080)
+          .withNetwork(network)
+          .withLogConsumer(new Slf4jLogConsumer(applicationLogger))
+          .withCopyFileToContainer(
+              MountableFile.forHostPath(AGENT_PATH), "/opentelemetry-javaagent-all.jar")
+          .withEnv("JAVA_TOOL_OPTIONS", "-javaagent:/opentelemetry-javaagent-all.jar")
+          .withEnv("OTEL_BSP_MAX_EXPORT_BATCH", "1")
+          .withEnv("OTEL_BSP_SCHEDULE_DELAY", "10")
+          .withEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://backend:8080")
+          .withEnv("OTEL_TRACES_SAMPLER", "xray");
+
   private static final TypeReference<List<ExportTraceServiceRequest>>
       EXPORT_TRACE_SERVICE_REQUEST_LIST = new TypeReference<>() {};
 
@@ -186,6 +201,35 @@ void hello() {
             });
   }
 
+  @Test
+  void defaultSamplingRate() {
+    int numRequests = 100;
+    for (int i = 0; i < numRequests; i++) {
+      appClient.get("/hello").aggregate().join();
+    }
+
+    var exported = getExported();
+    // 5 spans per request (1 CLIENT, 2 SERVER, 2 INTERNAL)
+    assertThat(exported).hasSize(numRequests * 5);
+  }
+
+  @Test
+  void remoteSamplingRate() {
+    // We just want to make sure OTEL_TRACES_SAMPLER=xray is using the XRay sampler. It will fail to
+    // fetch rules and fallback to the default, which will not sample all requests.
+    WebClient client =
+        WebClient.of("http://localhost:" + applicationXraySampler.getMappedPort(8080));
+    int numRequests = 100;
+    for (int i = 0; i < numRequests; i++) {
+      client.get("/hello").aggregate().join();
+    }
+
+    var exported = getExported();
+    // 5 spans per request (1 CLIENT, 2 SERVER, 2 INTERNAL) if sampled. The default sampler is
+    // around 5%, we do a very rough check that there are no more than 10% sampled.
+    assertThat(exported).hasSizeLessThanOrEqualTo(numRequests * 5 / 10);
+  }
+
   private List<Span> getExported() {
     List<ExportTraceServiceRequest> exported = ImmutableList.of();
     for (int i = 0; i < 100; i++) {