Improve virtual thread context tracking support

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java

@@ -0,0 +1,132 @@
+package datadog.trace.agent.test;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.instrument.classinject.ClassInjector;
+import datadog.trace.agent.tooling.AgentInstaller;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
+import datadog.trace.api.Config;
+import datadog.trace.api.IdGenerationStrategy;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.common.writer.ListWriter;
+import datadog.trace.core.CoreTracer;
+import datadog.trace.core.DDSpan;
+import datadog.trace.core.PendingTrace;
+import datadog.trace.core.TraceCollector;
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.Instrumentation;
+import java.util.ServiceLoader;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@ExtendWith(TestClassShadowingExtension.class)
+public abstract class AbstractInstrumentationTest {
+  static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
+
+  static final long TIMEOUT_MILLIS = TimeUnit.SECONDS.toMillis(20);
+
+  protected AgentTracer.TracerAPI tracer;
+
+  protected ListWriter writer;
+
+  protected ClassFileTransformer activeTransformer;
+  protected ClassFileTransformerListener transformerLister;
+
+  @BeforeEach
+  public void init() {
+    // If this fails, it's likely the result of another test loading Config before it can be
+    // injected into the bootstrap classpath.
+    // If one test extends AgentTestRunner in a module, all tests must extend
+    assertNull(Config.class.getClassLoader(), "Config must load on the bootstrap classpath.");
+
+    // Initialize test tracer
+    this.writer = new ListWriter();
+    // Initialize test tracer
+    CoreTracer tracer =
+        CoreTracer.builder()
+            .writer(this.writer)
+            .idGenerationStrategy(IdGenerationStrategy.fromName(idGenerationStrategyName()))
+            .strictTraceWrites(useStrictTraceWrites())
+            .build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+    this.tracer = tracer;
+
+    ClassInjector.enableClassInjection(INSTRUMENTATION);
+
+    // if a test enables the instrumentation it verifies,
+    // the cache needs to be recomputed taking into account that instrumentation's matchers
+    ClassLoaderMatchers.resetState();
+
+    assertTrue(
+        ServiceLoader.load(
+                InstrumenterModule.class, AbstractInstrumentationTest.class.getClassLoader())
+            .iterator()
+            .hasNext(),
+        "No instrumentation found");
+    this.transformerLister = new ClassFileTransformerListener();
+    this.activeTransformer =
+        AgentInstaller.installBytebuddyAgent(
+            INSTRUMENTATION, true, AgentInstaller.getEnabledSystems(), this.transformerLister);
+  }
+
+  protected String idGenerationStrategyName() {
+    return "SEQUENTIAL";
+  }
+
+  private boolean useStrictTraceWrites() {
+    return true;
+  }
+
+  @AfterEach
+  public void tearDown() {
+    this.tracer.close();
+    this.writer.close();
+    if (this.activeTransformer != null) {
+      INSTRUMENTATION.removeTransformer(this.activeTransformer);
+      this.activeTransformer = null;
+    }
+
+    // All cleanups should happen before these assertions.
+    // If not, a failing assertion may prevent cleanup
+    this.transformerLister.verify();
+    this.transformerLister = null;
+  }
+
+  protected void blockUntilChildSpansFinished(final int numberOfSpans) {
+    blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
+  }
+
+  static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
+    if (span instanceof DDSpan) {
+      TraceCollector traceCollector = ((DDSpan) span).context().getTraceCollector();
+      if (!(traceCollector instanceof PendingTrace)) {
+        throw new IllegalStateException(
+            "Expected $PendingTrace.name trace collector, got $traceCollector.class.name");
+      }
+
+      PendingTrace pendingTrace = (PendingTrace) traceCollector;
+      long deadline = System.currentTimeMillis() + TIMEOUT_MILLIS;
+
+      while (pendingTrace.size() < numberOfSpans) {
+        if (System.currentTimeMillis() > deadline) {
+          throw new RuntimeException(
+              new TimeoutException(
+                  "Timed out waiting for child spans. Received: " + pendingTrace.size()));
+        }
+        try {
+          Thread.sleep(10);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/ClassFileTransformerListener.java

@@ -0,0 +1,83 @@
+package datadog.trace.agent.test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.google.common.collect.Sets;
+import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.utility.JavaModule;
+import net.bytebuddy.utility.nullability.MaybeNull;
+
+public class ClassFileTransformerListener implements AgentBuilder.Listener {
+
+  final Set<String> transformedClassesNames = Sets.newConcurrentHashSet();
+  final Set<TypeDescription> transformedClassesTypes = Sets.newConcurrentHashSet();
+  final AtomicInteger instrumentationErrorCount = new AtomicInteger(0);
+
+  @Override
+  public void onTransformation(
+      TypeDescription typeDescription,
+      @MaybeNull ClassLoader classLoader,
+      @MaybeNull JavaModule module,
+      boolean loaded,
+      DynamicType dynamicType) {
+    this.transformedClassesNames.add(typeDescription.getActualName());
+    this.transformedClassesTypes.add(typeDescription);
+  }
+
+  @SuppressForbidden // Allows System.out.println
+  @Override
+  public void onError(
+      String typeName,
+      ClassLoader classLoader,
+      JavaModule module,
+      boolean loaded,
+      Throwable throwable) {
+    // Incorrect* classes assert on incorrect api usage. Error expected.
+    if (typeName.startsWith("context.FieldInjectionTestInstrumentation$Incorrect")
+        && throwable.getMessage().startsWith("Incorrect Context Api Usage detected.")) {
+      return;
+    }
+
+    System.out.println(
+        "Unexpected instrumentation error when instrumenting " + typeName + " on " + classLoader);
+    throwable.printStackTrace();
+    instrumentationErrorCount.incrementAndGet();
+  }
+
+  @Override
+  public void onDiscovery(
+      String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
+    // Nothing special to do
+  }
+
+  @Override
+  public void onIgnored(
+      TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, boolean loaded) {
+    // Nothing special to do
+  }
+
+  @Override
+  public void onComplete(
+      String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
+    // Nothing special to do
+  }
+
+  public void verify() {
+    // Check instrumentation errors
+    int errorCount = this.instrumentationErrorCount.get();
+    assertEquals(0, errorCount, errorCount + " instrumentation errors during test");
+    // Check effectively transformed classes that should have been ignored
+    assertTrue(
+        this.transformedClassesTypes.stream()
+            .map(TypeDescription::getActualName)
+            .noneMatch(GlobalIgnores::isAdditionallyIgnored),
+        "Transformed classes match global libraries ignore matcher");
+  }
+}

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle

@@ -22,9 +22,9 @@ idea {
   }
 }
 
-// Set all compile tasks to use JDK21 but let instrumentation code targets 1.8 compatibility
-tasks.withType(AbstractCompile).configureEach {
-  configureCompiler(it, 21, JavaVersion.VERSION_1_8)
+// Set test compile task to use JDK 21 to use the virtual threads API
+tasks.named("compileTestJava", JavaCompile) {
+  configureCompiler(it, 21)
 }
 
 dependencies {