diff --git a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/util/VirtualThreadUtil.java b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/util/VirtualThreadUtil.java new file mode 100644 index 000000000..f55770d52 --- /dev/null +++ b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/util/VirtualThreadUtil.java @@ -0,0 +1,47 @@ +package io.arex.agent.bootstrap.util; + +/** + * Loom(virtual thread) related helper. + * + *

Compiled with java 8, so neither {@code Thread.isVirtual()} nor + * {@code jdk.internal.misc.CarrierThread} can be referenced directly. + */ +public final class VirtualThreadUtil { + + /** + * The thread type of the default virtual thread scheduler's worker, + * see {@code java.lang.VirtualThread#createDefaultScheduler()}. + */ + private static final String CARRIER_THREAD = "jdk.internal.misc.CarrierThread"; + + private VirtualThreadUtil() { + } + + /** + * Identifies a default-scheduler worker while it is running as a platform + * thread. While a virtual thread is mounted, {@link Thread#currentThread()} + * returns the virtual thread, so this check does not skip its business tasks. + * + *

ForkJoinTask advice runs outside the virtual thread's continuation and + * must not transmit context to the carrier's separate thread locals. It must + * also avoid polling the shared captured-context cache's + * {@link java.lang.ref.ReferenceQueue}: on JDK 21, a virtual thread can be the + * first waiter for that queue's lock. If all carriers then block on the same + * lock, the virtual thread cannot resume even after the lock is released. + * + *

This identifies the JDK's default scheduler only, not custom schedulers. + * + * @return true if the current thread is a carrier thread of the default scheduler + */ + public static boolean isCarrierThread() { + return isCarrierThread(Thread.currentThread()); + } + + public static boolean isCarrierThread(Thread thread) { + return thread != null && isCarrierThread(thread.getClass().getName()); + } + + static boolean isCarrierThread(String threadClassName) { + return CARRIER_THREAD.equals(threadClassName); + } +} diff --git a/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/util/VirtualThreadUtilTest.java b/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/util/VirtualThreadUtilTest.java new file mode 100644 index 000000000..8d5fe80f3 --- /dev/null +++ b/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/util/VirtualThreadUtilTest.java @@ -0,0 +1,19 @@ +package io.arex.agent.bootstrap.util; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class VirtualThreadUtilTest { + + @Test + void isCarrierThread() { + // current thread of the unit test is a normal platform thread + assertFalse(VirtualThreadUtil.isCarrierThread()); + assertFalse(VirtualThreadUtil.isCarrierThread((Thread) null)); + assertFalse(VirtualThreadUtil.isCarrierThread(Thread.currentThread())); + assertFalse(VirtualThreadUtil.isCarrierThread("java.lang.Thread")); + assertTrue(VirtualThreadUtil.isCarrierThread("jdk.internal.misc.CarrierThread")); + } +} diff --git a/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentation.java b/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentation.java index 8144226b3..1786b1489 100644 --- a/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentation.java +++ b/arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentation.java @@ -3,6 +3,7 @@ import io.arex.agent.bootstrap.TraceContextManager; import io.arex.agent.bootstrap.ctx.ArexThreadLocal; import io.arex.agent.bootstrap.internal.Cache; +import io.arex.agent.bootstrap.util.VirtualThreadUtil; import io.arex.inst.extension.MethodInstrumentation; import io.arex.inst.extension.TypeInstrumentation; import net.bytebuddy.asm.Advice; @@ -32,6 +33,10 @@ public List methodAdvices() { public static class ConstructorAdvice { @Advice.OnMethodExit(suppress = Throwable.class) public static void onExit(@Advice.This Object task) { + // Do not capture carrier context or poll the shared reference queue while scheduling a VT. + if (VirtualThreadUtil.isCarrierThread()) { + return; + } if (TraceContextManager.get() == null) { return; } 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..8961e592c 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 @@ -2,6 +2,7 @@ import io.arex.agent.bootstrap.ctx.ArexThreadLocal; import io.arex.agent.bootstrap.internal.Cache; +import io.arex.agent.bootstrap.util.VirtualThreadUtil; import io.arex.inst.extension.MethodInstrumentation; import io.arex.inst.extension.TypeInstrumentation; import net.bytebuddy.asm.Advice; @@ -33,6 +34,10 @@ public static class ExecAdvice { public static void onEnter( @Advice.This Object task, @Advice.Local("backup") Object backup) { + // A carrier must stay available to resume VTs waiting on the shared reference queue's lock. + if (VirtualThreadUtil.isCarrierThread()) { + return; + } final Object captured = Cache.CAPTURED_CACHE.get(task); backup = ArexThreadLocal.Transmitter.replay(captured); } diff --git a/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinCarrierThreadTest.java b/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinCarrierThreadTest.java new file mode 100644 index 000000000..88028074f --- /dev/null +++ b/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinCarrierThreadTest.java @@ -0,0 +1,265 @@ +package io.arex.inst.executors; + +import io.arex.agent.bootstrap.TraceContextManager; +import io.arex.agent.bootstrap.ctx.ArexThreadLocal; +import io.arex.agent.bootstrap.internal.Cache; +import io.arex.agent.bootstrap.util.VirtualThreadUtil; +import java.io.File; +import java.lang.instrument.Instrumentation; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.lang.reflect.Field; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ForkJoinTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.BooleanSupplier; +import java.util.jar.Attributes; +import java.util.jar.JarOutputStream; +import java.util.jar.Manifest; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.asm.Advice; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import static net.bytebuddy.matcher.ElementMatchers.*; +import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +/** + * Uses separate JVMs because the reproducer deliberately starves every carrier. + * Reflection keeps the test sources compatible with the project's Java 8 target. + */ +class ForkJoinCarrierThreadTest { + @TempDir + Path directory; + + @Test + void reproducesStarvationWithoutCarrierGuard() throws Exception { + runScenario("before"); + } + + @Test + void carrierGuardsAllowVirtualThreadToResume() throws Exception { + runScenario("after"); + } + + private void runScenario(String mode) throws Exception { + assumeTrue("21".equals(System.getProperty("java.specification.version")), + "This reproducer targets JDK 21's ReentrantLock-based ReferenceQueue"); + assumeTrue(ReferenceQueue.class.getDeclaredField("lock").getType() == ReentrantLock.class); + Manifest manifest = new Manifest(); + Attributes attributes = manifest.getMainAttributes(); + attributes.putValue("Manifest-Version", "1.0"); + attributes.putValue("Premain-Class", TestAgent.class.getName()); + attributes.putValue("Can-Retransform-Classes", "true"); + Path agent = directory.resolve("carrier-test-agent.jar"); + // The agent class and its dependencies are supplied by the test classpath. + try (JarOutputStream ignored = new JarOutputStream(Files.newOutputStream(agent), manifest)) { + } + String bootstrap = new File(Cache.class.getProtectionDomain().getCodeSource().getLocation().toURI()).getPath(); + String classpath = System.getProperty("surefire.test.class.path", System.getProperty("java.class.path")); + Path log = directory.resolve(mode + ".log"); + Process process = new ProcessBuilder( + new File(System.getProperty("java.home"), "bin/java").getPath(), + "-Djdk.virtualThreadScheduler.parallelism=2", + "-Djdk.virtualThreadScheduler.maxPoolSize=2", + "--add-opens=java.base/java.lang=ALL-UNNAMED", + "--add-opens=java.base/java.lang.ref=ALL-UNNAMED", + "-Xbootclasspath/a:" + bootstrap, + "-javaagent:" + agent + "=" + mode, + "-cp", classpath, Scenario.class.getName(), mode) + .redirectErrorStream(true).redirectOutput(log.toFile()).start(); + try { + boolean finished = process.waitFor(20, TimeUnit.SECONDS); + String output = new String(Files.readAllBytes(log), StandardCharsets.UTF_8); + System.out.println(output); + assertTrue(finished, "Child JVM timed out:\n" + output); + assertEquals(0, process.exitValue(), output); + assertTrue(output.contains("RESULT " + mode + " PASS"), output); + } finally { + if (process.isAlive()) { + process.destroyForcibly(); + process.waitFor(5, TimeUnit.SECONDS); + } + } + } + + public static class TestAgent { + public static void premain(String mode, Instrumentation instrumentation) { + Class execution = "before".equals(mode) ? BeforeExec.class + : ForkJoinTaskInstrumentation.ExecAdvice.class; + Class constructor = "before".equals(mode) ? BeforeConstructor.class + : ForkJoinTaskConstructorInstrumentation.ConstructorAdvice.class; + new AgentBuilder.Default() + .disableClassFormatChanges() + .ignore(nameStartsWith("net.bytebuddy.").or(nameStartsWith("io.arex."))) + .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) + .with(AgentBuilder.Listener.StreamWriting.toSystemError().withTransformationsOnly()) + .type(new ForkJoinTaskInstrumentation().typeMatcher()) + .transform((builder, type, loader, module, domain) -> { + if (new ForkJoinTaskConstructorInstrumentation().typeMatcher().matches(type)) { + builder = builder.visit(Advice.to(constructor).on(isConstructor())); + } + return builder.visit(Advice.to(execution).on( + new ForkJoinTaskInstrumentation().methodAdvices().get(0).getMethodMatcher())); + }) + .assureReadEdgeFromAndTo(instrumentation, Cache.class) + .installOn(instrumentation); + } + } + + // The unguarded advice is a positive control: the same schedule must reproduce the old bug. + public static class BeforeExec { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.This Object task, @Advice.Local("backup") Object backup) { + backup = ArexThreadLocal.Transmitter.replay(Cache.CAPTURED_CACHE.get(task)); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit(@Advice.Local("backup") Object backup) { + ArexThreadLocal.Transmitter.restore(backup); + } + } + + public static class BeforeConstructor { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit(@Advice.This Object task) { + if (TraceContextManager.get() == null) { + return; + } + Object captured = ArexThreadLocal.Transmitter.capture(); + if (captured != null) { + Cache.CAPTURED_CACHE.put(task, captured); + } + } + } + + public static class Scenario { + private static volatile boolean enterQueue; + + public static void main(String[] args) throws Exception { + String mode = args[0]; + startVirtualThread(() -> {}).join(); + // Load the adapted-task class before arranging lock contention. + ForkJoinTask.adapt(() -> {}); + Object key = new Object(); + Cache.CAPTURED_CACHE.put(key, "cached-value"); + @SuppressWarnings("unchecked") + ReferenceQueue queue = (ReferenceQueue) Cache.CAPTURED_CACHE; + Field lockField = ReferenceQueue.class.getDeclaredField("lock"); + lockField.setAccessible(true); + ReentrantLock lock = (ReentrantLock) lockField.get(queue); + Field schedulerField = Class.forName("java.lang.VirtualThread").getDeclaredField("DEFAULT_SCHEDULER"); + schedulerField.setAccessible(true); + ForkJoinPool scheduler = (ForkJoinPool) schedulerField.get(null); + CountDownLatch running = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(1); + CountDownLatch probes = new CountDownLatch(2); + Thread vt = startVirtualThread(() -> { + check(!VirtualThreadUtil.isCarrierThread(), "VT must not be mistaken for a carrier"); + TraceContextManager.set("vt-context"); + running.countDown(); + while (!enterQueue) { + // Stay mounted until the main thread holds the ReferenceQueue lock. + } + check("cached-value".equals(Cache.CAPTURED_CACHE.get(key)), "cached value was lost"); + check("vt-context".equals(TraceContextManager.get()), "VT context was lost"); + done.countDown(); + }); + check(running.await(5, TimeUnit.SECONDS), "VT did not start"); + lock.lock(); + try { + // An empty ReferenceQueue.poll() never takes the lock. Explicit enqueue avoids GC races. + new WeakReference<>(new Object(), queue).enqueue(); + enterQueue = true; + await(() -> lock.hasQueuedThread(vt) && vt.getState() == Thread.State.WAITING, + "VT must be the first parked lock waiter"); + Runnable probe = () -> { + check(VirtualThreadUtil.isCarrierThread(), "probe must run as a real carrier"); + // Exercise the constructor guard with non-null carrier context as well. + TraceContextManager.set("carrier-context"); + try { + ForkJoinTask.adapt(() -> {}); + } finally { + TraceContextManager.remove(); + } + probes.countDown(); + }; + scheduler.execute(probe); + scheduler.execute(probe); + if ("before".equals(mode)) { + await(() -> lock.getQueueLength() == 3, "VT followed by two blocked carriers"); + } else { + check(probes.await(5, TimeUnit.SECONDS), "carrier advice blocked on the cache"); + check(lock.getQueueLength() == 1, "only the VT should wait for the lock"); + } + System.out.println("mode=" + mode + " before unlock: queueLength=" + lock.getQueueLength() + + " vtQueued=" + lock.hasQueuedThread(vt) + " poolSize=" + scheduler.getPoolSize() + + " probesRemaining=" + probes.getCount()); + } finally { + lock.unlock(); + } + boolean progressed = done.await(2, TimeUnit.SECONDS); + System.out.println("after unlock: progressed=" + progressed + " lockHeld=" + lock.isLocked() + + " vtState=" + vt.getState() + " queueLength=" + lock.getQueueLength() + + " probesRemaining=" + probes.getCount()); + if ("before".equals(mode)) { + check(!progressed && !lock.isLocked() && lock.getQueueLength() == 3 + && vt.getState() == Thread.State.RUNNABLE && probes.getCount() == 2, + "expected an unlocked queue with a runnable VT and both carriers blocked"); + Thread.getAllStackTraces().forEach((thread, stack) -> { + if (VirtualThreadUtil.isCarrierThread(thread)) { + System.out.println(thread); + for (StackTraceElement frame : stack) { + System.out.println(" " + frame); + } + } + }); + } else { + check(progressed, "VT did not resume after unlock"); + check(lock.getQueueLength() == 0, "queue did not drain"); + verifyOrdinaryForkJoinContext(); + } + System.out.println("RESULT " + mode + " PASS"); + // The before case deliberately leaves the scheduler's daemon carriers blocked. + } + + private static Thread startVirtualThread(Runnable task) throws Exception { + return (Thread) Thread.class.getMethod("startVirtualThread", Runnable.class).invoke(null, task); + } + + private static void verifyOrdinaryForkJoinContext() throws Exception { + ForkJoinPool pool = new ForkJoinPool(1); + TraceContextManager.set("parent-context"); + try { + String captured = pool.submit(() -> { + check(!VirtualThreadUtil.isCarrierThread(), "ordinary FJP worker mistaken for carrier"); + return TraceContextManager.get(); + }).get(5, TimeUnit.SECONDS); + check("parent-context".equals(captured), "ordinary FJP context propagation regressed"); + } finally { + TraceContextManager.remove(); + pool.shutdownNow(); + } + } + + private static void await(BooleanSupplier condition, String message) throws Exception { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!condition.getAsBoolean()) { + check(System.nanoTime() < deadline, message); + Thread.sleep(1); + } + } + + private static void check(boolean condition, String message) { + if (!condition) { + throw new AssertionError(message); + } + } + } +} diff --git a/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentationTest.java b/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentationTest.java index 706dda2e3..1a7a6163b 100644 --- a/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentationTest.java +++ b/arex-instrumentation/internal/arex-executors/src/test/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentationTest.java @@ -1,10 +1,14 @@ package io.arex.inst.executors; +import io.arex.agent.bootstrap.TraceContextManager; import io.arex.agent.bootstrap.ctx.ArexThreadLocal; +import io.arex.agent.bootstrap.internal.Cache; +import io.arex.agent.bootstrap.util.VirtualThreadUtil; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.MockedStatic; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; @@ -42,4 +46,18 @@ void onEnter() { Mockito.when(ArexThreadLocal.Transmitter.capture()).thenReturn("mock"); assertDoesNotThrow(() -> ForkJoinTaskConstructorInstrumentation.ConstructorAdvice.onExit(null)); } -} \ No newline at end of file + + @Test + void onExit_skipCarrierThread() { + try (MockedStatic mocked = Mockito.mockStatic(VirtualThreadUtil.class); + MockedStatic traces = Mockito.mockStatic(TraceContextManager.class)) { + mocked.when(VirtualThreadUtil::isCarrierThread).thenReturn(true); + traces.when(TraceContextManager::get).thenReturn("carrier-context"); + Object task = new Object(); + assertDoesNotThrow(() -> ForkJoinTaskConstructorInstrumentation.ConstructorAdvice.onExit(task)); + traces.verifyNoInteractions(); + // must not touch the weak cache(ReferenceQueue) on a carrier thread + assertFalse(Cache.CAPTURED_CACHE.contains(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..a91707bc1 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 @@ -5,6 +5,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import io.arex.agent.bootstrap.ctx.ArexThreadLocal; +import io.arex.agent.bootstrap.util.VirtualThreadUtil; import io.arex.inst.executors.ForkJoinTaskInstrumentation.ExecAdvice; import java.util.concurrent.CountedCompleter; import java.util.concurrent.ForkJoinTask; @@ -12,6 +13,8 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; class ForkJoinTaskInstrumentationTest { @@ -42,6 +45,17 @@ void ExecAdvice_onEnter() { assertDoesNotThrow(() -> ExecAdvice.onEnter("fork-test", new Object())); } + @Test + void ExecAdvice_onEnter_skipCarrierThread() { + try (MockedStatic mocked = Mockito.mockStatic(VirtualThreadUtil.class); + MockedStatic transmitter = Mockito.mockStatic(ArexThreadLocal.Transmitter.class)) { + mocked.when(VirtualThreadUtil::isCarrierThread).thenReturn(true); + assertDoesNotThrow(() -> ExecAdvice.onEnter("fork-test", new Object())); + // must not touch thread locals(nor the weak cache/ReferenceQueue) on a carrier thread + transmitter.verify(() -> ArexThreadLocal.Transmitter.replay(Mockito.any()), Mockito.never()); + } + } + @Test void ExecAdvice_onExit() { assertDoesNotThrow(() -> ExecAdvice.onExit(ArexThreadLocal.Transmitter.capture()));