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; + } }