Add support for Java 25 Structured Concurrency API

dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency/StructuredTaskScopeInstrumentation.java → dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency21/StructuredTaskScope21Instrumentation.java

@@ -1,31 +1,32 @@
-package datadog.trace.instrumentation.java.concurrent.structuredconcurrency;
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency21;
 
+import static datadog.environment.JavaVirtualMachine.isJavaVersionBetween;
+import static datadog.trace.bootstrap.InstrumentationContext.get;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
 import com.google.auto.service.AutoService;
-import datadog.environment.JavaVirtualMachine;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import java.util.Map;
-import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.OnMethodExit;
+import net.bytebuddy.asm.Advice.This;
 
 /**
  * This instrumentation captures the active span scope at StructuredTaskScope task creation
- * (SubtaskImpl). The scope is then activate and close through the Runnable instrumentation
- * (SubtaskImpl implementation Runnable).
+ * (SubtaskImpl). The scope is then activate and close through the {@link Runnable} instrumentation
+ * (SubtaskImpl implementing {@link Runnable}).
  */
 @SuppressWarnings("unused")
 @AutoService(InstrumenterModule.class)
-public class StructuredTaskScopeInstrumentation extends InstrumenterModule.Tracing
+public class StructuredTaskScope21Instrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
-  public StructuredTaskScopeInstrumentation() {
-    super("java_concurrent", "structured_task_scope");
+  public StructuredTaskScope21Instrumentation() {
+    super("java_concurrent", "structured-task-scope", "structured-task-scope-21");
   }
 
   @Override
@@ -35,13 +36,12 @@ public String instrumentedType() {
 
   @Override
   public boolean isEnabled() {
-    return JavaVirtualMachine.isJavaVersionAtLeast(21) && super.isEnabled();
+    return isJavaVersionBetween(21, 25) && super.isEnabled();
   }
 
   @Override
   public Map<String, String> contextStore() {
-    return singletonMap(
-        "java.util.concurrent.StructuredTaskScope.SubtaskImpl", State.class.getName());
+    return singletonMap(Runnable.class.getName(), State.class.getName());
   }
 
   @Override
@@ -50,15 +50,17 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static final class ConstructorAdvice {
-    @Advice.OnMethodExit
-    public static <T> void captureScope(
-        @Advice.This Object task // StructuredTaskScope.SubtaskImpl (can't use the type)
-        ) {
-      ContextStore<Object, State> contextStore =
-          InstrumentationContext.get(
-              "java.util.concurrent.StructuredTaskScope.SubtaskImpl",
-              "datadog.trace.bootstrap.instrumentation.java.concurrent.State");
-      capture(contextStore, task);
+    /**
+     * Captures task scope to be restored at the start of VirtualThread.run() method by {@link
+     * Runnable} instrumentation.
+     *
+     * @param subTaskImpl The StructuredTaskScope.SubtaskImpl object (the advice are compile against
+     *     Java 8 so the type from JDK25 can't be referred, using {@link Object} instead
+     */
+    @OnMethodExit(suppress = Throwable.class)
+    public static void captureScope(@This Object subTaskImpl) {
+      ContextStore<Runnable, State> contextStore = get(Runnable.class, State.class);
+      capture(contextStore, (Runnable) subTaskImpl);
     }
   }
 }

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

@@ -0,0 +1,7 @@
+apply from: "$rootDir/gradle/java.gradle"
+
+muzzle {
+  pass {
+    coreJdk('25')
+  }
+}

dd-java-agent/instrumentation/java-concurrent/java-concurrent-25/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency25/StructuredTaskScope25Instrumentation.java

@@ -0,0 +1,65 @@
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency25;
+
+import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast;
+import static datadog.trace.bootstrap.InstrumentationContext.get;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import net.bytebuddy.asm.Advice.OnMethodExit;
+import net.bytebuddy.asm.Advice.This;
+
+// WARNING:
+// This instrumentation is tested using smoke tests as instrumented tests cannot run using Java 25.
+// Instrumented tests rely on Spock / Groovy which cannot run using Java 25 due to byte-code
+// compatibility. Check dd-java-agent/instrumentation/java-concurrent/java-concurrent-25 for this
+// instrumentation test suite.
+
+/**
+ * This instrumentation captures the active span scope at StructuredTaskScope task creation
+ * (SubtaskImpl). The scope is then activate and close through the {@link Runnable} instrumentation
+ * (SubtaskImpl implementing {@link Runnable}).
+ */
+@SuppressWarnings("unused")
+@AutoService(InstrumenterModule.class)
+public class StructuredTaskScope25Instrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public StructuredTaskScope25Instrumentation() {
+    super("java_concurrent", "structured-task-scope", "structured-task-scope-25");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "java.util.concurrent.StructuredTaskScopeImpl$SubtaskImpl";
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return isJavaVersionAtLeast(25) && super.isEnabled();
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(isConstructor(), getClass().getName() + "$ConstructorAdvice");
+  }
+
+  public static final class ConstructorAdvice {
+    /**
+     * Captures task scope to be restored at the start of VirtualThread.run() method by {@link
+     * Runnable} instrumentation.
+     *
+     * @param subTaskImpl The StructuredTaskScopeImpl.SubtaskImpl object (the advice are compile
+     *     against Java 8 so the type from JDK25 can't be referred, using {@link Object} instead
+     */
+    @OnMethodExit(suppress = Throwable.class)
+    public static void captureScope(@This Object subTaskImpl) {
+      ContextStore<Runnable, State> contextStore = get(Runnable.class, State.class);
+      capture(contextStore, (Runnable) subTaskImpl);
+    }
+  }
+}

dd-smoke-tests/concurrent/java-21/build.gradle

@@ -5,6 +5,7 @@ plugins {
 
 ext {
   minJavaVersionForTests = JavaVersion.VERSION_21
+  maxJavaVersionForTests = JavaVersion.VERSION_25
 }
 
 apply from: "$rootDir/gradle/java.gradle"

dd-smoke-tests/concurrent/java-25/build.gradle

@@ -0,0 +1,62 @@
+plugins {
+  id 'application'
+  id 'com.gradleup.shadow'
+}
+
+ext {
+  // This smoke test should be limited to Java 25 and above
+  // But the groovy testing framework cannot run on Java 25
+  // Using Java 8 for now and runs a JVM 25 when forking tests process.
+  // Relying on forked JVM 25 is hardcoded in the test suites (see createProcessBuilder()).
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+  maxJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+description = 'JDK 25 Concurrent Integration Tests'
+
+tasks.named('compileJava').configure {
+  setJavaVersion(it, 25)
+  sourceCompatibility = JavaVersion.VERSION_25
+  targetCompatibility = JavaVersion.VERSION_25
+  options.compilerArgs.add("--enable-preview")
+  javaToolchains {
+    java {
+      toolchain {
+        languageVersion = JavaLanguageVersion.of(25)
+      }
+    }
+  }
+}
+
+tasks.named('compileTestGroovy').configure {
+  setJavaVersion(it, 8)
+  sourceCompatibility = JavaVersion.VERSION_1_8
+  targetCompatibility = JavaVersion.VERSION_1_8
+}
+
+// Disable plugin tasks that do not support Java 25:
+// * forbiddenApis is missing classes
+// * spotless as the google-java-format version does not support Java 25 and can't be changed once applied
+// * spotbugs failed to read class using newer bytecode versions
+forbiddenApisMain {
+  failOnMissingClasses = false
+}
+['spotlessApply', 'spotlessCheck', 'spotlessJava', 'spotbugsMain'].each {
+  tasks.named(it).configure { enabled = false }
+}
+
+application {
+  mainClass = 'datadog.smoketest.concurrent.ConcurrentApp'
+}
+
+dependencies {
+  implementation group: 'io.opentelemetry.instrumentation', name: 'opentelemetry-instrumentation-annotations', version: '2.19.0'
+  testImplementation project(':dd-smoke-tests')
+}
+
+tasks.withType(Test).configureEach {
+  dependsOn "shadowJar"
+  jvmArgs "-Ddatadog.smoketest.shadowJar.path=${tasks.shadowJar.archiveFile.get()}"
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/ConcurrentApp.java

@@ -0,0 +1,26 @@
+package datadog.smoketest.concurrent;
+
+public class ConcurrentApp {
+  void main(String[] args) {
+    if (args.length != 1) {
+      System.err.println("Missing test case argument");
+      System.exit(1);
+    }
+    String name = args[0];
+    try {
+      fromName(name).run();
+    } catch (InterruptedException e) {
+      throw new RuntimeException("Failed to test case " + name, e);
+    }
+  }
+
+  private static TestCase fromName(String name) {
+    return switch (name) {
+      case "NestedTasks" -> new NestedTasks();
+      case "MultipleTasks" -> new MultipleTasks();
+      case "SimpleCallableTask" -> new SimpleCallableTask();
+      case "SimpleRunnableTask" -> new SimpleRunnableTask();
+      default -> throw new IllegalArgumentException("Invalid test case name " + name);
+    };
+  }
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/MultipleTasks.java

@@ -0,0 +1,32 @@
+package datadog.smoketest.concurrent;
+
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import java.util.concurrent.StructuredTaskScope;
+
+public final class MultipleTasks implements TestCase {
+  @WithSpan("parent")
+  public void run() throws InterruptedException {
+    try (var scope = StructuredTaskScope.open()) {
+      scope.fork(this::runnableTask1);
+      scope.fork(this::callableTask2);
+      scope.fork(this::runnableTask3);
+      scope.join();
+    }
+  }
+
+  @WithSpan("child1")
+  void runnableTask1() {
+    // Some basic computations here
+  }
+
+  @WithSpan("child2")
+  Boolean callableTask2() {
+    // Some basic computations here
+    return true;
+  }
+
+  @WithSpan("child3")
+  void runnableTask3() {
+    // Some basic computations here
+  }
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/NestedTasks.java

@@ -0,0 +1,40 @@
+package datadog.smoketest.concurrent;
+
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import java.util.concurrent.StructuredTaskScope;
+
+public final class NestedTasks implements TestCase {
+  @WithSpan("parent")
+  public void run() throws InterruptedException {
+    try (var scope = StructuredTaskScope.open()) {
+      scope.fork(this::task1);
+      scope.fork(this::task2);
+      scope.join();
+    }
+  }
+
+  @WithSpan("child1")
+  void task1() {
+    try (var scope = StructuredTaskScope.open()) {
+      scope.fork(this::subTask1);
+      scope.fork(this::subTask2);
+      scope.join();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+  }
+
+  @WithSpan("child2")
+  void task2() {
+    // Some basic computations here
+  }
+
+  @WithSpan("great-child1-1")
+  void subTask1() {
+    // Some basic computations here
+  }
+  @WithSpan("great-child1-2")
+  void subTask2() {
+    // Some basic computations here
+  }
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/SimpleCallableTask.java

@@ -0,0 +1,20 @@
+package datadog.smoketest.concurrent;
+
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import java.util.concurrent.StructuredTaskScope;
+
+public final class SimpleCallableTask implements TestCase {
+  @WithSpan("parent")
+  public void run() throws InterruptedException {
+    try (var scope = StructuredTaskScope.open()) {
+      scope.fork(this::doSomething);
+      scope.join();
+    }
+  }
+
+  @WithSpan("child")
+  Boolean doSomething() {
+    // Some basic computations here
+    return true;
+  }
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/SimpleRunnableTask.java

@@ -0,0 +1,19 @@
+package datadog.smoketest.concurrent;
+
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import java.util.concurrent.StructuredTaskScope;
+
+public final class SimpleRunnableTask implements TestCase {
+  @WithSpan("parent")
+  public void run() throws InterruptedException {
+    try (var scope = StructuredTaskScope.open()) {
+      scope.fork(this::doSomething);
+      scope.join();
+    }
+  }
+
+  @WithSpan("child")
+  void doSomething() {
+    // Some basic computations here
+  }
+}

dd-smoke-tests/concurrent/java-25/src/main/java/datadog/smoketest/concurrent/TestCase.java

@@ -0,0 +1,6 @@
+package datadog.smoketest.concurrent;
+
+@FunctionalInterface
+public interface TestCase {
+  void run() throws InterruptedException;
+}