Add support for Java 25 Structured Concurrency API

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

@@ -0,0 +1,55 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_25
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'idea'
+
+muzzle {
+  pass {
+    coreJdk('25')
+  }
+}
+
+idea {
+  module {
+    jdkName = '25'
+  }
+}
+
+/*
+ * Declare previewTest, a test suite that requires the Javac/Java --enable-preview feature flag
+ */
+addTestSuite('previewTest')
+// Configure groovy test file compilation
+compilePreviewTestGroovy.configure {
+  javaLauncher = javaToolchains.launcherFor {
+    languageVersion = JavaLanguageVersion.of(25)
+  }
+  options.compilerArgs.add("--enable-preview")
+}
+// Configure Java test files compilation
+compilePreviewTestJava.configure {
+  options.compilerArgs.add("--enable-preview")
+}
+// Configure tests execution
+previewTest.configure {
+  jvmArgs = ['--enable-preview']
+}
+// Require the preview test suite to run as part of module check
+tasks.named("check").configure {
+  dependsOn "previewTest"
+}
+
+dependencies {
+  testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
+}
+
+// Set all compile tasks to use JDK21 but let instrumentation code targets 1.8 compatibility
+project.tasks.withType(AbstractCompile).configureEach {
+  setJavaVersion(it, 25)
+}
+compileJava.configure {
+  sourceCompatibility = JavaVersion.VERSION_1_8
+  targetCompatibility = JavaVersion.VERSION_1_8
+}

dd-java-agent/instrumentation/java-concurrent/java-concurrent-25/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency/StructuredTaskScopeInstrumentation.java

@@ -0,0 +1,67 @@
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency;
+
+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;
+
+/**
+ * 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).
+ */
+@SuppressWarnings("unused")
+@AutoService(InstrumenterModule.class)
+public class StructuredTaskScopeInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  private static final String SUBTASK_IMPL_CLASS_NAME =
+      "java.util.concurrent.StructuredTaskScopeImpl.SubtaskImpl";
+
+  public StructuredTaskScopeInstrumentation() {
+    super("java_concurrent", "structured_task_scope");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return SUBTASK_IMPL_CLASS_NAME;
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return JavaVirtualMachine.isJavaVersionAtLeast(25) && super.isEnabled();
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(SUBTASK_IMPL_CLASS_NAME, State.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(isConstructor(), getClass().getName() + "$ConstructorAdvice");
+  }
+
+  public static final class ConstructorAdvice {
+    @Advice.OnMethodExit
+    public static <T> void captureScope(
+        @Advice.This Object task // StructuredTaskScopeImpl.SubtaskImpl
+        // (the advice are compile against Java 8 so the type from JDK25 can't be referred)
+        ) {
+      ContextStore<Object, State> contextStore =
+          InstrumentationContext.get(
+              SUBTASK_IMPL_CLASS_NAME,
+              "datadog.trace.bootstrap.instrumentation.java.concurrent.State");
+      capture(contextStore, task);
+    }
+  }
+}

dd-java-agent/instrumentation/java-concurrent/java-concurrent-25/src/previewTest/groovy/StructuredConcurrencyTest.groovy

@@ -0,0 +1,170 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.Trace
+import spock.lang.Timeout
+
+import java.util.concurrent.Callable
+import java.util.concurrent.StructuredTaskScope
+
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace
+
+class StructuredConcurrencyTest extends AgentTestRunner {
+  /**
+   * Tests the structured task scope with a single task.
+   */
+  @Timeout(10)
+  def "test single task"() {
+    setup:
+    def taskScope = StructuredTaskScope.open()
+    def result = false
+
+    when:
+    runUnderTrace("parent") {
+      def task = taskScope.fork(new Callable<Boolean>() {
+          @Trace(operationName = "child")
+          @Override
+          Boolean call() throws Exception {
+            return true
+          }
+        })
+      taskScope.join()
+      result = task.get()
+    }
+    taskScope.close()
+
+    then:
+    result
+    assertTraces(1) {
+      sortSpansByStart()
+      trace(2) {
+        span(0) {
+          parent()
+          operationName "parent"
+        }
+        span(1) {
+          childOfPrevious()
+          operationName "child"
+        }
+      }
+    }
+  }
+
+  /**
+   * Tests the structured task scope with a multiple tasks.
+   * Here is the expected task/span structure:
+   * <pre>
+   *   parent
+   *   |-- child1
+   *   |-- child2
+   *   \-- child3
+   * </pre>
+   */
+  @Timeout(10)
+  def "test multiple tasks"() {
+    setup:
+    def taskScope = StructuredTaskScope.open()
+
+    when:
+    runUnderTrace("parent") {
+      taskScope.fork {
+        runnableUnderTrace("child1") {}
+      }
+      taskScope.fork {
+        runnableUnderTrace("child2") {}
+      }
+      taskScope.fork {
+        runnableUnderTrace("child3") {}
+      }
+      taskScope.join()
+    }
+    taskScope.close()
+
+    then:
+    assertTraces(1) {
+      sortSpansByStart()
+      trace(4) {
+        span {
+          parent()
+          operationName "parent"
+        }
+        def parent = span(0)
+        span {
+          childOf(parent)
+          assert span.operationName.toString().startsWith("child")
+        }
+        span {
+          childOf(parent)
+          assert span.operationName.toString().startsWith("child")
+        }
+        span {
+          childOf(parent)
+          assert span.operationName.toString().startsWith("child")
+        }
+      }
+    }
+  }
+
+  /**
+   * Tests the structured task scope with a multiple nested tasks.
+   * Here is the expected task/span structure:
+   * <pre>
+   *   parent
+   *   |-- child1
+   *   |   |-- great-child1-1
+   *   |   \-- great-child1-2
+   *   \-- child2
+   * </pre>
+   */
+  @Timeout(10)
+  def "test nested tasks"() {
+    setup:
+    def taskScope = StructuredTaskScope.open()
+
+    when:
+    runUnderTrace("parent") {
+      taskScope.fork {
+        runnableUnderTrace("child1") {
+          taskScope.fork {
+            runnableUnderTrace("great-child1-1") {}
+          }
+          taskScope.fork {
+            runnableUnderTrace("great-child1-2") {}
+          }
+        }
+      }
+      taskScope.fork {
+        runnableUnderTrace("child2") {}
+      }
+      taskScope.join()
+    }
+    taskScope.close()
+
+    then:
+    assertTraces(1) {
+      sortSpansByStart()
+      trace(5) {
+        // Check parent span
+        span {
+          parent()
+          operationName "parent"
+        }
+        def parent = span(0)
+        // Check child and great child spans
+        def child1 = null
+        for (i in 0..<4) {
+          span {
+            def name = span.operationName.toString()
+            if (name.startsWith("child")) {
+              childOf(parent)
+              if (name == "child1") {
+                child1 = span
+              }
+            } else if (name.startsWith("great-child1")) {
+              childOf(child1) // We can assume child1 will be set as spans are sorted by start time
+            }
+          }
+        }
+      }
+    }
+  }
+}

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,JAVA_21_HOME
+org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME,JAVA_25_HOME

settings.gradle.kts

@@ -352,6 +352,7 @@ include(
   ":dd-java-agent:instrumentation:java-concurrent",
   ":dd-java-agent:instrumentation:java-concurrent:java-completablefuture",
   ":dd-java-agent:instrumentation:java-concurrent:java-concurrent-21",
+  ":dd-java-agent:instrumentation:java-concurrent:java-concurrent-25",
   ":dd-java-agent:instrumentation:java-concurrent:lambda-testing",
   ":dd-java-agent:instrumentation:java-directbytebuffer",
   ":dd-java-agent:instrumentation:java-http-client",