feat(concurrency): Improve JDK 21 structured task scope instrumentation

Add instrumentation alias
Limit activation to Java < 25
Fix context store declaration and usage to match RunnableInstrumentation
Remove unused generic

Author: PerfectSlayer

…/StructuredTaskScopeInstrumentation.java …tructuredTaskScope21Instrumentation.javadd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency/StructuredTaskScopeInstrumentation.java renamed to dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency21/StructuredTaskScope21Instrumentation.java dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/main/java/datadog/trace/instrumentation/java/concurrent/structuredconcurrency/StructuredTaskScopeInstrumentation.java renamed to 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-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"