diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpcProcessManager.java b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpcProcessManager.java index f3e121af473..a09e2ecd3b0 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpcProcessManager.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpcProcessManager.java @@ -17,14 +17,45 @@ import org.jspecify.annotations.Nullable; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Supplier; /** * Manages the lifecycle of {@link RewriteRpc} instances, allowing for a single * Rewrite RPC subprocess per thread. + *

+ * Each thread that does RPC work (parsing or printing a Python/JavaScript/… LST) + * lazily gets its own subprocess via {@link #getOrStart()}, which is reused for + * the life of that thread so a thread never spawns more than one server. Because + * each server is an out-of-process node process with no finalizer, it must be + * torn down explicitly — letting the owning {@link RewriteRpc} become garbage + * does not stop the OS process (it lingers until the JVM-exit shutdown + * hook on {@link RewriteRpcProcess} fires). + *

+ * The hazard this guards against is orphaned servers: a server started + * on thread X cannot be reached by {@link #shutdown()} called on thread Y, so + * any server whose owning thread terminates before calling {@code shutdown()} + * (ForkJoinPool ManagedBlocker compensation threads, cached/elastic pool threads + * that die on their idle timeout, …) would otherwise survive — and accumulate + * without bound — until JVM exit. To keep the number of live servers in check, + * this manager tracks every started server in a process-wide, thread-keyed + * registry and {@linkplain #reapDeadThreads() reaps} servers whose owning thread + * is no longer alive. Reaping is activity-driven: it runs when a new RPC thread + * starts a server (in {@link #getOrStart()}) and on {@link #shutdown()}, so a + * dead thread's server lingers only until the next such event — the residual is + * bounded by the number of RPC threads that have retired since the last reap, + * not by the total number of threads ever created. Reaping a dead thread's + * server is always safe because a dead thread cannot be mid-RPC, which is what + * makes it sound to do even while other threads are running concurrent RPC work. */ public class RewriteRpcProcessManager { - private final ThreadLocal rpc = new ThreadLocal<>(); + /** + * Process-wide registry of live servers keyed by their owning thread. Replaces a + * bare {@link ThreadLocal} so that, in addition to the per-thread fast path, the + * full set of live servers can be enumerated for reaping and definitive shutdown. + */ + private final Map rpcByThread = new ConcurrentHashMap<>(); private final ThreadLocal> factory; public RewriteRpcProcessManager(Supplier defaultFactory) { @@ -32,17 +63,31 @@ public RewriteRpcProcessManager(Supplier defaultFactory) { } public @Nullable R get() { - return rpc.get(); + return rpcByThread.get(Thread.currentThread()); } public R getOrStart() { - R current = rpc.get(); - //noinspection ConstantValue - if (current == null) { - current = factory.get().get(); - rpc.set(current); + Thread current = Thread.currentThread(); + R existing = rpcByThread.get(current); + if (existing != null) { + // Fast path: this thread already owns a server. Avoid the reap sweep so + // hot parse/print paths don't pay for it on every call. + return existing; + } + // A new RPC thread is appearing — exactly when orphaned servers from + // previously-retired threads should be cleaned up before adding another. + reapDeadThreads(); + R started = factory.get().get(); + // Construct the server before touching the map. Only the current thread ever + // inserts its own key, so putIfAbsent's loser branch can fire only if this + // thread re-entered getOrStart() during construction; if it did, shut the + // duplicate down and keep the already-registered server. + R prior = rpcByThread.putIfAbsent(current, started); + if (prior != null) { + started.shutdown(); + return prior; } - return current; + return started; } public void setFactory(Supplier factory) { @@ -50,19 +95,83 @@ public void setFactory(Supplier factory) { } public void reset() { - R current = rpc.get(); - //noinspection ConstantValue + R current = get(); if (current != null) { current.reset(); } } + /** + * Shut down the calling thread's server (if any) and reap any servers whose + * owning thread has since died. Callers already invoke this at run boundaries, + * so routing dead-thread reaping through it keeps orphaned servers from + * accumulating across runs without requiring any change at the call sites. + */ public void shutdown() { - R current = rpc.get(); - //noinspection ConstantValue - if (current != null) { - current.shutdown(); - rpc.remove(); + R current = rpcByThread.remove(Thread.currentThread()); + try { + if (current != null) { + current.shutdown(); + } + } finally { + reapDeadThreads(); + } + } + + /** + * Shut down every live server across all threads, driving the live-server + * count to zero. Unlike {@link #reapDeadThreads()} this also stops servers owned by + * currently-live threads, so it must only be called once no thread is doing — or + * about to do — RPC work, e.g. on JVM/service shutdown. A server is registered only + * after it is constructed in {@link #getOrStart()}, so calling this concurrently with + * {@code getOrStart()} could let a just-created server register after the sweep and + * survive; calling it mid parse/print would tear that run's server out from under it. + *

+ * This is not required for correctness on JVM exit: each RPC subprocess registers its + * own JVM-exit shutdown hook, so {@code shutdownAll()} only bounds the live-server + * count earlier within a long-running JVM. + */ + public void shutdownAll() { + for (Thread t : rpcByThread.keySet()) { + R r = rpcByThread.remove(t); + if (r != null) { + try { + r.shutdown(); + } catch (Throwable ignored) { + // Keep going so one bad teardown can't strand the rest. + } + } + } + } + + /** + * The number of live servers currently tracked by this manager. Intended for + * diagnostics and tests asserting that the count stays bounded. + */ + public int liveCount() { + return rpcByThread.size(); + } + + /** + * Shut down servers whose owning thread is no longer alive. Safe to call at any + * time, including while other threads are doing RPC work, because a terminated + * thread cannot be mid-RPC on its server. + */ + private void reapDeadThreads() { + for (Map.Entry e : rpcByThread.entrySet()) { + if (!e.getKey().isAlive()) { + // remove() returns the value to exactly one racing reaper, so the + // server is shut down once even under concurrent reaping. + R dead = rpcByThread.remove(e.getKey()); + if (dead != null) { + try { + dead.shutdown(); + } catch (Throwable ignored) { + // Best-effort: the owning thread is gone, so there is no one + // to surface this to; keep reaping the remaining orphans. + } + } + } } } } diff --git a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcProcessManagerTest.java b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcProcessManagerTest.java new file mode 100644 index 00000000000..b256aef69a4 --- /dev/null +++ b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcProcessManagerTest.java @@ -0,0 +1,311 @@ +/* + * Copyright 2025 the original author or authors. + *

+ * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + *

+ * https://www.apache.org/licenses/LICENSE-2.0 + *

+ * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openrewrite.rpc; + +import com.fasterxml.jackson.module.paramnames.ParameterNamesModule; +import io.moderne.jsonrpc.JsonRpc; +import io.moderne.jsonrpc.formatter.JsonMessageFormatter; +import io.moderne.jsonrpc.handler.HeaderDelimitedMessageHandler; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.openrewrite.marketplace.RecipeMarketplace; + +import java.io.ByteArrayInputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Lifecycle accounting for {@link RewriteRpcProcessManager}. These verify that the + * number of live RPC "servers" stays bounded by the number of currently-live RPC + * threads — not by how many threads, files, or runs are processed — and is driven to + * zero on demand. + *

+ * A {@link CountingRpc} stands in for a real server: it increments a shared counter on + * construction and decrements it on {@link RewriteRpc#shutdown()}, so the counter mirrors + * the count of live node subprocesses without spawning any. + */ +// The shared static CountingRpc.live counter, reset per test by newManager(), is only +// safe under sequential execution; pin it so enabling method-level parallelism elsewhere +// can't corrupt these assertions. +@Execution(ExecutionMode.SAME_THREAD) +class RewriteRpcProcessManagerTest { + + /** + * A no-op {@link RewriteRpc} that tracks how many instances are currently "alive". + * Stands in for a spawned node server so the manager's create/dispose accounting can be + * asserted without real subprocesses. Its backing {@link JsonRpc} reads an + * already-at-EOF stream and spins up a {@link java.util.concurrent.ForkJoinPool} on + * construction; {@link #shutdown()} chains to {@link RewriteRpc#shutdown()} to tear that + * pool down so its worker threads don't accumulate across the many stubs this suite builds. + */ + static class CountingRpc extends RewriteRpc { + static final AtomicInteger live = new AtomicInteger(); + + private boolean shutdown; + + CountingRpc() { + super(new JsonRpc(new HeaderDelimitedMessageHandler( + new JsonMessageFormatter(new ParameterNamesModule()), + new ByteArrayInputStream(new byte[0]), + new OutputStream() { + @Override + public void write(int b) { + // discard + } + })), new RecipeMarketplace()); + live.incrementAndGet(); + } + + @Override + public void shutdown() { + // Idempotent: a server must only ever be counted down once even if it is + // both reaped and explicitly shut down. + if (!shutdown) { + shutdown = true; + live.decrementAndGet(); + // Tear down the backing JsonRpc ForkJoinPool so its worker threads don't + // accumulate across the hundreds of stubs this suite constructs. + super.shutdown(); + } + } + } + + /** + * A {@link CountingRpc} whose {@link #shutdown()} decrements the live counter (so accounting + * stays accurate) and then throws, modelling a server whose teardown fails — e.g. a node + * process reporting a non-zero exit code. Used to assert that one failed teardown does not + * strand the remaining servers. + */ + static class ThrowingRpc extends CountingRpc { + @Override + public void shutdown() { + super.shutdown(); + throw new RuntimeException("simulated shutdown failure"); + } + } + + private static RewriteRpcProcessManager newManager() { + CountingRpc.live.set(0); + return new RewriteRpcProcessManager<>(CountingRpc::new); + } + + /** + * Run {@code task} on a fresh thread and join it, so the thread is guaranteed dead + * (its {@link Thread#isAlive()} is {@code false}) by the time this returns. + */ + private static void runOnDeadThread(Runnable task) throws InterruptedException { + Thread t = new Thread(task, "rpc-test-worker"); + t.start(); + t.join(); + } + + @Test + void getOrStartReusesOnePerThread() { + RewriteRpcProcessManager manager = newManager(); + + CountingRpc first = manager.getOrStart(); + CountingRpc second = manager.getOrStart(); + + assertThat(second).as("the same thread must reuse its server, never spawn a second") + .isSameAs(first); + assertThat(manager.liveCount()).isEqualTo(1); + assertThat(CountingRpc.live).hasValue(1); + } + + @Test + void manySequentialRunsStayBoundedAndEndAtZero() { + RewriteRpcProcessManager manager = newManager(); + + // Each "run" lives on its own thread that starts a server then tears it down, + // mirroring the worker's per-run shutdownCurrent(). The live count must never + // exceed one regardless of how many runs go through, and must be zero between runs. + AtomicInteger peak = new AtomicInteger(); + for (int run = 0; run < 200; run++) { + int finalRun = run; + try { + runOnDeadThread(() -> { + manager.getOrStart(); + int live = manager.liveCount(); + peak.accumulateAndGet(live, Math::max); + assertThat(live) + .as("run %d should only have its own server live", finalRun) + .isEqualTo(1); + manager.shutdown(); + }); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + assertThat(CountingRpc.live) + .as("server torn down between runs") + .hasValue(0); + } + + assertThat(peak).as("never more than one live server across 200 sequential runs").hasValue(1); + assertThat(manager.liveCount()).isZero(); + assertThat(CountingRpc.live).hasValue(0); + } + + @Test + void deadThreadOrphansDoNotAccumulate() throws InterruptedException { + RewriteRpcProcessManager manager = newManager(); + + // Simulate transient threads (FJP compensation, cached/elastic pool threads that + // die on idle timeout) that each start a server and then terminate WITHOUT calling + // shutdown — the exact orphaning that used to leak node processes until JVM exit. + // Each new RPC thread reaps prior dead-thread orphans before adding its own, so the + // live count stays bounded instead of climbing to `orphans`. Pre-fix this would be 50. + int orphans = 50; + for (int i = 0; i < orphans; i++) { + runOnDeadThread(manager::getOrStart); + assertThat(manager.liveCount()) + .as("orphans must not accumulate as transient threads come and go") + .isLessThanOrEqualTo(1); + } + + // A new RPC thread appearing reaps any remaining dead-thread orphan before adding + // its own, so the live count is exactly this thread's server. + manager.getOrStart(); + assertThat(manager.liveCount()) + .as("dead-thread orphans reaped; only the live thread's server remains") + .isEqualTo(1); + assertThat(CountingRpc.live).hasValue(1); + + manager.shutdownAll(); + assertThat(CountingRpc.live).hasValue(0); + } + + @Test + void shutdownReapsDeadThreadOrphans() throws InterruptedException { + RewriteRpcProcessManager manager = newManager(); + + // The orchestrator thread (this one) owns a server, as does a transient thread + // that dies without cleaning up — the worker's shutdownCurrent() on the + // orchestrator thread must reach the orphan too, not just its own server. + manager.getOrStart(); + runOnDeadThread(manager::getOrStart); + assertThat(CountingRpc.live).hasValue(2); + + manager.shutdown(); + + assertThat(manager.liveCount()).isZero(); + assertThat(CountingRpc.live) + .as("shutdown() tears down the calling thread's server and reaps dead-thread orphans") + .hasValue(0); + } + + @Test + void shutdownAllDrivesLiveServersToZeroAcrossLiveThreads() throws InterruptedException { + RewriteRpcProcessManager manager = newManager(); + + // Hold several threads alive and parked, each owning a server, to model concurrent + // in-flight runs. shutdownAll() must reach servers on live threads too. + int concurrent = 8; + List threads = new ArrayList<>(); + AtomicInteger started = new AtomicInteger(); + Object park = new Object(); + AtomicBoolean released = new AtomicBoolean(); + for (int i = 0; i < concurrent; i++) { + Thread t = new Thread(() -> { + manager.getOrStart(); + started.incrementAndGet(); + synchronized (park) { + // Guarded wait: if notifyAll() fires before this thread parks, the + // released flag is already set, so we never block on a missed wakeup. + while (!released.get()) { + try { + park.wait(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + } + }, "rpc-test-live-" + i); + t.start(); + threads.add(t); + } + while (started.get() < concurrent) { + Thread.sleep(5); + } + assertThat(manager.liveCount()).isEqualTo(concurrent); + + manager.shutdownAll(); + + assertThat(manager.liveCount()).isZero(); + assertThat(CountingRpc.live) + .as("shutdownAll() stops every server, including those owned by live threads") + .hasValue(0); + + synchronized (park) { + released.set(true); + park.notifyAll(); + } + for (Thread t : threads) { + // Bounded join so a regression that strands a worker fails fast instead of hanging the suite. + t.join(10_000); + assertThat(t.isAlive()).as("parked worker %s should have been released", t.getName()).isFalse(); + } + } + + @Test + void shutdownAllKeepsGoingWhenAServerShutdownThrows() throws InterruptedException { + RewriteRpcProcessManager manager = newManager(); + + // One healthy server on this thread and one whose teardown throws, owned by a dead thread. + // shutdownAll() must tear down both — a single failing shutdown cannot strand the rest. + manager.getOrStart(); + runOnDeadThread(() -> { + manager.setFactory(ThrowingRpc::new); + manager.getOrStart(); + }); + assertThat(CountingRpc.live).hasValue(2); + + manager.shutdownAll(); + + assertThat(manager.liveCount()).isZero(); + assertThat(CountingRpc.live) + .as("a throwing teardown must not strand the healthy server") + .hasValue(0); + } + + @Test + void shutdownReapsDeadOrphanWhoseShutdownThrows() throws InterruptedException { + RewriteRpcProcessManager manager = newManager(); + + // This thread owns a healthy server; a dead thread left behind an orphan whose teardown + // throws. shutdown() reaps the orphan in its finally block and must swallow the failure + // while still tearing down the calling thread's own server. + manager.getOrStart(); + runOnDeadThread(() -> { + manager.setFactory(ThrowingRpc::new); + manager.getOrStart(); + }); + assertThat(CountingRpc.live).hasValue(2); + + manager.shutdown(); + + assertThat(manager.liveCount()).isZero(); + assertThat(CountingRpc.live) + .as("reaping a dead orphan whose shutdown throws must not strand other servers") + .hasValue(0); + } +} diff --git a/rewrite-csharp/src/main/java/org/openrewrite/csharp/rpc/CSharpRewriteRpc.java b/rewrite-csharp/src/main/java/org/openrewrite/csharp/rpc/CSharpRewriteRpc.java index 3c218740366..345313a9533 100644 --- a/rewrite-csharp/src/main/java/org/openrewrite/csharp/rpc/CSharpRewriteRpc.java +++ b/rewrite-csharp/src/main/java/org/openrewrite/csharp/rpc/CSharpRewriteRpc.java @@ -85,10 +85,24 @@ public void shutdown() { process.shutdown(); } + /** + * Shut down the current thread's C# RPC server, and reap any servers whose + * owning thread has since died. Reaping is scoped to dead threads only — a server + * owned by a still-live thread is never touched. + */ public static void shutdownCurrent() { MANAGER.shutdown(); } + /** + * Shut down every C# RPC server across all threads. Only call this when no + * C# parsing or printing is in flight (e.g. on JVM/service shutdown); + * {@link #shutdownCurrent()} already reaps servers orphaned by dead threads. + */ + public static void shutdownAll() { + MANAGER.shutdownAll(); + } + public InstallRecipesResponse installRecipes(java.io.File recipes) { return send( "InstallRecipes", diff --git a/rewrite-go/src/main/java/org/openrewrite/golang/rpc/GoRewriteRpc.java b/rewrite-go/src/main/java/org/openrewrite/golang/rpc/GoRewriteRpc.java index f82396d4fbe..252c3a39305 100644 --- a/rewrite-go/src/main/java/org/openrewrite/golang/rpc/GoRewriteRpc.java +++ b/rewrite-go/src/main/java/org/openrewrite/golang/rpc/GoRewriteRpc.java @@ -90,10 +90,24 @@ public void shutdown() { process.shutdown(); } + /** + * Shut down the current thread's Go RPC server, and reap any servers whose + * owning thread has since died. Reaping is scoped to dead threads only — a server + * owned by a still-live thread is never touched. + */ public static void shutdownCurrent() { MANAGER.shutdown(); } + /** + * Shut down every Go RPC server across all threads. Only call this when no + * Go parsing or printing is in flight (e.g. on JVM/service shutdown); + * {@link #shutdownCurrent()} already reaps servers orphaned by dead threads. + */ + public static void shutdownAll() { + MANAGER.shutdownAll(); + } + public static void resetCurrent() { MANAGER.reset(); } diff --git a/rewrite-javascript/src/main/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpc.java b/rewrite-javascript/src/main/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpc.java index 49cc1a20392..290699e2bf6 100644 --- a/rewrite-javascript/src/main/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpc.java +++ b/rewrite-javascript/src/main/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpc.java @@ -90,10 +90,24 @@ public void shutdown() { process.shutdown(); } + /** + * Shut down the current thread's JavaScript RPC server, and reap any servers whose + * owning thread has since died. Reaping is scoped to dead threads only — a server + * owned by a still-live thread is never touched. + */ public static void shutdownCurrent() { MANAGER.shutdown(); } + /** + * Shut down every JavaScript RPC server across all threads. Only call this when no + * JavaScript parsing or printing is in flight (e.g. on JVM/service shutdown); + * {@link #shutdownCurrent()} already reaps servers orphaned by dead threads. + */ + public static void shutdownAll() { + MANAGER.shutdownAll(); + } + public static void resetCurrent() { MANAGER.reset(); } diff --git a/rewrite-python/src/main/java/org/openrewrite/python/rpc/PythonRewriteRpc.java b/rewrite-python/src/main/java/org/openrewrite/python/rpc/PythonRewriteRpc.java index 0625477463c..960f7e057f5 100644 --- a/rewrite-python/src/main/java/org/openrewrite/python/rpc/PythonRewriteRpc.java +++ b/rewrite-python/src/main/java/org/openrewrite/python/rpc/PythonRewriteRpc.java @@ -87,10 +87,24 @@ public void shutdown() { process.shutdown(); } + /** + * Shut down the current thread's Python RPC server, and reap any servers whose + * owning thread has since died. Reaping is scoped to dead threads only — a server + * owned by a still-live thread is never touched. + */ public static void shutdownCurrent() { MANAGER.shutdown(); } + /** + * Shut down every Python RPC server across all threads. Only call this when no + * Python parsing or printing is in flight (e.g. on JVM/service shutdown); + * {@link #shutdownCurrent()} already reaps servers orphaned by dead threads. + */ + public static void shutdownAll() { + MANAGER.shutdownAll(); + } + public static void resetCurrent() { MANAGER.reset(); }