feat(concurrent): Improve loom support with structured concurrency

Add structured concurrency support
Improve virtual thread tests
Create a dedicated concurrency module for JDK21+

Author: PerfectSlayer

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

@@ -1,7 +1,3 @@
-ext {
-  latestDepTestMinJavaVersionForTests = JavaVersion.VERSION_21
-}
-
 muzzle {
   pass {
     coreJdk()
@@ -10,17 +6,11 @@ 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/java-concurrent-21/build.gradle

@@ -0,0 +1,55 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_21
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'idea'
+
+muzzle {
+  pass {
+    coreJdk('21')
+  }
+}
+
+idea {
+  module {
+    jdkName = '21'
+  }
+}
+
+/*
+ * 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(21)
+  }
+  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, 21)
+}
+compileJava.configure {
+  sourceCompatibility = JavaVersion.VERSION_1_8
+  targetCompatibility = JavaVersion.VERSION_1_8
+}

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

@@ -0,0 +1,77 @@
+package datadog.trace.instrumentation.java.concurrent.structuredconcurrency;
+
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.FORK_JOIN_TASK;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.ExcludeFilterProvider;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Platform;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import java.util.Collection;
+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, ExcludeFilterProvider {
+
+  public StructuredTaskScopeInstrumentation() {
+    super("java_concurrent", "structured_task_scope");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "java.util.concurrent.StructuredTaskScope$SubtaskImpl";
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return Platform.isJavaVersionAtLeast(21) && super.isEnabled();
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "java.util.concurrent.StructuredTaskScope.SubtaskImpl", State.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(isConstructor(), getClass().getName() + "$ConstructorAdvice");
+  }
+
+  @Override
+  public Map<ExcludeFilter.ExcludeType, ? extends Collection<String>> excludedClasses() {
+    // Prevent the ForkJoinPool instrumentation to enable the task scope too early on the carrier
+    // thread rather than on the expected running thread, which is virtual by default.
+    return singletonMap(
+        FORK_JOIN_TASK, singleton("java.util.concurrent.ForkJoinTask$RunnableExecuteAction"));
+  }
+
+  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, true);
+    }
+  }
+}

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskRunnerInstrumentation.java renamed to dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/virtualthread/TaskRunnerInstrumentation.java

@@ -1,4 +1,4 @@
-package datadog.trace.instrumentation.java.concurrent;
+package datadog.trace.instrumentation.java.concurrent.virtualthread;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
@@ -11,12 +11,17 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Platform;
 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;
 
+/**
+ * Instruments {@code TaskRunner}, internal runnable for {@code ThreadPerTaskExecutor} (JDK 19+ as
+ * preview, 21+ as stable), the executor with default virtual thread factory.
+ */
 @AutoService(InstrumenterModule.class)
 public final class TaskRunnerInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType {
@@ -29,6 +34,11 @@ public String instrumentedType() {
     return "java.util.concurrent.ThreadPerTaskExecutor$TaskRunner";
   }
 
+  @Override
+  public boolean isEnabled() {
+    return Platform.isJavaVersionAtLeast(19) && super.isEnabled();
+  }
+
   @Override
   public Map<String, String> contextStore() {
     return singletonMap("java.lang.Runnable", State.class.getName());

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

@@ -0,0 +1,167 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.Trace
+
+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
+import static java.time.Instant.now
+
+class StructuredConcurrencyTest extends AgentTestRunner {
+  /**
+   * Tests the structured task scope with a single task.
+   */
+  def "test single task"() {
+    setup:
+    def taskScope = new StructuredTaskScope.ShutdownOnFailure()
+    def result = false
+
+    when:
+    runUnderTrace("parent") {
+      def task = taskScope.fork(new Callable<Boolean>() {
+          @Trace(operationName = "child")
+          @Override
+          Boolean call() throws Exception {
+            return true
+          }
+        })
+      taskScope.joinUntil(now() + 1) // Wait for a second at maximum
+      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>
+   */
+  def "test multiple tasks"() {
+    setup:
+    def taskScope = new StructuredTaskScope.ShutdownOnFailure()
+
+    when:
+    runUnderTrace("parent") {
+      taskScope.fork {
+        runnableUnderTrace("child1") {}
+      }
+      taskScope.fork {
+        runnableUnderTrace("child2") {}
+      }
+      taskScope.fork {
+        runnableUnderTrace("child3") {}
+      }
+      taskScope.joinUntil(now() + 2) // Wait for two seconds at maximum
+    }
+    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>
+   */
+  def "test nested tasks"() {
+    setup:
+    def taskScope = new StructuredTaskScope.ShutdownOnFailure()
+
+    when:
+    runUnderTrace("parent") {
+      taskScope.fork {
+        runnableUnderTrace("child1") {
+          taskScope.fork {
+            runnableUnderTrace("great-child1-1") {}
+          }
+          taskScope.fork {
+            runnableUnderTrace("great-child1-2") {}
+          }
+        }
+      }
+      taskScope.fork {
+        runnableUnderTrace("child2") {}
+      }
+      taskScope.joinUntil(now() + 2) // Wait for two seconds at maximum
+    }
+    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
+            }
+          }
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/java-concurrent/src/latestDepTest/groovy/VirtualThreadTest.groovy renamed to dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/test/groovy/VirtualThreadTest.groovy

@@ -11,7 +11,6 @@ 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