Support JDK-21 virtual thread executor

.github/workflows/codeql-analysis.yml

@@ -33,7 +33,7 @@ jobs:
           # queries: ./path/to/local/query, your-org/your-repo/queries@main
 
       - name: Build dd-trace-java for creating the CodeQL database
-        run: JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 ./gradlew clean :dd-java-agent:shadowJar --build-cache --parallel --stacktrace --no-daemon --max-workers=8
+        run: JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 JAVA_21_HOME=$JAVA_HOME_21_X64 ./gradlew clean :dd-java-agent:shadowJar --build-cache --parallel --stacktrace --no-daemon --max-workers=8
 
       - name: Perform CodeQL Analysis
         uses: github/codeql-action/analyze@1a927e9307bc11970b2c679922ebc4d03a5bd980 # 1.0.31

.github/workflows/trivy-analysis.yml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Build and publish artifacts locally
         run: |
-          GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx2G -Xms2G'" JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 ./gradlew clean publishToMavenLocal --build-cache --parallel --stacktrace --no-daemon --max-workers=4
+          GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx2G -Xms2G'" JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 JAVA_21_HOME=$JAVA_HOME_21_X64 ./gradlew clean publishToMavenLocal --build-cache --parallel --stacktrace --no-daemon --max-workers=4
 
       - name: Copy published artifacts
         run: |

dd-java-agent/instrumentation/java-concurrent/build.gradle

@@ -1,3 +1,7 @@
+ext {
+  latestDepTestMinJavaVersionForTests = JavaVersion.VERSION_21
+}
+
 muzzle {
   pass {
     coreJdk()
@@ -6,11 +10,17 @@ muzzle {
 
 apply from: "$rootDir/gradle/java.gradle"
 
+addTestSuite('latestDepTest')
+
+compileLatestDepTestGroovy.configure {
+  javaLauncher = getJavaLauncherFor(21)
+}
 dependencies {
   testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
 
   // test dependencies required for testing the executors we permit
   testImplementation 'org.apache.tomcat.embed:tomcat-embed-core:7.0.0'
   testImplementation deps.guava
   testImplementation group: 'io.netty', name: 'netty-all', version: '4.1.9.Final'
+  latestDepTestImplementation group: 'io.netty', name: 'netty-all', version: '4.+'
 }

dd-java-agent/instrumentation/java-concurrent/src/latestDepTest/groovy/VirtualThreadTest.groovy

@@ -0,0 +1,79 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.Trace
+import datadog.trace.core.DDSpan
+import spock.lang.Shared
+
+import java.util.concurrent.Callable
+import java.util.concurrent.ExecutorCompletionService
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
+
+class VirtualThreadTest extends AgentTestRunner {
+
+  @Shared
+  def executeRunnable = { e, c -> e.execute((Runnable) c) }
+  @Shared
+  def submitRunnable = { e, c -> e.submit((Runnable) c) }
+  @Shared
+  def submitCallable = { e, c -> e.submit((Callable) c) }
+  @Shared
+  def submitRunnableExecutorCompletionService = { ecs, c -> ecs.submit((Runnable) c, null) }
+  @Shared
+  def invokeAll = { e, c -> e.invokeAll([(Callable) c]) }
+  @Shared
+  def invokeAllTimeout = { e, c -> e.invokeAll([(Callable) c], 10, TimeUnit.SECONDS) }
+  @Shared
+  def invokeAny = { e, c -> e.invokeAny([(Callable) c]) }
+  @Shared
+  def invokeAnyTimeout = { e, c -> e.invokeAny([(Callable) c], 10, TimeUnit.SECONDS) }
+
+  def "virtualThreadPool #name"() {
+    setup:
+    def pool = poolImpl
+    def m = method
+
+    new Runnable() {
+        @Override
+        @Trace(operationName = "parent")
+        void run() {
+          activeScope().setAsyncPropagation(true)
+          // this child will have a span
+          m(pool, new JavaAsyncChild())
+          // this child won't
+          m(pool, new JavaAsyncChild(false, false))
+          blockUntilChildSpansFinished(1)
+        }
+      }.run()
+
+    TEST_WRITER.waitForTraces(1)
+    List<DDSpan> trace = TEST_WRITER.get(0)
+
+    expect:
+    TEST_WRITER.size() == 1
+    trace.size() == 2
+    trace.get(0).operationName == "parent"
+    trace.get(1).operationName == "asyncChild"
+    trace.get(1).parentId == trace.get(0).spanId
+
+    cleanup:
+    if (pool?.hasProperty("shutdown")) {
+      pool?.shutdown()
+    }
+
+    where:
+    // spotless:off
+    name                     | method              | poolImpl
+    "execute Runnable"       | executeRunnable     | Executors.newVirtualThreadPerTaskExecutor()
+    "submit Runnable"        | submitRunnable      | Executors.newVirtualThreadPerTaskExecutor()
+    "submit Callable"        | submitCallable      | Executors.newVirtualThreadPerTaskExecutor()
+    "submit Runnable ECS"    | submitRunnableExecutorCompletionService | new ExecutorCompletionService<>(Executors.newVirtualThreadPerTaskExecutor())
+    "submit Callable ECS"    | submitCallable      | new ExecutorCompletionService<>(Executors.newVirtualThreadPerTaskExecutor())
+    "invokeAll"              | invokeAll           | Executors.newVirtualThreadPerTaskExecutor()
+    "invokeAll with timeout" | invokeAllTimeout    | Executors.newVirtualThreadPerTaskExecutor()
+    "invokeAny"              | invokeAny           | Executors.newVirtualThreadPerTaskExecutor()
+    "invokeAny with timeout" | invokeAnyTimeout    | Executors.newVirtualThreadPerTaskExecutor()
+     // spotless:on
+  }
+}

dd-java-agent/instrumentation/java-concurrent/src/latestDepTest/java/JavaAsyncChild.java

@@ -0,0 +1,59 @@
+import datadog.trace.api.Trace;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ForkJoinTask;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class JavaAsyncChild extends ForkJoinTask implements Runnable, Callable {
+  private final AtomicBoolean blockThread;
+  private final boolean doTraceableWork;
+
+  public JavaAsyncChild() {
+    this(true, false);
+  }
+
+  @Override
+  public Object getRawResult() {
+    return null;
+  }
+
+  @Override
+  protected void setRawResult(final Object value) {}
+
+  @Override
+  protected boolean exec() {
+    runImpl();
+    return true;
+  }
+
+  public JavaAsyncChild(final boolean doTraceableWork, final boolean blockThread) {
+    this.doTraceableWork = doTraceableWork;
+    this.blockThread = new AtomicBoolean(blockThread);
+  }
+
+  public void unblock() {
+    blockThread.set(false);
+  }
+
+  @Override
+  public void run() {
+    runImpl();
+  }
+
+  @Override
+  public Object call() throws Exception {
+    runImpl();
+    return null;
+  }
+
+  private void runImpl() {
+    while (blockThread.get()) {
+      // busy-wait to block thread
+    }
+    if (doTraceableWork) {
+      asyncChild();
+    }
+  }
+
+  @Trace(operationName = "asyncChild")
+  private void asyncChild() {}
+}

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskRunnerInstrumentation.java

@@ -0,0 +1,61 @@
+package datadog.trace.instrumentation.java.concurrent;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class TaskRunnerInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType {
+  public TaskRunnerInstrumentation() {
+    super("java_concurrent", "task-runner");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "java.util.concurrent.ThreadPerTaskExecutor$TaskRunner";
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("java.lang.Runnable", State.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
+    transformer.applyAdvice(isMethod().and(named("run")), getClass().getName() + "$Run");
+  }
+
+  public static final class Construct {
+    @Advice.OnMethodExit
+    public static void captureScope(@Advice.This Runnable task) {
+      capture(InstrumentationContext.get(Runnable.class, State.class), task, true);
+    }
+  }
+
+  public static final class Run {
+    @Advice.OnMethodEnter
+    public static AgentScope activate(@Advice.This Runnable task) {
+      return startTaskScope(InstrumentationContext.get(Runnable.class, State.class), task);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void close(@Advice.Enter AgentScope scope) {
+      endTaskScope(scope);
+    }
+  }
+}

gradle.properties

@@ -4,4 +4,4 @@ org.gradle.jvmargs=-XX:MaxMetaspaceSize=1g
 org.gradle.java.installations.auto-detect=false
 org.gradle.java.installations.auto-download=false
 # 8 and 11 is needed to build
-org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME
+org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME

lib-injection/build_java_agent.sh

@@ -1,8 +1,8 @@
 #!/bin/sh
 
-JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 ./gradlew dd-java-agent:build dd-java-agent:shadowJar --build-cache --parallel --no-daemon --max-workers=8
+JAVA_HOME=$JAVA_HOME_8_X64 JAVA_8_HOME=$JAVA_HOME_8_X64 JAVA_11_HOME=$JAVA_HOME_11_X64 JAVA_17_HOME=$JAVA_HOME_17_X64 JAVA_21_HOME=$JAVA_HOME_21_X64 ./gradlew dd-java-agent:build dd-java-agent:shadowJar --build-cache --parallel --no-daemon --max-workers=8
 cp workspace/dd-java-agent/build/libs/dd-java-agent-*.jar lib-injection/
 rm lib-injection/*-sources.jar lib-injection/*-javadoc.jar
 mv lib-injection/*.jar lib-injection/dd-java-agent.jar
 echo "Java tracer copied to lib-injection folder"
-ls lib-injection/
+ls lib-injection/