From db77df51cb2ca74d07d0a814a517ca50fd5ce725 Mon Sep 17 00:00:00 2001 From: zq Date: Wed, 2 Sep 2026 20:41:18 +0800 Subject: [PATCH] fix: ForkJoinTask captured context is retained until GC (#626) The constructor advice captures a context snapshot for every ForkJoinTask, but the exec/run exit only restores and never removes, so eviction depends entirely on the task being collected, the weak key being enqueued, and some later put/get happening to drain the queue. That is not a leak in the strict sense -- the snapshot does not reference the task back, so the weak key is still collectible. It behaves like one, though: the snapshot is the map value and stays strongly reachable from the static cache, so it survives every young GC and is promoted to the old generation, and entries linger once traffic stops because nothing calls check() any more. Worst on JDK 21 virtual threads. VirtualThread.runContinuation is a plain Runnable, so ForkJoinPool.execute(Runnable) allocates a new RunnableExecuteAction on every submit -- one entry per park/unpark, and a full Transmitter.capture() (a fresh HashMap plus deep-copied CallDepth values) on each one. Changes: - Drop the entry at the exec/run exit. This is that task's terminal execution: Completion.run()/exec() are final on the base class and are both tryFire(ASYNC), and the ASYNC branch skips claim(), runs the function and then nulls src/dep/fn, so any later tryFire returns at the entry null check. Same on JDK 8 and 21. - Guard the removal on backup != null, which holds exactly when captured was non-null on entry because replay(null) returns null. Nothing was recorded otherwise, the constructor advice is gated on TraceContextManager and never put an entry, and exec is one of the hottest methods in the JVM. - Add onThrowable to the exit advice. Throwing out of exec is a normal path: RunnableExecuteAction.exec() is a bare runnable.run(), AdaptedCallable.exec() rethrows, and AdaptedXxx.run() is invoke(), which reports by rethrowing, so both the run and exec exits are skipped. Without it the replayed context is never restored and the worker picks up its next task still carrying the previous trace, recording downstream calls under a foreign traceId. That is a pre-existing defect; the new removal would be skipped on the same paths. - CAPTURED_CACHE is declared as WeakCache so the advice can call remove. Cache itself does not get a remove method: TrieCache is a prefix tree and has no matching semantics, and CAPTURED_CACHE is the only caller. Not covered, and called out in the class javadoc: only tasks that reach exec/run get their entry dropped. Completions built by the non-async operators have a null executor and are only driven through tryFire(SYNC)/tryFire(NESTED), and CompletableFuture$Signaller is captured but never submitted to the pool; both still wait for the weak key. There is no safe removal signal inside tryFire -- a null return means both "spun without firing" and "fired with no dependent to propagate", and Completion.isLive() is package-private so inlined advice cannot reach it. Binding the snapshot to a field on the task would remove the map altogether and cover those too, at the cost of a field-injection mechanism. Behaviour change: after reinitialize() the same instance can be forked again and the second execution no longer replays. It previously replayed the snapshot taken at construction, which was already the wrong context. The JDK never calls reinitialize() itself. Verified: arex-agent-bootstrap 138/138, arex-executors 23/23. Mutation checked -- dropping the backup guard fails execAdviceSkipsRemoveWhenNothingWasReplayed, deleting the remove fails execAdviceRemovesCapturedEntryOnExit, and dropping onThrowable fails exitAdviceAlsoRunsOnTheExceptionPath. --- .../arex/agent/bootstrap/internal/Cache.java | 12 ++- .../agent/bootstrap/internal/WeakCache.java | 12 +++ .../bootstrap/internal/WeakCacheTest.java | 32 ++++++++ .../ForkJoinTaskInstrumentation.java | 50 +++++++++++- .../ForkJoinTaskInstrumentationTest.java | 78 ++++++++++++++++++- 5 files changed, 180 insertions(+), 4 deletions(-) diff --git a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/Cache.java b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/Cache.java index a38d28ae7..c9def11b2 100644 --- a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/Cache.java +++ b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/Cache.java @@ -2,7 +2,17 @@ public interface Cache { - Cache 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. + * + *

Weak keys alone are not enough here. The captured snapshot is the map value, + * 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 CAPTURED_CACHE = new WeakCache<>(); static Cache weakMap() { return new WeakCache<>(); diff --git a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/WeakCache.java b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/WeakCache.java index 4e9df7a12..ea14201db 100644 --- a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/WeakCache.java +++ b/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. + * + *

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(); } diff --git a/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/internal/WeakCacheTest.java b/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/internal/WeakCacheTest.java index 8fe070dbd..91503151a 100644 --- a/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/internal/WeakCacheTest.java +++ b/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 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 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 key = new WeakCache.WeakReferenceKey<>("test", new ReferenceQueue<>()); diff --git a/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskInstrumentation.java b/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskInstrumentation.java index 571c499ff..8d868f024 100644 --- a/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskInstrumentation.java +++ b/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. + * + *

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. + * + *

What this does not 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. + * + *

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); + } } } } diff --git a/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskInstrumentationTest.java b/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskInstrumentationTest.java index 9086efa32..df272e86f 100644 --- a/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskInstrumentationTest.java +++ b/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 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 threadLocal) { + threadLocal.set("value"); + Object backup = ArexThreadLocal.Transmitter.capture(); + assertNotNull(backup); + return backup; + } }