Add custom distro instrumentation unit test example

examples/distro/agent/build.gradle

@@ -1,9 +1,15 @@
 plugins {
   id("com.github.johnrengelman.shadow") version "6.0.0"
 }
+
+apply from: "$rootDir/gradle/shadow.gradle"
+
+def relocatePackages = ext.relocatePackages
+
 configurations {
   customShadow
 }
+
 dependencies {
   customShadow project(path: ":custom", configuration: "shadow")
   customShadow project(path: ":instrumentation", configuration: "shadow")
@@ -35,17 +41,7 @@ tasks {
     }
     exclude("**/module-info.class")
 
-    // Prevents conflict with other SLF4J instances. Important for premain.
-    relocate("org.slf4j", "io.opentelemetry.javaagent.slf4j")
-    // rewrite dependencies calling Logger.getLogger
-    relocate("java.util.logging.Logger", "io.opentelemetry.javaagent.bootstrap.PatchLogger")
-
-    // prevents conflict with library instrumentation
-    relocate("io.opentelemetry.instrumentation.api", "io.opentelemetry.javaagent.shaded.instrumentation.api")
-
-    // relocate OpenTelemetry API
-    relocate("io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api")
-    relocate("io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context")
+    relocatePackages(it)
 
     manifest {
       attributes.put("Main-Class", "io.opentelemetry.javaagent.OpenTelemetryAgent")

examples/distro/custom/build.gradle

@@ -3,6 +3,10 @@ plugins {
   id("com.github.johnrengelman.shadow") version "6.0.0"
 }
 
+apply from: "$rootDir/gradle/shadow.gradle"
+
+def relocatePackages = ext.relocatePackages
+
 dependencies {
   compileOnly("io.opentelemetry:opentelemetry-sdk:${versions.opentelemetry}")
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:${versions.opentelemetryAlpha}")
@@ -15,16 +19,6 @@ tasks {
 
     exclude("**/module-info.class")
 
-    // Prevents conflict with other SLF4J instances. Important for premain.
-    relocate("org.slf4j", "io.opentelemetry.javaagent.slf4j")
-    // rewrite dependencies calling Logger.getLogger
-    relocate("java.util.logging.Logger", "io.opentelemetry.javaagent.bootstrap.PatchLogger")
-
-    // prevents conflict with library instrumentation
-    relocate("io.opentelemetry.instrumentation.api", "io.opentelemetry.javaagent.shaded.instrumentation.api")
-
-    // relocate OpenTelemetry API
-    relocate("io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api")
-    relocate("io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context")
+    relocatePackages(it)
   }
 }

examples/distro/gradle/instrumentation.gradle

@@ -0,0 +1,63 @@
+apply plugin: 'java'
+apply plugin: 'com.github.johnrengelman.shadow'
+
+apply from: "$rootDir/gradle/shadow.gradle"
+
+def relocatePackages = ext.relocatePackages
+
+configurations {
+  testInstrumentation
+  testAgent
+}
+
+dependencies {
+  compileOnly("io.opentelemetry:opentelemetry-sdk:${versions.opentelemetry}")
+  compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-api:${versions.opentelemetryJavaagentAlpha}")
+  compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:${versions.opentelemetryJavaagentAlpha}")
+
+  compileOnly deps.bytebuddy
+  compileOnly deps.bytebuddyagent
+  annotationProcessor deps.autoservice
+  compileOnly deps.autoservice
+
+  // the javaagent that is going to be used when running instrumentation unit tests
+  testAgent("io.opentelemetry.javaagent:opentelemetry-agent-for-testing:${versions.opentelemetryJavaagentAlpha}")
+  // test dependencies
+  testImplementation("io.opentelemetry.javaagent:opentelemetry-testing-common:${versions.opentelemetryJavaagentAlpha}")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-testing:${versions.opentelemetry}")
+  testImplementation("org.assertj:assertj-core:3.19.0")
+}
+
+shadowJar {
+  configurations = [project.configurations.runtimeClasspath, project.configurations.testInstrumentation]
+  mergeServiceFiles()
+
+  archiveFileName = 'agent-testing.jar'
+
+  relocatePackages(it)
+}
+
+tasks.withType(Test).configureEach {
+  jvmArgs "-Dotel.javaagent.debug=true"
+  jvmArgs "-javaagent:${configurations.testAgent.files.first().absolutePath}"
+  jvmArgs "-Dotel.initializer.jar=${shadowJar.archiveFile.get().asFile.absolutePath}"
+  jvmArgs "-Dinternal.testing.disable.global.library.ignores=true"
+  // always run with the thread propagation debugger to help track down sporadic test failures
+  jvmArgs "-Dotel.threadPropagationDebugger=true"
+  jvmArgs "-Dotel.internal.failOnContextLeak=true"
+  // always print muzzle warnings
+  jvmArgs "-Dio.opentelemetry.javaagent.slf4j.simpleLogger.log.muzzleMatcher=warn"
+  // prevent sporadic gradle deadlocks, see SafeLogger for more details
+  jvmArgs "-Dotel.internal.enableTransformSafeLogging=true"
+
+  dependsOn shadowJar
+
+  // The sources are packaged into the testing jar so we need to make sure to exclude from the test
+  // classpath, which automatically inherits them, to ensure our shaded versions are used.
+  classpath = classpath.filter {
+    if (it == file("$buildDir/resources/main") || it == file("$buildDir/classes/java/main")) {
+      return false
+    }
+    return true
+  }
+}

examples/distro/gradle/shadow.gradle

@@ -0,0 +1,21 @@
+ext.relocatePackages = { shadowJar ->
+  // Prevents conflict with other SLF4J instances. Important for premain.
+  shadowJar.relocate 'org.slf4j', 'io.opentelemetry.javaagent.slf4j'
+  // rewrite dependencies calling Logger.getLogger
+  shadowJar.relocate 'java.util.logging.Logger', 'io.opentelemetry.javaagent.bootstrap.PatchLogger'
+
+  // rewrite library instrumentation dependencies
+  shadowJar.relocate "io.opentelemetry.instrumentation", "io.opentelemetry.javaagent.shaded.instrumentation"
+
+  // relocate OpenTelemetry API usage
+  shadowJar.relocate "io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api"
+  shadowJar.relocate "io.opentelemetry.semconv", "io.opentelemetry.javaagent.shaded.io.opentelemetry.semconv"
+  shadowJar.relocate "io.opentelemetry.spi", "io.opentelemetry.javaagent.shaded.io.opentelemetry.spi"
+  shadowJar.relocate "io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context"
+
+  // relocate the OpenTelemetry extensions that are used by instrumentation modules
+  // these extensions live in the AgentClassLoader, and are injected into the user's class loader
+  // by the instrumentation modules that use them
+  shadowJar.relocate "io.opentelemetry.extension.aws", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.aws"
+  shadowJar.relocate "io.opentelemetry.extension.kotlin", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.kotlin"
+}

examples/distro/instrumentation/build.gradle

@@ -2,24 +2,17 @@ plugins {
   id("com.github.johnrengelman.shadow") version "6.0.0"
 }
 
+apply from: "$rootDir/gradle/shadow.gradle"
+
+def relocatePackages = ext.relocatePackages
+
 Project instr_project = project
 subprojects {
   afterEvaluate { Project subProj ->
     if (subProj.getPlugins().hasPlugin('java')) {
       // Make it so all instrumentation subproject tests can be run with a single command.
       instr_project.tasks.test.dependsOn(subProj.tasks.test)
 
-      dependencies {
-        compileOnly("io.opentelemetry:opentelemetry-sdk:${versions.opentelemetry}")
-        compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-api:${versions.opentelemetryJavaagentAlpha}")
-        compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:${versions.opentelemetryJavaagentAlpha}")
-
-        compileOnly deps.bytebuddy
-        compileOnly deps.bytebuddyagent
-        annotationProcessor deps.autoservice
-        compileOnly deps.autoservice
-      }
-
       instr_project.dependencies {
         implementation(project(subProj.getPath()))
       }
@@ -28,35 +21,11 @@ subprojects {
 }
 
 shadowJar {
-
   mergeServiceFiles()
 
   exclude '**/module-info.class'
 
-  // Prevents conflict with other SLF4J instances. Important for premain.
-  relocate 'org.slf4j', 'io.opentelemetry.javaagent.slf4j'
-
   duplicatesStrategy = DuplicatesStrategy.FAIL
 
-  // rewrite library instrumentation dependencies
-  relocate("io.opentelemetry.instrumentation", "io.opentelemetry.javaagent.shaded.instrumentation") {
-    exclude "io.opentelemetry.javaagent.instrumentation.**"
-  }
-
-  // rewrite dependencies calling Logger.getLogger
-  relocate 'java.util.logging.Logger', 'io.opentelemetry.javaagent.bootstrap.PatchLogger'
-
-  // prevents conflict with library instrumentation
-  relocate 'io.opentelemetry.instrumentation.api', 'io.opentelemetry.javaagent.shaded.instrumentation.api'
-
-  // relocate OpenTelemetry API usage
-  relocate "io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api"
-  relocate "io.opentelemetry.spi", "io.opentelemetry.javaagent.shaded.io.opentelemetry.spi"
-  relocate "io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context"
-
-  // relocate the OpenTelemetry extensions that are used by instrumentation modules
-  // these extensions live in the AgentClassLoader, and are injected into the user's class loader
-  // by the instrumentation modules that use them
-  relocate "io.opentelemetry.extension.aws", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.aws"
-  relocate "io.opentelemetry.extension.kotlin", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.kotlin"
+  relocatePackages(it)
 }

examples/distro/instrumentation/servlet-3/build.gradle

@@ -1,7 +1,17 @@
-plugins {
-  id "java"
-}
+apply from: "$rootDir/gradle/instrumentation.gradle"
 
 dependencies {
   compileOnly group: 'javax.servlet', name: 'javax.servlet-api', version: '3.0.1'
+
+  testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-servlet-common', version: versions.opentelemetryJavaagentAlpha
+  testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-servlet-2.2', version: versions.opentelemetryJavaagentAlpha
+  testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-servlet-3.0', version: versions.opentelemetryJavaagentAlpha
+
+  testImplementation("io.opentelemetry.javaagent:opentelemetry-testing-common:${versions.opentelemetryJavaagentAlpha}") {
+    exclude group: 'org.eclipse.jetty', module: 'jetty-server'
+  }
+
+  testImplementation group: 'javax.servlet', name: 'javax.servlet-api', version: '3.0.1'
+  testImplementation group: 'org.eclipse.jetty', name: 'jetty-server', version: '8.0.0.v20110901'
+  testImplementation group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '8.0.0.v20110901'
 }

examples/distro/instrumentation/servlet-3/src/test/java/com/example/javaagent/instrumentation/DemoServlet3InstrumentationTest.java

@@ -0,0 +1,95 @@
+package com.example.javaagent.instrumentation;
+
+import static io.opentelemetry.sdk.testing.assertj.TracesAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.instrumentation.test.utils.OkHttpUtils;
+import io.opentelemetry.instrumentation.test.utils.PortUtils;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import java.io.IOException;
+import java.io.Writer;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import okhttp3.HttpUrl;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.servlet.DefaultServlet;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * This is a demo instrumentation test that verifies that the custom servlet instrumentation was applied.
+ */
+class DemoServlet3InstrumentationTest {
+  @RegisterExtension
+  static final AgentInstrumentationExtension instrumentation = AgentInstrumentationExtension
+      .create();
+
+  static final OkHttpClient httpClient = OkHttpUtils.client();
+
+  static int port;
+  static Server server;
+
+  @BeforeAll
+  static void startServer() throws Exception {
+    port = PortUtils.randomOpenPort();
+    server = new Server(port);
+    for (var connector : server.getConnectors()) {
+      connector.setHost("localhost");
+    }
+
+    var servletContext = new ServletContextHandler(null, null);
+    servletContext.addServlet(DefaultServlet.class, "/");
+    servletContext.addServlet(TestServlet.class, "/servlet");
+    server.setHandler(servletContext);
+
+    server.start();
+  }
+
+  @AfterAll
+  static void stopServer() throws Exception {
+    server.stop();
+    server.destroy();
+  }
+
+  @Test
+  void shouldAddCustomHeader() throws Exception {
+    // given
+    var request =
+        new Request.Builder()
+            .url(HttpUrl.get("http://localhost:" + port + "/servlet"))
+            .get()
+            .build();
+
+    // when
+    var response = httpClient.newCall(request).execute();
+
+    // then
+    assertEquals(200, response.code());
+    assertEquals("result", response.body().string());
+
+    assertThat(instrumentation.waitForTraces(1))
+        .hasSize(1)
+        .hasTracesSatisfyingExactly(trace -> trace.hasSize(1)
+            .hasSpansSatisfyingExactly(span -> span.hasName("/servlet").hasKind(SpanKind.SERVER)));
+
+    var traceId = instrumentation.spans().get(0).getTraceId();
+    assertEquals(traceId, response.header("X-server-id"));
+  }
+
+  public static class TestServlet extends HttpServlet {
+    protected void doGet(HttpServletRequest request, HttpServletResponse response)
+        throws IOException {
+      try (Writer writer = response.getWriter()) {
+        writer.write("result");
+        response.setStatus(200);
+      }
+    }
+  }
+}