fix: ForkJoinTask captured context is retained until GC (#626)

arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/Cache.java

@@ -2,7 +2,17 @@
 
 public interface Cache<K, V> {
 
-    Cache<Object, Object> CAPTURED_CACHE = weakMap();
+    /**
+     * Declared as the concrete {@link WeakCache} rather than {@code Cache} so the ForkJoinTask
+     * exec/run exit advice can drop its entry explicitly. {@code Cache} deliberately does not
+     * expose remove: {@link TrieCache} is a prefix tree and has no matching semantics for it.
+     *
+     * <p>Weak keys alone are not enough here. The captured snapshot is the map <em>value</em>,
+     * so it stays strongly reachable from this static field: it survives every young GC and is
+     * promoted to the old generation. And an entry only leaves the map when a later put/get
+     * happens to drain the reference queue, so entries linger once traffic stops.
+     */
+    WeakCache<Object, Object> CAPTURED_CACHE = new WeakCache<>();
 
     static <K, V> Cache<K, V> weakMap() {
         return new WeakCache<>();

arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/WeakCache.java

@@ -41,6 +41,18 @@ public void put(K key, V value) {
         target.put(new WeakReferenceKey<>(key, this), value);
     }
 
+    /**
+     * Drops the entry for {@code key} and returns its previous value, or null if there was none.
+     *
+     * <p>Weak keys are still collected on their own, but an explicit remove lets a caller that
+     * knows the key is done reclaim the entry immediately instead of waiting for a GC cycle
+     * followed by whatever put/get happens to drain the queue next.
+     */
+    public V remove(K key) {
+        check();
+        return target.remove(new WeakReferenceKey<>(key));
+    }
+
     public void clear() {
         target.clear();
     }

arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/internal/WeakCacheTest.java

@@ -58,6 +58,38 @@ void testNormalKeyValue() throws InterruptedException {
         assertFalse(Cache.CAPTURED_CACHE.contains(null));
     }
 
+    /**
+     * Explicit remove: the ForkJoinTask exec/run exit uses it to hand the snapshot back right
+     * away instead of waiting for a GC cycle. Also pins the identity semantics -- remove must
+     * only drop the entry for that same object, never one that merely compares equal.
+     */
+    @Test
+    void testRemove() {
+        WeakCache<Object, String> cache = new WeakCache<>();
+        Object key = new Object();
+        cache.put(key, "value");
+        assertTrue(cache.contains(key));
+
+        assertEquals("value", cache.remove(key));
+        assertFalse(cache.contains(key));
+        assertNull(cache.remove(key));
+    }
+
+    @Test
+    void testRemoveIsIdentityBasedNotEquals() {
+        WeakCache<Object, String> cache = new WeakCache<>();
+        String first = new StringBuilder("same").toString();
+        String second = new StringBuilder("same").toString();
+        cache.put(first, "first-value");
+        cache.put(second, "second-value");
+
+        assertEquals("first-value", cache.remove(first));
+
+        assertFalse(cache.contains(first));
+        assertTrue(cache.contains(second));
+        assertEquals("second-value", cache.get(second));
+    }
+
     @Test
     void testWeakReferenceKeyEqualsReturnsFalse() {
         WeakCache.WeakReferenceKey<String> key = new WeakCache.WeakReferenceKey<>("test", new ReferenceQueue<>());

arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskInstrumentation.java

@@ -13,6 +13,33 @@
 
 import static net.bytebuddy.matcher.ElementMatchers.*;
 
+/**
+ * Replays the context captured when the task was constructed, and drops the captured entry
+ * once the task has run.
+ *
+ * <p>Removing at the exec/run exit is safe. For {@code CompletableFuture} nodes,
+ * {@code Completion.run()} and {@code exec()} are final methods on the base class and both are
+ * just {@code tryFire(ASYNC)}; the ASYNC branch skips {@code claim()}, runs the function, and
+ * then nulls out {@code src}, {@code dep} and {@code fn}, so it is that node's one terminal
+ * fire and any later {@code tryFire} returns at the null check on entry. Same shape on JDK 8
+ * and 21. For plain ForkJoinTasks, {@code doExec()} calls {@code exec()} once.
+ *
+ * <p>What this does <em>not</em> cover: the constructor advice captures for every ForkJoinTask,
+ * but only tasks that reach exec/run get their entry dropped explicitly. Completions built by
+ * the non-async operators ({@code thenApply}, {@code thenCompose}, ...) have a null executor and
+ * are only ever driven through {@code tryFire(SYNC)} / {@code tryFire(NESTED)};
+ * {@code CompletableFuture$Signaller} is captured too but is never submitted to the pool. Those
+ * still wait for the weak key to be collected. There is no safe removal signal inside
+ * {@code tryFire}: a null return means both "spun without firing" and "fired with no dependent
+ * to propagate", and {@code Completion.isLive()} is package-private so inlined advice cannot
+ * reach it. Binding the snapshot to a field on the task itself would remove the map entirely
+ * and cover those cases, at the cost of adding a field-injection mechanism.
+ *
+ * <p>Known behaviour change: after {@code ForkJoinTask.reinitialize()} the same instance can be
+ * forked again, and the second execution no longer replays. It previously replayed the snapshot
+ * taken at construction time, which was already the wrong context for a fresh run. The JDK never
+ * calls {@code reinitialize()} itself.
+ */
 public class ForkJoinTaskInstrumentation extends TypeInstrumentation {
 
     @Override
@@ -37,9 +64,28 @@ public static void onEnter(
             backup = ArexThreadLocal.Transmitter.replay(captured);
         }
 
-        @Advice.OnMethodExit(suppress = Throwable.class)
-        public static void onExit(@Advice.Local("backup") Object backup) {
+        /**
+         * {@code onThrowable} is required, not optional. Throwing out of exec is a normal path,
+         * not a corner case: {@code RunnableExecuteAction.exec()} is a bare {@code runnable.run()},
+         * {@code AdaptedCallable.exec()} rethrows explicitly, and {@code AdaptedXxx.run()} is
+         * {@code invoke()}, which reports the exception by rethrowing it -- so the run and exec
+         * exits are both skipped. Without it the enter advice has already replayed a context onto
+         * this ForkJoin worker and nothing restores it, so the worker picks up the next task still
+         * carrying the previous trace and its downstream calls are recorded under a foreign
+         * traceId. The remove below would be skipped on the same paths.
+         */
+        @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+        public static void onExit(
+                @Advice.This Object task,
+                @Advice.Local("backup") Object backup) {
             ArexThreadLocal.Transmitter.restore(backup);
+            // A non-null backup means captured was non-null on entry, because replay(null)
+            // always returns null. Use it to skip the map lookup when nothing was recorded:
+            // the constructor advice is gated on TraceContextManager and never put an entry in
+            // that case, and exec is one of the hottest methods in the JVM.
+            if (backup != null) {
+                Cache.CAPTURED_CACHE.remove(task);
+            }
         }
     }
 }

arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskInstrumentationTest.java

@@ -2,12 +2,18 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 import io.arex.agent.bootstrap.ctx.ArexThreadLocal;
+import io.arex.agent.bootstrap.internal.Cache;
 import io.arex.inst.executors.ForkJoinTaskInstrumentation.ExecAdvice;
+import java.lang.reflect.Method;
 import java.util.concurrent.CountedCompleter;
 import java.util.concurrent.ForkJoinTask;
+import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -23,6 +29,7 @@ void setUp() {
 
     @AfterEach
     void tearDown() {
+        Cache.CAPTURED_CACHE.clear();
     }
 
     @Test
@@ -44,11 +51,80 @@ void ExecAdvice_onEnter() {
 
     @Test
     void ExecAdvice_onExit() {
-        assertDoesNotThrow(() -> ExecAdvice.onExit(ArexThreadLocal.Transmitter.capture()));
+        assertDoesNotThrow(() -> ExecAdvice.onExit("fork-test", ArexThreadLocal.Transmitter.capture()));
     }
 
     @Test
     void ConstructorAdvice_onEnter() {
         assertDoesNotThrow(() -> ForkJoinTaskConstructorInstrumentation.ConstructorAdvice.onExit(new Object()));
     }
+
+    /**
+     * The exit advice must drop the captured entry. Leaving it behind is the bug: the snapshot is
+     * the map value, so it stays strongly reachable from the static cache and is promoted to the
+     * old generation, and nothing evicts it until a later put/get drains the reference queue.
+     */
+    @Test
+    void execAdviceRemovesCapturedEntryOnExit() {
+        ArexThreadLocal<String> threadLocal = new ArexThreadLocal<>();
+        try {
+            Object task = new Object();
+            Cache.CAPTURED_CACHE.put(task, "snapshot");
+            assertTrue(Cache.CAPTURED_CACHE.contains(task));
+
+            ExecAdvice.onExit(task, backupOf(threadLocal));
+
+            assertFalse(Cache.CAPTURED_CACHE.contains(task));
+        } finally {
+            threadLocal.remove();
+        }
+    }
+
+    /**
+     * A null backup means captured was null on entry, i.e. nothing was being recorded and the
+     * constructor advice never put an entry. Touching the map there would cost a hash and a bin
+     * walk on one of the hottest methods in the JVM for nothing. Dropping the guard fails this.
+     */
+    @Test
+    void execAdviceSkipsRemoveWhenNothingWasReplayed() {
+        Object task = new Object();
+        Cache.CAPTURED_CACHE.put(task, "snapshot");
+
+        ExecAdvice.onExit(task, null);
+
+        assertTrue(Cache.CAPTURED_CACHE.contains(task));
+    }
+
+    /**
+     * Throwing out of exec is a normal path: RunnableExecuteAction.exec() is a bare
+     * runnable.run(). Without onThrowable the whole exit advice is skipped, the replayed context
+     * is never restored, and the worker carries the previous trace into its next task.
+     */
+    @Test
+    void exitAdviceAlsoRunsOnTheExceptionPath() {
+        for (Method method : ExecAdvice.class.getDeclaredMethods()) {
+            Advice.OnMethodExit exit = method.getAnnotation(Advice.OnMethodExit.class);
+            if (exit != null) {
+                assertEquals(Throwable.class, exit.onThrowable(),
+                        "ExecAdvice.onExit is missing onThrowable: when the instrumented method "
+                                + "throws, the exit advice does not run and the replayed context "
+                                + "leaks onto the worker thread.");
+                return;
+            }
+        }
+        fail("ExecAdvice has no @Advice.OnMethodExit method");
+    }
+
+    /**
+     * Builds a non-null backup. Transmitter.capture() returns null when no ArexThreadLocal is
+     * registered in the holder, so one has to be set first -- otherwise the test would silently
+     * slide onto the "nothing recorded" branch and assert nothing. restore(capture()) is an
+     * identity operation for the current thread, so this does not disturb other tests.
+     */
+    private static Object backupOf(ArexThreadLocal<String> threadLocal) {
+        threadLocal.set("value");
+        Object backup = ArexThreadLocal.Transmitter.capture();
+        assertNotNull(backup);
+        return backup;
+    }
 }