Improve structured concurrency instrumentation

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

@@ -0,0 +1,42 @@
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import static java.lang.invoke.MethodHandles.publicLookup;
+import static java.lang.invoke.MethodType.methodType;
+
+import java.lang.invoke.MethodHandle;
+
+/** Helper for the java-concurrent-25.0 {@code StructuredTaskScopeForkInstrumentation}. */
+public final class StructuredTaskScopeHelper {
+  private static final MethodHandle IS_CANCELLED = resolveIsCancelled();
+
+  private StructuredTaskScopeHelper() {}
+
+  private static MethodHandle resolveIsCancelled() {
+    try {
+      ClassLoader classLoader = StructuredTaskScopeHelper.class.getClassLoader();
+      Class<?> scopeClass =
+          Class.forName("java.util.concurrent.StructuredTaskScope", false, classLoader);
+      return publicLookup().findVirtual(scopeClass, "isCancelled", methodType(boolean.class));
+    } catch (Throwable ignored) {
+      return null;
+    }
+  }
+
+  /**
+   * Returns whether the given {@code StructuredTaskScope} has been cancelled.
+   *
+   * @param scope a {@code java.util.concurrent.StructuredTaskScope} instance.
+   * @return {@code true} if the scope is cancelled; {@code false} otherwise, including when {@code
+   *     isCancelled()} could not be resolved or invoked.
+   */
+  public static boolean isCancelled(Object scope) {
+    if (IS_CANCELLED == null || scope == null) {
+      return false;
+    }
+    try {
+      return (boolean) IS_CANCELLED.invoke(scope);
+    } catch (Throwable ignored) {
+      return false;
+    }
+  }
+}

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

@@ -1,3 +1,6 @@
+import static org.gradle.api.JavaVersion.VERSION_1_8
+import static org.gradle.api.JavaVersion.VERSION_24
+
 plugins {
   id 'idea'
 }
@@ -31,13 +34,13 @@ tasks.named("previewTest").configure {
   testJvmConstraints {
     // Structured concurrency is a preview feature in Java 21. Methods (e.g. ShutdownOnFailure) used in this instrumentation test are no longer available in Java 25, so we set the max version to 24.
     // See: https://download.java.net/java/early_access/loom/docs/api/java.base/java/util/concurrent/StructuredTaskScope.html
-    maxJavaVersion = JavaVersion.VERSION_24
+    maxJavaVersion = VERSION_24
   }
 }
 
 // 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)
+  configureCompiler(it, 21, VERSION_1_8)
 }
 
 // Configure groovy test file compilation

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

@@ -49,7 +49,7 @@ class StructuredConcurrencyTest extends InstrumentationSpecification {
   }
 
   /**
-   * Tests the structured task scope with a multiple tasks.
+   * Tests the structured task scope with multiple tasks.
    * Here is the expected task/span structure:
    * <pre>
    *   parent

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

@@ -1,3 +1,7 @@
+import static org.gradle.api.JavaVersion.VERSION_25
+import static org.gradle.api.JavaVersion.VERSION_26
+import static org.gradle.api.JavaVersion.VERSION_27
+
 plugins {
   id 'idea'
 }
@@ -6,6 +10,11 @@ apply from: "$rootDir/gradle/java.gradle"
 // Use slf4j-simple as default; logback has a high chance of getting stuck in a deadlock on CI.
 apply from: "$rootDir/gradle/slf4j-simple.gradle"
 
+testJvmConstraints {
+  // The new StructuredTaskScope API (JEP 505) only exists from Java 25
+  minJavaVersion = VERSION_25
+}
+
 muzzle {
   pass {
     coreJdk('25')
@@ -17,3 +26,65 @@ idea {
     jdkName = '25'
   }
 }
+
+/*
+ * Declare previewTest, a test suite that requires the Javac/Java --enable-preview feature flag.
+ * --enable-preview class files are version-locked to the compiling JDK, so the suite is compiled with
+ * the SAME JDK that runs it (the selected test JVM, clamped to >= 25).
+ */
+addTestSuite('previewTest')
+
+def previewTestJvm = tasks.named("previewTest", Test)
+  .flatMap { it.javaLauncher }
+  .map { Math.max(25, it.metadata.languageVersion.asInt()) }
+  .orElse(25)
+
+tasks.named("compilePreviewTestJava", JavaCompile) {
+  javaCompiler = javaToolchains.compilerFor {
+    languageVersion = previewTestJvm.map {
+      JavaLanguageVersion.of(it)
+    }
+  }
+  options.release = previewTestJvm
+  options.compilerArgs.add("--enable-preview")
+}
+tasks.named("previewTest", Test) {
+  jvmArgs = ['--enable-preview']
+}
+
+/*
+ * Per-JDK suites test:
+ * Each suite has its own source set, is compiled for one JDK,
+ * and is pinned to it (preview class files are version-locked).
+ */
+def previewCheckSuites = ['previewTest']
+def addJoinerLeakSuite = { int jdk, JavaVersion version ->
+  String suite = "previewTestJdk$jdk"
+  addTestSuite(suite)
+  tasks.named("compile${suite.capitalize()}Java", JavaCompile) {
+    configureCompiler(it, jdk)
+    options.compilerArgs.add("--enable-preview")
+  }
+  tasks.named(suite, Test) {
+    jvmArgs = ['--enable-preview']
+    testJvmConstraints {
+      minJavaVersion = version
+      maxJavaVersion = version
+    }
+  }
+  previewCheckSuites << suite
+}
+
+addJoinerLeakSuite(25, VERSION_25)
+if (System.getenv('JAVA_26_HOME') != null) {
+  addJoinerLeakSuite(26, VERSION_26)
+}
+if (System.getenv('JAVA_27_HOME') != null) {
+  addJoinerLeakSuite(27, VERSION_27)
+}
+
+// Require the preview test suites to run as part of module check
+tasks.named("check") {
+  dependsOn previewCheckSuites
+}
+

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

@@ -1,66 +1,60 @@
 package datadog.trace.instrumentation.java.concurrent.structuredconcurrency25;
 
-import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 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 static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.cancelTask;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
-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 datadog.trace.bootstrap.instrumentation.java.concurrent.StructuredTaskScopeHelper;
+import java.util.concurrent.Callable;
 import net.bytebuddy.asm.Advice.OnMethodExit;
+import net.bytebuddy.asm.Advice.Return;
 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/java-concurrent/java-concurrent-25.0 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}).
+ * This instrumentation cancels the continuation captured by {@link
+ * StructuredTaskScope25TaskInstrumentation} for a subtask forked into an already-canceled scope.
+ *
+ * <p>In that case, the subtask's thread is never started, so {@code SubtaskImpl.run()} never runs,
+ * and the {@link Runnable} instrumentation never activates/closes the captured continuation. It
+ * leads to continuation leak. This instrumentation verifies the subtask was canceled and releases
+ * the related continuation.
  */
 @SuppressWarnings("unused")
-@AutoService(InstrumenterModule.class)
-public class StructuredTaskScope25Instrumentation extends InstrumenterModule.ContextTracking
+public class StructuredTaskScope25Instrumentation
     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();
+    return "java.util.concurrent.StructuredTaskScopeImpl";
   }
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isConstructor(), getClass().getName() + "$ConstructorAdvice");
+    transformer.applyAdvice(
+        named("fork").and(takesArgument(0, Callable.class)), getClass().getName() + "$ForkAdvice");
   }
 
-  public static final class ConstructorAdvice {
+  public static final class ForkAdvice {
     /**
-     * Captures task scope to be restored at the start of VirtualThread.run() method by {@link
-     * Runnable} instrumentation.
+     * Cancels the task scope continuation captured at task creation when the subtask is forked into
+     * an already-canceled scope: its thread never starts, so the {@link Runnable} instrumentation
+     * never releases the continuation.
      *
-     * @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
+     * @param scope The StructuredTaskScopeImpl object (the advice is compiled against Java 8 so the
+     *     type from JDK25 can't be referred, using {@link Object} instead).
+     * @param subtask The StructuredTaskScopeImpl.SubtaskImpl object (the advice is compiled 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);
+    public static void afterFork(@This Object scope, @Return Object subtask) {
+      if (subtask instanceof Runnable && StructuredTaskScopeHelper.isCancelled(scope)) {
+        ContextStore<Runnable, State> contextStore = get(Runnable.class, State.class);
+        cancelTask(contextStore, (Runnable) subtask);
+      }
     }
   }
 }

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-25.0/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency25/StructuredTaskScope25Module.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency25;
+
+import static datadog.environment.JavaVirtualMachine.isJavaVersionAtLeast;
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonMap;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * This module propagates context across {@code java.util.concurrent.StructuredTaskScope} forked
+ * subtasks (JDK 25+).
+ */
+@SuppressWarnings("unused")
+@AutoService(InstrumenterModule.class)
+public class StructuredTaskScope25Module extends InstrumenterModule.ContextTracking {
+  public StructuredTaskScope25Module() {
+    super("java_concurrent", "structured-task-scope", "structured-task-scope-25");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return isJavaVersionAtLeast(25) && super.isEnabled();
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(Runnable.class.getName(), State.class.getName());
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    return asList(
+        new StructuredTaskScope25TaskInstrumentation(), new StructuredTaskScope25Instrumentation());
+  }
+}

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-25.0/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency25/StructuredTaskScope25TaskInstrumentation.java

@@ -0,0 +1,46 @@
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency25;
+
+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 datadog.trace.agent.tooling.Instrumenter;
+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;
+
+/**
+ * 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")
+public class StructuredTaskScope25TaskInstrumentation
+    implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "java.util.concurrent.StructuredTaskScopeImpl$SubtaskImpl";
+  }
+
+  @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 is compiled
+     *     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);
+    }
+  }
+}