fix(virtual-thread): Add support for early VirtualThread afterTerminate method

PerfectSlayer · PerfectSlayer · commit d223f7924b13 · 2026-04-16T10:49:54.000+02:00
Support both afterDone() and afterTerminate() cleanup methods to handle different JDK 21 virtual thread implementations. Both methods cancel the help continuation and release the context scope when the virtual thread completes.
diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java
@@ -45,8 +45,8 @@
  *   <li>{@code unmount()}: swaps the carrier thread's original context back, saving the virtual
  *       thread's (possibly modified) context for the next mount.
  *   <li>Steps 2-3 repeat on each park/unpark cycle, potentially on different carrier threads.
- *   <li>{@code afterDone()}: cancels the help continuation, releasing the context scope to be
- *       closed.
+ *   <li>{@code afterDone()} / {@code afterTerminate()} for early VirtualThread support: cancels the
+ *       help continuation, releasing the context scope to be closed.
  * </ol>
  *
  * <p>Instrumenting the internal {@code VirtualThread.runContinuation()} method does not work as the
@@ -111,6 +111,9 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod().and(named("afterDone")).and(takesArguments(boolean.class)),
         getClass().getName() + "$AfterDone");
+    transformer.applyAdvice(
+        isMethod().and(named("afterTerminate")).and(takesArguments(boolean.class, boolean.class)),
+        getClass().getName() + "$AfterDone");
   }
 
   public static final class Construct {
@@ -153,7 +156,7 @@ public static void onUnmount(@Advice.This Object virtualThread) {
 
   public static final class AfterDone {
     @OnMethodEnter(suppress = Throwable.class)
-    public static void onTerminate(@Advice.This Object virtualThread) {
+    public static void onDone(@Advice.This Object virtualThread) {
       ContextStore<Object, VirtualThreadState> store =
           InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
       VirtualThreadState state = store.remove(virtualThread);
