From 28ea479f67af1c5e7a6ec6be086c8eacab9b99e5 Mon Sep 17 00:00:00 2001 From: Jonathan Schneider Date: Sun, 16 Aug 2026 13:22:32 -0400 Subject: [PATCH 1/2] Keep interned RPC refs alive across source files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `RewriteRpc.evict` dropped a source file's tree from both peers and *also* rolled each side's ref maps back to a checkpoint taken before that file. Ref ids are assigned sequentially and every file's checkpoint is effectively 0, so nothing interned survived from one source file to the next: each file re-sent the full transitive closure of everything it referenced. Measured over 25 consecutive spring-boot files from one source set, evicting refs costs 3,716,349,765 bytes against 324,869,841 with refs retained — 11.4x. The ratio grows with file count; by file 25 the marginal cost is 149 MB versus 216 KB. Tree eviction stays. `Evict`'s payload does not change, so there is no wire-format change and the five language servers are independently landable. Retention is cheap on the receiver: peak today is one file, almost entirely a `JavaSourceSet` it reconstructs and discards once per file. Keeping refs it holds one copy. The interned type set grows about 2x (132,439 types across 25 files against 64,573 for the largest single file), roughly +3 MB against a 150 MB peak. Three rollbacks are deliberately left in place. They undo the refs a single *failed transfer* allocated, which is orthogonal to eviction: - `rewrite-core/.../rpc/request/GetObject.java` (`savedRefCount` / `removeIf`) - `rewrite-javascript/rewrite/src/rpc/request/get-object.ts` (`snapshot`/`rollbackTo`), which is why `ReferenceMap.snapshot`/`rollbackTo` stay in `reference.ts` - `rewrite-go/cmd/rpc/main.go` `SendQueue.DiscardNewReferences` Python had two layers that had to go together — the child's receive-side `_ref_checkpoints` and the facade's send-side `_hub_send_checkpoint`. Removing either alone produces `Received reference to unknown object` in facade mode only, which the leaf-server tests do not reproduce. The monotonic `_hub_send_next` counter is kept: it is what makes cross-file reuse work. `RewriteRpc.evict(String, int, int)` becomes `evict(String)`. Its only callers are `RecipeRunCycle` and `RewriteRpcTest`, and it shipped only in #8297, so there is no deprecated shim. The three tests that asserted lockstep rollback now assert the opposite invariant rather than being deleted, so they still distinguish the new behavior from the old. --- .../java/org/openrewrite/rpc/RewriteRpc.java | 33 +++--------- .../scheduling/RecipeRunCycle.java | 37 ++++---------- .../org/openrewrite/rpc/RewriteRpcTest.java | 17 ++++--- .../CSharp/Rpc/RewriteRpcServer.cs | 51 ++----------------- rewrite-go/cmd/rpc/main.go | 48 ++--------------- rewrite-go/pkg/rpc/reference.go | 22 -------- rewrite-javascript/rewrite/src/reference.ts | 4 ++ .../rewrite/src/rpc/request/batch-visit.ts | 2 - .../rewrite/src/rpc/request/visit.ts | 2 - .../rewrite/src/rpc/rewrite-rpc.ts | 37 ++------------ .../rewrite/src/rewrite/rpc/server.py | 37 ++------------ .../rewrite/tests/rpc/test_bundle_children.py | 2 +- .../rewrite/tests/rpc/test_facade.py | 32 ++++++------ .../rewrite/tests/rpc/test_server.py | 25 ++++----- 14 files changed, 71 insertions(+), 278 deletions(-) diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java index 1cd283153e7..dd19d714ec1 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java @@ -258,7 +258,7 @@ protected Boolean handle(Void noParams) { jsonRpc.rpc("Evict", new JsonRpcMethod() { @Override protected Boolean handle(Evict request) { - // Inbound side has no per-file checkpoint, so refs are left for Reset. + // Refs are left for Reset, so the next file reuses them. remoteObjects.remove(request.getId()); localObjects.remove(request.getId()); return true; @@ -367,37 +367,16 @@ public void reset() { } /** - * Ref high-water marks (send-side count, receive-side max key) captured before a file is - * visited so {@link #evict} rolls back exactly that file's refs. Receive-side uses the max - * key, not the size, because remote ids may be zero-based. + * Drop a file's tree from both peers. Interned refs deliberately survive so the next file + * reuses them rather than re-sending the objects they stand for; only {@link #reset} clears + * them. Notified fire-and-forget — under source-outer iteration the file's transfer is + * already complete. */ - public int[] refCheckpoint() { - int remoteRefsMax = -1; - for (Integer ref : remoteRefs.keySet()) { - if (ref > remoteRefsMax) { - remoteRefsMax = ref; - } - } - return new int[]{localRefs.size(), remoteRefsMax}; - } - - /** - * Drop a file's tree from both peers and roll their refs back to the pre-file checkpoint. - * Symmetric by design: dropping the send-side ref forces the next file to re-{@code ADD} the - * interned object instead of a {@code REF_USE} the rolled-back receiver would reject. Notified - * fire-and-forget — under source-outer iteration the file's transfer is already complete. - * - * @param localRefsCheckpoint {@code refCheckpoint()[0]} captured before the file was visited - * @param remoteRefsCheckpoint {@code refCheckpoint()[1]} captured before the file was visited - */ - public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint) { + public void evict(String id) { jsonRpc.notify(new JsonRpcRequest(null, "Evict", RawJson.of(new Evict(id)))); remoteObjects.remove(id); localObjects.remove(id); - - localRefs.values().removeIf(ref -> ref > localRefsCheckpoint); - remoteRefs.keySet().removeIf(ref -> ref > remoteRefsCheckpoint); } public

@Nullable Tree visit(SourceFile sourceFile, String visitorName, P p) { diff --git a/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java b/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java index 11bf935e7bc..a03f4517241 100644 --- a/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java +++ b/rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java @@ -125,7 +125,6 @@ public LSS scanSources(LSS sourceSet) { return sourceSetEditor.apply(sourceSet, sourceFile -> { BatchState scanBatch = new BatchState(); Set touched = newSetFromMap(new IdentityHashMap<>()); - Map refCheckpoints = new IdentityHashMap<>(); SourceFile result = allRecipeStack.reduce(sourceSet, recipe, ctx, (source, recipeStack) -> { Recipe recipe = leaf(recipeStack); @@ -140,11 +139,9 @@ public LSS scanSources(LSS sourceSet) { RewriteRpc currentRpc = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getRpc() : null; String scanVisitorName = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getScanVisitor() : null; - if (scanVisitorName != null) { - captureRpc(currentRpc, touched, refCheckpoints); - } - if (currentRpc != null && scanVisitorName != null) { + touched.add(currentRpc); + // Flush if switching to a different RPC instance if (scanBatch.rpc != null && scanBatch.rpc != currentRpc) { flushScanBatch(scanBatch, source); @@ -204,7 +201,7 @@ public LSS scanSources(LSS sourceSet) { flushScanBatch(scanBatch, result); } - evictSourceFile(sourceFile, touched, refCheckpoints); + evictSourceFile(sourceFile, touched); return result; }); } @@ -230,29 +227,16 @@ private void flushScanBatch(BatchState batch, SourceFile source) { } /** - * Record a peer this file touched, snapshotting its ref high-water on first sight so - * {@link #evictSourceFile} can roll back exactly the refs this file introduced. + * Drop this source file's tree from every RPC peer that visited it, bounding each peer's tree + * cache to ~one file at a time. Interned refs survive so the next file reuses them. */ - private static void captureRpc(@Nullable RewriteRpc rpc, Set touched, - Map refCheckpoints) { - if (rpc != null && touched.add(rpc)) { - refCheckpoints.put(rpc, rpc.refCheckpoint()); - } - } - - /** - * Drop this source file's tree from every RPC peer that visited it, rolling each peer's - * ref maps back to the pre-file checkpoint. Bounds RPC-server memory to ~one file at a time. - */ - private static void evictSourceFile(@Nullable SourceFile sourceFile, Set touched, - Map refCheckpoints) { + private static void evictSourceFile(@Nullable SourceFile sourceFile, Set touched) { if (sourceFile == null || touched.isEmpty()) { return; } String id = sourceFile.getId().toString(); for (RewriteRpc rpc : touched) { - int[] cp = refCheckpoints.get(rpc); - rpc.evict(id, cp[0], cp[1]); + rpc.evict(id); } } @@ -355,7 +339,6 @@ void clear() { recipeRunStats.recordSourceVisited(sourceFile); BatchState batch = new BatchState(); Set touched = newSetFromMap(new IdentityHashMap<>()); - Map refCheckpoints = new IdentityHashMap<>(); SourceFile result = allRecipeStack.reduce(sourceSet, recipe, ctx, (source, recipeStack) -> { Recipe recipe = leaf(recipeStack); @@ -364,7 +347,9 @@ void clear() { } RewriteRpc currentRpc = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getRpc() : null; - captureRpc(currentRpc, touched, refCheckpoints); + if (currentRpc != null) { + touched.add(currentRpc); + } // Flush batch if switching to a different RPC or non-RPC recipe if (batch.rpc != null && batch.rpc != currentRpc) { @@ -488,7 +473,7 @@ void clear() { } // Recipe errors are handled inside the reduce, so this runs on every normal return. - evictSourceFile(sourceFile, touched, refCheckpoints); + evictSourceFile(sourceFile, touched); return result; } diff --git a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java index f08b61ad165..9c44676dd9a 100644 --- a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java +++ b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java @@ -45,8 +45,10 @@ import java.io.PipedInputStream; import java.io.PipedOutputStream; import java.nio.file.Path; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.CountDownLatch; @@ -199,8 +201,8 @@ void sendFailureCleansUpRemoteObjects() { } /** - * {@link RewriteRpc#evict} drops the tree from both peers and rolls the client's ref maps - * back to the pre-file checkpoint. + * {@link RewriteRpc#evict} drops the tree from both peers but leaves every interned ref in + * place, so the next source file reuses them instead of re-sending what they stand for. */ @SneakyThrows @Test @@ -212,9 +214,6 @@ void evictDropsTreeFromBothPeers() { String id = original.getId().toString(); String sourceFileType = PlainText.class.getName(); - // High-water before the client fetches anything, so evict rolls back exactly this exchange. - int[] checkpoint = client.refCheckpoint(); - // Server holds the tree; client fetches it → both peers cache it. server.localObjects.put(id, original); client.getObject(id, sourceFileType); @@ -223,12 +222,14 @@ void evictDropsTreeFromBothPeers() { assertThat(server.localObjects).containsKey(id); assertThat(server.remoteObjects).containsKey(id); - client.evict(id, checkpoint[0], checkpoint[1]); + Set refsBefore = new HashSet<>(client.remoteRefs.keySet()); + + client.evict(id); - // Client cleared synchronously, including refs rolled back to the checkpoint. + // Trees cleared synchronously; refs survive, which is what makes cross-file interning work. assertThat(client.localObjects).doesNotContainKey(id); assertThat(client.remoteObjects).doesNotContainKey(id); - assertThat(client.remoteRefs.keySet()).allMatch(ref -> ref <= checkpoint[1]); + assertThat(client.remoteRefs.keySet()).containsExactlyInAnyOrderElementsOf(refsBefore); // The Evict notification is fire-and-forget; wait for the server to apply it. long deadline = System.currentTimeMillis() + 5_000; diff --git a/rewrite-csharp/csharp/OpenRewrite/CSharp/Rpc/RewriteRpcServer.cs b/rewrite-csharp/csharp/OpenRewrite/CSharp/Rpc/RewriteRpcServer.cs index b3bb5a37a7d..130522637db 100644 --- a/rewrite-csharp/csharp/OpenRewrite/CSharp/Rpc/RewriteRpcServer.cs +++ b/rewrite-csharp/csharp/OpenRewrite/CSharp/Rpc/RewriteRpcServer.cs @@ -86,12 +86,6 @@ public class RewriteRpcServer /// private readonly ConcurrentDictionary _remoteRefs = new(); - ///

- /// Ref high-water per source file (send-side _localRefs count, receive-side max _remoteRefs - /// key), captured before first visit so rolls back exactly its refs. - /// - private readonly ConcurrentDictionary _refCheckpoints = new(); - /// /// DependencyTypes pages its (potentially hundreds-of-MB) response: the full RpcObjectData list /// is built once, cached keyed by coordinate, and handed back one @@ -1448,7 +1442,6 @@ public async Task Visit(VisitRequest request) } // Fetch tree from the remote (Java) process - CaptureRefCheckpoint(request.TreeId); var tree = await GetObjectFromRemoteAsync(request.TreeId, request.SourceFileType); if (phase != "scan" && phase != "edit") @@ -1503,7 +1496,6 @@ public async Task BatchVisit(BatchVisitRequest request) } var sw = Stopwatch.StartNew(); - CaptureRefCheckpoint(request.TreeId); var tree = await GetObjectFromRemoteAsync(request.TreeId, request.SourceFileType); var fetchMs = sw.ElapsedMilliseconds; @@ -1806,35 +1798,15 @@ private void ClearLocalState() _remoteObjects.Clear(); _localRefs.Clear(); _remoteRefs.Clear(); - _refCheckpoints.Clear(); _preparedRecipes.Clear(); _recipeAccumulators.Clear(); _executionContexts.Clear(); } /// - /// Records the ref high-water before a source file is first visited (first visit wins), so - /// can roll back exactly the refs that file introduced. - /// - private void CaptureRefCheckpoint(string treeId) - { - _refCheckpoints.GetOrAdd(treeId, _ => - { - var remoteMax = -1; - foreach (var key in _remoteRefs.Keys) - { - if (key > remoteMax) - { - remoteMax = key; - } - } - return (_localRefs.Count, remoteMax); - }); - } - - /// - /// Drops one source file's tree and rolls back the refs it introduced; recipe/accumulator/ - /// context state (keyed separately) is preserved. Fire-and-forget, so it returns no response. + /// Drops one source file's tree. Interned refs are deliberately kept so the next file reuses + /// them; only Reset clears them. Recipe/accumulator/context state (keyed separately) is + /// preserved. Fire-and-forget, so it returns no response. /// [JsonRpcMethod("Evict", UseSingleObjectParameterDeserialization = true)] public void Evict(EvictRequest request) @@ -1845,23 +1817,6 @@ public void Evict(EvictRequest request) } _localObjects.TryRemove(request.Id, out _); _remoteObjects.TryRemove(request.Id, out _); - if (_refCheckpoints.TryRemove(request.Id, out var cp)) - { - foreach (var kv in _localRefs) - { - if (kv.Value > cp.LocalRefs) - { - _localRefs.TryRemove(kv.Key, out _); - } - } - foreach (var key in _remoteRefs.Keys) - { - if (key > cp.RemoteRefsMax) - { - _remoteRefs.TryRemove(key, out _); - } - } - } } /// diff --git a/rewrite-go/cmd/rpc/main.go b/rewrite-go/cmd/rpc/main.go index 13f816717e3..00a4bf5cc4b 100644 --- a/rewrite-go/cmd/rpc/main.go +++ b/rewrite-go/cmd/rpc/main.go @@ -74,13 +74,6 @@ type rpcError struct { Data string `json:"data,omitempty"` } -// evictCheckpoint is a peer's ref high-water before a file is visited: localRefsNext (send, -// Go→Java) and reverseRemoteRefsMax (receive, Java→Go), so evict rolls back exactly its refs. -type evictCheckpoint struct { - localRefsNext int - reverseRemoteRefsMax int -} - type server struct { localObjects map[string]any remoteObjects map[string]any // forward direction: tracks what Java has from Go @@ -107,10 +100,6 @@ type server struct { reverseTypePool map[string]java.JavaType - // Ref high-water marks captured before a source file is first visited, keyed by tree id, - // so handleEvict can roll back exactly the refs that file introduced (see handleEvict). - refCheckpoints map[string]evictCheckpoint - // Prepared recipe instances keyed by unique ID preparedRecipes map[string]recipe.Recipe @@ -236,7 +225,6 @@ func newServer(cfg serverConfig) *server { reverseRemoteObjects: make(map[string]any), reverseRemoteRefs: make(map[int]any), reverseTypePool: make(map[string]java.JavaType), - refCheckpoints: make(map[string]evictCheckpoint), preparedRecipes: make(map[string]recipe.Recipe), preparedRecipeNames: make(map[string]string), preparedEditorOverrides: make(map[string]recipe.TreeVisitor), @@ -1183,36 +1171,17 @@ func (s *server) handleReset() bool { s.reverseRemoteObjects = make(map[string]any) s.reverseRemoteRefs = make(map[int]any) s.reverseTypePool = make(map[string]java.JavaType) - s.refCheckpoints = make(map[string]evictCheckpoint) s.preparedRecipes = make(map[string]recipe.Recipe) s.preparedRecipeNames = make(map[string]string) s.preparedEditorOverrides = make(map[string]recipe.TreeVisitor) s.preparedAccumulators = make(map[string]any) s.preparedContexts = make(map[string]*recipe.ExecutionContext) - s.refCheckpoints = make(map[string]evictCheckpoint) return true } -// captureRefCheckpoint records the ref high-water before a file is first visited (first visit -// wins), keyed by tree id, so handleEvict rolls back exactly the refs that file introduced. -func (s *server) captureRefCheckpoint(treeID string) { - if _, ok := s.refCheckpoints[treeID]; ok { - return - } - maxKey := -1 - for k := range s.reverseRemoteRefs { - if k > maxKey { - maxKey = k - } - } - s.refCheckpoints[treeID] = evictCheckpoint{ - localRefsNext: s.localRefs.NextID(), - reverseRemoteRefsMax: maxKey, - } -} - -// handleEvict drops one source file's tree and rolls back the refs it introduced. Recipe/ -// accumulator/context state (keyed separately) is preserved. Fire-and-forget, so it never errors. +// handleEvict drops one source file's tree. Interned refs are deliberately kept so the next +// file reuses them; only Reset clears them. Recipe/accumulator/context state (keyed +// separately) is preserved. Fire-and-forget, so it never errors. func (s *server) handleEvict(params json.RawMessage) bool { var req struct { ID string `json:"id"` @@ -1226,15 +1195,6 @@ func (s *server) handleEvict(params json.RawMessage) bool { } delete(s.localObjects, req.ID) delete(s.remoteObjects, req.ID) - if cp, ok := s.refCheckpoints[req.ID]; ok { - s.localRefs.RollbackTo(cp.localRefsNext) - for k := range s.reverseRemoteRefs { - if k > cp.reverseRemoteRefsMax { - delete(s.reverseRemoteRefs, k) - } - } - delete(s.refCheckpoints, req.ID) - } return true } @@ -1920,7 +1880,6 @@ func (s *server) handleVisit(params json.RawMessage) (any, *rpcError) { } // Get the tree from Java via bidirectional RPC - s.captureRefCheckpoint(req.TreeID) treeObj := s.getObjectFromJava(req.TreeID, req.SourceFileType) if treeObj == nil { return &visitResponse{Modified: false}, nil @@ -2087,7 +2046,6 @@ func (s *server) handleBatchVisit(params json.RawMessage) (any, *rpcError) { ctx := s.resolveExecutionContext(req.PID) - s.captureRefCheckpoint(req.TreeID) treeObj := s.getObjectFromJava(req.TreeID, req.SourceFileType) current, _ := treeObj.(java.Tree) if current == nil { diff --git a/rewrite-go/pkg/rpc/reference.go b/rewrite-go/pkg/rpc/reference.go index 2e2c5a2ecbd..ab10cb0bc97 100644 --- a/rewrite-go/pkg/rpc/reference.go +++ b/rewrite-go/pkg/rpc/reference.go @@ -61,28 +61,6 @@ func (m *ReferenceMap) Len() int { return len(m.refs) } -// NextID returns the id that will be assigned to the next new reference. Capture it -// before a source file is visited to use as a rollback checkpoint (see RollbackTo). -func (m *ReferenceMap) NextID() int { - m.mu.Lock() - defer m.mu.Unlock() - return m.nextID -} - -// RollbackTo drops every reference assigned at or after the checkpoint and resets the id -// counter so the next file re-allocates the same ids. Unlike DiscardNewReferences, this is -// safe only because the remote receiver rolls back in lockstep (per-source-file evict). -func (m *ReferenceMap) RollbackTo(checkpoint int) { - m.mu.Lock() - defer m.mu.Unlock() - for obj, ref := range m.refs { - if ref >= checkpoint { - delete(m.refs, obj) - } - } - m.nextID = checkpoint -} - func (m *ReferenceMap) deleteIfMatches(obj any, ref int) { m.mu.Lock() defer m.mu.Unlock() diff --git a/rewrite-javascript/rewrite/src/reference.ts b/rewrite-javascript/rewrite/src/reference.ts index 8de75125d30..e783d696652 100644 --- a/rewrite-javascript/rewrite/src/reference.ts +++ b/rewrite-javascript/rewrite/src/reference.ts @@ -93,6 +93,10 @@ export class ReferenceMap { return this.refsById.size; } + /** + * Undo the refs a single failed transfer allocated. Only `GetObject` calls this — refs + * survive an `Evict` so the next source file reuses them. + */ rollbackTo(savedRefCount: number): void { for (let i = savedRefCount; i < this.refCount; i++) { const obj = this.refsById.get(i); diff --git a/rewrite-javascript/rewrite/src/rpc/request/batch-visit.ts b/rewrite-javascript/rewrite/src/rpc/request/batch-visit.ts index eec26bc023b..6848dc76cb6 100644 --- a/rewrite-javascript/rewrite/src/rpc/request/batch-visit.ts +++ b/rewrite-javascript/rewrite/src/rpc/request/batch-visit.ts @@ -68,7 +68,6 @@ export class BatchVisit { preparedRecipes: Map, recipeCursors: WeakMap, getObject: (id: string, sourceFileType?: string) => any, - captureRefCheckpoint: (treeId: string) => void, getCursor: (cursorIds: string[] | undefined, sourceFileType?: string) => Promise, dataTableStore: () => DataTableStore | undefined, metricsCsv?: string): void { @@ -83,7 +82,6 @@ export class BatchVisit { if (store && p instanceof ExecutionContext) { p.messages[DATA_TABLE_STORE] = store; } - captureRefCheckpoint(request.treeId); let tree: Tree = await getObject(request.treeId, request.sourceFileType); const cursor = await getCursor(request.cursor, request.sourceFileType); diff --git a/rewrite-javascript/rewrite/src/rpc/request/visit.ts b/rewrite-javascript/rewrite/src/rpc/request/visit.ts index 2101547b992..d0d5571ed4d 100644 --- a/rewrite-javascript/rewrite/src/rpc/request/visit.ts +++ b/rewrite-javascript/rewrite/src/rpc/request/visit.ts @@ -44,7 +44,6 @@ export class Visit { preparedRecipes: Map, recipeCursors: WeakMap, getObject: (id: string, sourceFileType?: string) => any, - captureRefCheckpoint: (treeId: string) => void, getCursor: (cursorIds: string[] | undefined, sourceFileType?: string) => Promise, dataTableStore: () => DataTableStore | undefined, metricsCsv?: string): void { @@ -59,7 +58,6 @@ export class Visit { if (store && p instanceof ExecutionContext) { p.messages[DATA_TABLE_STORE] = store; } - captureRefCheckpoint(request.treeId); const before: Tree = await getObject(request.treeId, request.sourceFileType); const cursor = await getCursor(request.cursor, request.sourceFileType); context.target = extractSourcePath(before, cursor); diff --git a/rewrite-javascript/rewrite/src/rpc/rewrite-rpc.ts b/rewrite-javascript/rewrite/src/rpc/rewrite-rpc.ts index 8e13ff2f6fc..7da2a0b1bac 100644 --- a/rewrite-javascript/rewrite/src/rpc/rewrite-rpc.ts +++ b/rewrite-javascript/rewrite/src/rpc/rewrite-rpc.ts @@ -75,10 +75,6 @@ export class RewriteRpc { readonly remoteRefs: Map = new Map(); readonly localRefs: ReferenceMap = new ReferenceMap(); - // Ref high-water per source file, captured before it is first visited so an Evict rolls - // back exactly the refs it introduced. `send` = localRefs snapshot, `recvMax` = max remoteRefs key. - readonly refCheckpoints: Map = new Map(); - private remoteLanguages?: string[]; private readonly logger?: rpc.Logger; private traceGetObject: TraceGetObject = {receive: false, send: false}; @@ -111,19 +107,6 @@ export class RewriteRpc { // Need this indirection, otherwise `this` will be undefined when executed in the handlers. const getObject = (id: string, sourceFileType?: string) => this.getObject(id, sourceFileType); const getCursor = (cursorIds: string[] | undefined, sourceFileType?: string) => this.getCursor(cursorIds, sourceFileType); - // First visit of the file wins. - const captureRefCheckpoint = (treeId: string) => { - if (this.refCheckpoints.has(treeId)) { - return; - } - let recvMax = -1; - for (const k of this.remoteRefs.keys()) { - if (k > recvMax) { - recvMax = k; - } - } - this.refCheckpoints.set(treeId, {send: this.localRefs.snapshot(), recvMax}); - }; const traceGetObject = () => this.traceGetObject.send; const dataTableStore = () => this.configuredDataTableStore; @@ -132,8 +115,8 @@ export class RewriteRpc { // GetMarketplace builds rows so the host can attribute each recipe to its own bundle. const recipeOrigin: Map = new Map(); - Visit.handle(this.connection, this.localObjects, preparedRecipes, recipeCursors, getObject, captureRefCheckpoint, getCursor, dataTableStore, options.metricsCsv); - BatchVisit.handle(this.connection, this.localObjects, preparedRecipes, recipeCursors, getObject, captureRefCheckpoint, getCursor, dataTableStore, options.metricsCsv); + Visit.handle(this.connection, this.localObjects, preparedRecipes, recipeCursors, getObject, getCursor, dataTableStore, options.metricsCsv); + BatchVisit.handle(this.connection, this.localObjects, preparedRecipes, recipeCursors, getObject, getCursor, dataTableStore, options.metricsCsv); Generate.handle(this.connection, this.localObjects, preparedRecipes, recipeCursors, getObject, dataTableStore, options.metricsCsv); SetDataTableStore.handle(this.connection, store => this.configuredDataTableStore = store, options.metricsCsv); GetObject.handle(this.connection, this.remoteObjects, this.localObjects, @@ -164,7 +147,6 @@ export class RewriteRpc { this.remoteObjects.clear(); this.remoteRefs.clear(); this.localRefs.clear(); - this.refCheckpoints.clear(); preparedRecipes.clear(); this.remoteLanguages = undefined; }; @@ -181,24 +163,15 @@ export class RewriteRpc { } ) - // Drop one source file's tree + roll back the refs it introduced. Fire-and-forget - // notification (no reply), so recipe/accumulator/context state is left intact. + // Drop one source file's tree. Interned refs are deliberately kept so the next file + // reuses them; only Reset clears them. Fire-and-forget notification (no reply), so + // recipe/accumulator/context state is left intact. this.connection.onNotification( new rpc.NotificationType<{ id: string }>("Evict"), (params) => { const id = params.id; this.localObjects.delete(id); this.remoteObjects.delete(id); - const cp = this.refCheckpoints.get(id); - if (cp !== undefined) { - this.localRefs.rollbackTo(cp.send); - for (const k of [...this.remoteRefs.keys()]) { - if (k > cp.recvMax) { - this.remoteRefs.delete(k); - } - } - this.refCheckpoints.delete(id); - } } ) diff --git a/rewrite-python/rewrite/src/rewrite/rpc/server.py b/rewrite-python/rewrite/src/rewrite/rpc/server.py index 50ae0a4a459..c84e71e32c8 100644 --- a/rewrite-python/rewrite/src/rewrite/rpc/server.py +++ b/rewrite-python/rewrite/src/rewrite/rpc/server.py @@ -67,9 +67,6 @@ remote_objects: Dict[str, Any] = {} # Remote refs - maps reference IDs to objects for cyclic graph handling remote_refs: Dict[int, Any] = {} -# Per-source-file remote_refs high-water, captured before a file is first visited so -# handle_evict can roll back exactly the refs that file introduced. Keyed by tree id. -_ref_checkpoints: Dict[str, int] = {} # Per-call metrics CSV (--metrics-csv), same schema as Go: cache-size ramp vs per-file-Evict sawtooth. _metrics_file = None @@ -1110,15 +1107,14 @@ def handle_reset(params: dict) -> bool: _local_object_ids.clear() _recipe_accumulators.clear() _recipe_phases.clear() - _ref_checkpoints.clear() logger.info("Reset: cleared all cached state") return True def handle_evict(params: dict) -> bool: - """Handle an Evict RPC notification - drop one source file's tree and roll back the - refs it introduced, bounding memory to roughly one source file at a time. Recipe, + """Handle an Evict RPC notification - drop one source file's tree. Interned refs are + deliberately kept so the next file reuses them; only Reset clears them. Recipe, accumulator, and execution-context state (keyed separately) is left intact. """ obj_id = params.get('id') @@ -1126,10 +1122,6 @@ def handle_evict(params: dict) -> bool: return True local_objects.pop(obj_id, None) remote_objects.pop(obj_id, None) - checkpoint = _ref_checkpoints.pop(obj_id, None) - if checkpoint is not None: - for ref_id in [k for k in remote_refs if k > checkpoint]: - del remote_refs[ref_id] return True @@ -2062,9 +2054,6 @@ def handle_visit(params: dict) -> dict: ctx = _context_for(p_id) - # Snapshot the remote_refs high-water for this file before fetching its tree (first visit wins). - _ref_checkpoints.setdefault(tree_id, max(remote_refs.keys(), default=-1)) - # Always fetch the tree from Java to ensure we have the latest version. # Java may have modified the tree (e.g., via a Java-side recipe) since our last sync. tree = get_object_from_java(tree_id, source_file_type) @@ -2120,9 +2109,6 @@ def handle_batch_visit(params: dict) -> dict: ctx = _context_for(p_id) - # Snapshot the remote_refs high-water for this file before fetching its tree. - _ref_checkpoints.setdefault(tree_id, max(remote_refs.keys(), default=-1)) - # Fetch tree once from Java tree = get_object_from_java(tree_id, source_file_type) if tree is None: @@ -2345,7 +2331,6 @@ def handle_generate(params: dict) -> dict: _hub_send_refs: Dict[str, Dict] = {} # bundle -> send ref map (facade -> child) _hub_send_next: Dict[str, int] = {} # bundle -> next send ref number _hub_served: Dict[tuple, Any] = {} # (bundle, obj_id) -> what that child was last served -_hub_send_checkpoint: Dict[tuple, int] = {} # (bundle, obj_id) -> send ref counter before this file def _hub_acquire(obj_id: str, source_file_type: Optional[str]): """The facade's copy of the in-flight tree, fetched from Java (over the facade<->Java table) the @@ -2376,9 +2361,6 @@ def _hub_serve_child(bundle: str, obj_id: str, source_file_type: Optional[str]) q = RpcSendQueue(source_file_type) q.refs = _hub_send_refs.setdefault(bundle, {}) q.next_ref = _hub_send_next.get(bundle, 0) - # Remember where this child's ref numbering stood before this file, so Evict can roll it back - # in lockstep with the child's own rollback (see _hub_release). - _hub_send_checkpoint.setdefault((bundle, obj_id), q.next_ref) data = q.generate(tree, _hub_served.get((bundle, obj_id))) _hub_send_next[bundle] = q.next_ref _hub_served[(bundle, obj_id)] = tree @@ -2409,21 +2391,12 @@ def pull(): def _hub_release(obj_id: str) -> None: - """The rollback must be symmetric with the child's own Evict: the child drops the refs this file - introduced from its receive map, so if the facade kept them in its send map it would emit a - GET_REF for a ref the child no longer has ("Received reference to unknown object"). + """Drop the facade's copy of one file's tree. The send ref map is deliberately left alone and + stays symmetric with the child's receive map, which no longer drops refs on Evict either. """ _hub_tree.pop(obj_id, None) for key in [k for k in _hub_served if k[1] == obj_id]: del _hub_served[key] - for key in [k for k in _hub_send_checkpoint if k[1] == obj_id]: - bundle = key[0] - checkpoint = _hub_send_checkpoint.pop(key) - refs = _hub_send_refs.get(bundle) - if refs is not None: - for ref_key in [k for k, (_, num) in refs.items() if num > checkpoint]: - del refs[ref_key] - _hub_send_next[bundle] = checkpoint def _hub_is_builtin_visitor(visitor_name: Optional[str]) -> bool: @@ -2869,7 +2842,7 @@ def _rss_bytes(): def _record_metric(method: str, duration_ms: float, error: str) -> None: """Append one row of timing + cache residency. refs counts remote_refs only: Python's send-side - refs live on a per-call RpcSendQueue, the only cross-call ref cache handle_evict rolls back.""" + refs live on a per-call RpcSendQueue, so remote_refs is the only cross-call ref cache.""" if _metrics_writer is None: return used, peak = _rss_bytes() diff --git a/rewrite-python/rewrite/tests/rpc/test_bundle_children.py b/rewrite-python/rewrite/tests/rpc/test_bundle_children.py index b9ebb8c869e..8d31ca55384 100644 --- a/rewrite-python/rewrite/tests/rpc/test_bundle_children.py +++ b/rewrite-python/rewrite/tests/rpc/test_bundle_children.py @@ -122,7 +122,7 @@ def fake_spawn(cmd, upstream=None, exclude_paths=()): bc.install("pkgb", "pkgb") bc.broadcast_evict({"id": "tree-1"}) - # every live child is told to evict the file so each rolls back its own ref map + # every live child is told to evict the file so each drops its own copy of the tree assert ("Evict", {"id": "tree-1"}) in bc._children["pkga"].requests assert ("Evict", {"id": "tree-1"}) in bc._children["pkgb"].requests diff --git a/rewrite-python/rewrite/tests/rpc/test_facade.py b/rewrite-python/rewrite/tests/rpc/test_facade.py index 699398ea079..464d7b419d3 100644 --- a/rewrite-python/rewrite/tests/rpc/test_facade.py +++ b/rewrite-python/rewrite/tests/rpc/test_facade.py @@ -295,7 +295,7 @@ def generate(self, params): ... assert routed["install"] == r -def test_hub_release_rolls_each_childs_ref_table_back_in_lockstep(): +def test_hub_release_keeps_each_childs_ref_table_intact(): import rewrite.rpc.server as server server._hub_tree["T"] = object() @@ -304,16 +304,15 @@ def test_hub_release_rolls_each_childs_ref_table_back_in_lockstep(): server._hub_send_refs["A"] = {10: ("before-1", 1), 11: ("before-2", 2), 12: ("from-file", 3), 13: ("from-file", 4)} server._hub_send_next["A"] = 4 - server._hub_send_checkpoint[("A", "T")] = 2 server._hub_release("T") - # only the refs this file introduced are dropped; the pre-file ones survive - assert sorted(n for _, n in server._hub_send_refs["A"].values()) == [1, 2] - assert server._hub_send_next["A"] == 2 # counter rewound so the next file re-ADDs + # The child no longer drops refs on Evict, so the facade must not either -- rewinding here + # would re-ADD ref 3 for a different object while the child still holds the old one. + assert sorted(n for _, n in server._hub_send_refs["A"].values()) == [1, 2, 3, 4] + assert server._hub_send_next["A"] == 4 assert "T" not in server._hub_tree assert ("A", "T") not in server._hub_served - assert ("A", "T") not in server._hub_send_checkpoint class _PreconditionChildren(_FakeChildren): @@ -404,8 +403,7 @@ def _print_python(cu) -> str: def _isolated_hub(monkeypatch, server): """Fresh hub state so this test neither sees nor leaves module-global tables.""" - for name in ("_hub_tree", "_hub_served", "_hub_send_refs", "_hub_send_next", - "_hub_send_checkpoint", "local_objects"): + for name in ("_hub_tree", "_hub_served", "_hub_send_refs", "_hub_send_next", "local_objects"): monkeypatch.setattr(server, name, {}) @@ -467,10 +465,9 @@ def test_local_visit_advances_the_hub_tree_and_the_next_serve_carries_it_to_the_ def test_a_built_in_visitor_that_deletes_the_file_releases_it_from_the_hub(monkeypatch): """A built-in visitor may delete the file outright. The facade has to let go of it the same way - a broadcast Evict would: drop the tree, forget what each child was served, and rewind that - child's send-ref numbering — otherwise the next file reuses ref numbers the child no longer - holds. Nothing after the delete may run, and the child must be told DELETE, not served a - resurrected tree.""" + a broadcast Evict would: drop the tree and forget what each child was served, while leaving that + child's interned refs in place. Nothing after the delete may run, and the child must be told + DELETE, not served a resurrected tree.""" import rewrite.rpc.server as server from rewrite.python.visitor import PythonVisitor @@ -496,7 +493,9 @@ def visit_compilation_unit(self, cu, p): server._hub_acquire(tree_id, sft) server._hub_serve_child(bundle, tree_id, sft) # child now holds this file and its refs - assert server._hub_send_next[bundle] > 0 + served_refs = dict(server._hub_send_refs[bundle]) + served_next = server._hub_send_next[bundle] + assert served_next > 0 results = server._hub_local_visit([{"visitor": "Delete"}, {"visitor": "Later"}], {"treeId": tree_id, "sourceFileType": sft}) @@ -506,12 +505,11 @@ def visit_compilation_unit(self, cu, p): "hasNewMessages": False, "searchResultIds": []}] assert ran_after == [] - # The facade no longer owns the file, and every per-child table it seeded is rolled back. + # The facade no longer owns the file, but the child's refs stay valid for the next one. assert tree_id not in server._hub_tree assert (bundle, tree_id) not in server._hub_served - assert (bundle, tree_id) not in server._hub_send_checkpoint - assert server._hub_send_refs[bundle] == {} - assert server._hub_send_next[bundle] == 0 + assert server._hub_send_refs[bundle] == served_refs + assert server._hub_send_next[bundle] == served_next # So a child asking for it again is told the file is gone rather than served a stale tree. assert server._hub_serve_child(bundle, tree_id, sft) == [ diff --git a/rewrite-python/rewrite/tests/rpc/test_server.py b/rewrite-python/rewrite/tests/rpc/test_server.py index 8e3062a36ee..085570b4ffe 100644 --- a/rewrite-python/rewrite/tests/rpc/test_server.py +++ b/rewrite-python/rewrite/tests/rpc/test_server.py @@ -525,19 +525,17 @@ def text_of(child): assert [text_of(c) for c in children] == ["a", "b", "c"] -def test_hub_release_rewinds_send_refs_in_lockstep_with_the_child(): - """A child drops its receive refs for a file when the broadcast Evict reaches it, so the facade - must return its send-ref numbering to exactly the pre-file value. If the facade kept advancing, - it would emit a GET_REF for a ref the child no longer holds; if it rewound while the child did - not, it would reuse a number still bound to the old object and serve a wrong tree silently.""" +def test_hub_release_leaves_send_refs_for_the_next_file(): + """A child keeps its receive refs when the broadcast Evict reaches it, so the facade must keep + its send-ref numbering advancing. Rewinding would reuse a number still bound to the old object + on the child and serve a wrong tree silently.""" import rewrite.rpc.server as server - bundle, first, second = "pkg", "file-1", "file-2" + bundle, first = "pkg", "file-1" server._hub_send_refs[bundle] = {} server._hub_send_next[bundle] = 0 - # Serving the first file advances this child's numbering and records where it started. - server._hub_send_checkpoint.setdefault((bundle, first), server._hub_send_next[bundle]) + # Serving the first file advances this child's numbering. server._hub_send_refs[bundle].update({"obj-a": (object(), 1), "obj-b": (object(), 2)}) server._hub_send_next[bundle] = 2 server._hub_served[(bundle, first)] = object() @@ -545,13 +543,8 @@ def test_hub_release_rewinds_send_refs_in_lockstep_with_the_child(): server._hub_release(first) - # Everything that file introduced is gone, and the counter is back where it began. - assert server._hub_send_next[bundle] == 0 - assert server._hub_send_refs[bundle] == {} + # The tree is gone but the refs it interned are still there for the next file to hit. + assert server._hub_send_next[bundle] == 2 + assert sorted(n for _, n in server._hub_send_refs[bundle].values()) == [1, 2] assert (bundle, first) not in server._hub_served - assert (bundle, first) not in server._hub_send_checkpoint assert first not in server._hub_tree - - # So the next file reuses the same ref numbers rather than continuing past them. - server._hub_send_checkpoint.setdefault((bundle, second), server._hub_send_next[bundle]) - assert server._hub_send_checkpoint[(bundle, second)] == 0 From e299bde1a1f0485ba64afce13651193f4088739d Mon Sep 17 00:00:00 2001 From: Jonathan Schneider Date: Sun, 16 Aug 2026 13:42:14 -0400 Subject: [PATCH 2/2] Give `JavaSourceSet` an RPC codec in Java, TypeScript and C# MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `JavaSourceSet` did not implement `RpcCodec`, so `RpcSendQueue.add` took the inline path and emitted the whole marker as the value of a *single* `RpcObjectData`. On a spring-boot source set that is one JSON document of about 128 MB, built whole and parsed whole. `RpcSendQueue` batches by count, not bytes, so no batching bounds it. Bucketed by `valueType` over three files, that one marker was 89.6% of all bytes crossing the boundary — `JavaType` payloads were 1.9%. Size is not the only reason for a codec. The marker's classpath is a `JavaType` graph, and the tree is already sending `JavaType`s ref-deduplicated on the same `localRefs` map. Inlining means the receiver reconstructs a second, non-identical type universe it cannot share with the first. `gavToTypes` needed care. `RpcSendQueue.getValueType` returns null for a `Map`, so a naive `getAndSend` would inline the entire type graph a second time — a pessimization no round-trip test would catch, because the values would still arrive correct. It is decomposed into a key list plus one `getAndSendListAsRef` per key instead. Because the classpath goes first and the bucket values are the *same instances* (see `JavaSourceSet#build`), every nested element is a ref hit: ref-only ADDs rather than a second copy. ## Parity is a hard gate The marker rides on every file in a source set, resource files included, so it reaches any peer that accepts a resource type: | peer | accepts | codec | |---|---|---| | JavaScript | PlainText, JSON, YAML, JS | added | | C# | Cs, **Xml** | added | | Python | Py only | not reachable; follow-up | | Go | Go, GoMod, GoSum | not reachable; follow-up | Landing the Java codec alone would desynchronize both reachable peers. Two integration tests hold that gate, and each was confirmed to fail without its peer's codec: - `JavaScriptRewriteRpcTest#javaSourceSetMarkerAcrossRpcBoundary` — fails with "No RPC codec registered on the TypeScript side". - `CSharpJavaSourceSetRpcTest#javaSourceSetMarkerAcrossRpcBoundary` — fails with "No RPC codec registered on the C# side". Left alone, C# would have resolved the marker to `UnknownMarker(Guid.NewGuid())` and desynchronized with no diagnostic at all. `TypeSender`/`TypeReceiver` in `rewrite-javascript` were module-private and are now exported. On the C# side `org.openrewrite.java.marker.*` already resolves by convention, so only the send-direction name needed an entry — without it `OpenRewrite.Java.JavaSourceSet` would go back as `J$JavaSourceSet`. `JavaSourceSetRpcTest` covers the Java side, including the assertion that actually proves the decomposition works: a received `gavToTypes` bucket element is the *same instance* as its classpath entry. --- .../OpenRewrite/Core/Rpc/RpcSendQueue.cs | 1 + .../csharp/OpenRewrite/Java/JavaSourceSet.cs | 85 +++++++++++++ .../rpc/CSharpJavaSourceSetRpcTest.java | 94 +++++++++++++++ .../java/marker/JavaSourceSet.java | 53 +++++++- .../java/marker/JavaSourceSetRpcTest.java | 114 ++++++++++++++++++ .../rewrite/src/java/markers.ts | 53 +++++++- rewrite-javascript/rewrite/src/java/rpc.ts | 4 +- .../rpc/JavaScriptRewriteRpcTest.java | 23 ++++ 8 files changed, 421 insertions(+), 6 deletions(-) create mode 100644 rewrite-csharp/csharp/OpenRewrite/Java/JavaSourceSet.cs create mode 100644 rewrite-csharp/src/integTest/java/org/openrewrite/csharp/rpc/CSharpJavaSourceSetRpcTest.java create mode 100644 rewrite-java/src/test/java/org/openrewrite/java/marker/JavaSourceSetRpcTest.java diff --git a/rewrite-csharp/csharp/OpenRewrite/Core/Rpc/RpcSendQueue.cs b/rewrite-csharp/csharp/OpenRewrite/Core/Rpc/RpcSendQueue.cs index 7d4057706ad..40452de26e0 100644 --- a/rewrite-csharp/csharp/OpenRewrite/Core/Rpc/RpcSendQueue.cs +++ b/rewrite-csharp/csharp/OpenRewrite/Core/Rpc/RpcSendQueue.cs @@ -365,6 +365,7 @@ private void DoChange(object? after, object? before, Action? onChange, IRpcCodec { "JRightPadded" or "JLeftPadded" or "JContainer" or "JavaType" => $"org.openrewrite.java.tree.{name}", + "JavaSourceSet" => "org.openrewrite.java.marker.JavaSourceSet", _ => $"org.openrewrite.java.tree.J${name}", }, "OpenRewrite.CSharp" => $"org.openrewrite.csharp.tree.Cs${name}", diff --git a/rewrite-csharp/csharp/OpenRewrite/Java/JavaSourceSet.cs b/rewrite-csharp/csharp/OpenRewrite/Java/JavaSourceSet.cs new file mode 100644 index 00000000000..08961cf0993 --- /dev/null +++ b/rewrite-csharp/csharp/OpenRewrite/Java/JavaSourceSet.cs @@ -0,0 +1,85 @@ +/* + * Copyright 2026 the original author or authors. + *

+ * Licensed under the Moderne Source Available License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + *

+ * https://docs.moderne.io/licensing/moderne-source-available-license + *

+ * 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. + */ +using OpenRewrite.Core; +using OpenRewrite.Core.Rpc; +using OpenRewrite.Java.Rpc; + +namespace OpenRewrite.Java; + +///

+/// Mirrors org.openrewrite.java.marker.JavaSourceSet. It rides on every source file in a Java +/// source set, resource files included, so it reaches this peer through Xml as well as C#. +/// +/// Field order is the protocol; see JavaSourceSet#rpcSend on the Java side. Without this codec the +/// marker resolves to , which consumes one message where Java sends +/// many, and the queue then desynchronizes with no diagnostic. +/// +/// +public sealed class JavaSourceSet( + Guid id, + string name, + IList classpath, + IDictionary> gavToTypes +) : Marker, IRpcCodec +{ + public Guid Id { get; } = id; + public string Name { get; } = name; + public IList Classpath { get; } = classpath; + public IDictionary> GavToTypes { get; } = gavToTypes; + + public void RpcSend(JavaSourceSet after, RpcSendQueue q) + { + var typeSender = new JavaSender(); + q.GetAndSend(after, s => s.Id); + q.GetAndSend(after, s => s.Name); + q.GetAndSendListAsRef(after, s => s.Classpath, TypeKey, t => typeSender.VisitType(t, q)); + + var gavs = after.GavToTypes.Keys.ToList(); + q.GetAndSendList(after, _ => gavs, gav => gav, null); + foreach (var gav in gavs) + { + q.GetAndSendListAsRef(after, s => s.GavToTypes[gav], TypeKey, + t => typeSender.VisitType(t, q)); + } + } + + // Classpath entries are Class instances in practice; the key only feeds the sender's own diff. + private static object TypeKey(JavaType.FullyQualified type) => + type is JavaType.Class cls ? cls.FullyQualifiedName : type; + + public JavaSourceSet RpcReceive(JavaSourceSet before, RpcReceiveQueue q) + { + var typeReceiver = new JavaReceiver(); + var id = q.ReceiveAndGet(before.Id, Guid.Parse); + var name = q.Receive(before.Name)!; + var classpath = q.ReceiveList(before.Classpath, + t => (JavaType.FullyQualified)typeReceiver.VisitType(t, q)!)!; + + // The uninitialized instance the queue hands back on an ADD has null collections, so the + // before state is never dereferenced without a guard. + var beforeGavs = before.GavToTypes; + var gavs = q.ReceiveList(beforeGavs?.Keys.ToList(), null); + var gavToTypes = new Dictionary>(); + foreach (var gav in gavs ?? []) + { + IList? beforeBucket = null; + beforeGavs?.TryGetValue(gav, out beforeBucket); + gavToTypes[gav] = q.ReceiveList(beforeBucket, + t => (JavaType.FullyQualified)typeReceiver.VisitType(t, q)!)!; + } + return new JavaSourceSet(id, name, classpath, gavToTypes); + } +} diff --git a/rewrite-csharp/src/integTest/java/org/openrewrite/csharp/rpc/CSharpJavaSourceSetRpcTest.java b/rewrite-csharp/src/integTest/java/org/openrewrite/csharp/rpc/CSharpJavaSourceSetRpcTest.java new file mode 100644 index 00000000000..2d383a04cc9 --- /dev/null +++ b/rewrite-csharp/src/integTest/java/org/openrewrite/csharp/rpc/CSharpJavaSourceSetRpcTest.java @@ -0,0 +1,94 @@ +/* + * Copyright 2026 the original author or authors. + *

+ * Licensed under the Moderne Source Available License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + *

+ * https://docs.moderne.io/licensing/moderne-source-available-license + *

+ * 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.csharp.rpc; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.openrewrite.ExecutionContext; +import org.openrewrite.InMemoryExecutionContext; +import org.openrewrite.Tree; +import org.openrewrite.java.marker.JavaSourceSet; +import org.openrewrite.java.tree.JavaType; +import org.openrewrite.xml.XmlParser; +import org.openrewrite.xml.tree.Xml; + +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * {@link JavaSourceSet} rides on every source file in a Java source set, resource files included, + * so it reaches the C# peer through {@code Xml$Document} even though C# never sees a Java + * compilation unit. Without a codec on the C# side it resolves to {@code UnknownMarker}, which + * consumes one message where Java sends many, and the queue then desynchronizes. + */ +@Tag("slow") +@Timeout(value = 180, unit = TimeUnit.SECONDS) +class CSharpJavaSourceSetRpcTest { + + @BeforeAll + static void setUpFactory() { + Path basePath = Paths.get(System.getProperty("user.dir")); + Path[] searchPaths = { + basePath.resolve("csharp"), + basePath.resolve("rewrite-csharp/csharp"), + }; + for (Path searchPath : searchPaths) { + Path csproj = searchPath.resolve("OpenRewrite.Tool/OpenRewrite.Tool.csproj"); + if (csproj.toFile().exists()) { + CSharpRewriteRpc.setFactory( + CSharpRewriteRpc.builder() + .csharpServerEntry(csproj.toAbsolutePath().normalize()) + .log(Paths.get(System.getProperty("java.io.tmpdir"), "csharp-rpc-java-source-set.log"))); + return; + } + } + throw new IllegalStateException("Could not find C# Rewrite project"); + } + + @AfterAll + static void tearDown() { + CSharpRewriteRpc.shutdownCurrent(); + } + + @Test + void javaSourceSetMarkerAcrossRpcBoundary() { + ExecutionContext ctx = new InMemoryExecutionContext(); + Xml.Document document = XmlParser.builder().build() + .parse(ctx, "test") + .findFirst() + .map(Xml.Document.class::cast) + .orElseThrow(); + + JavaType.FullyQualified a = JavaType.ShallowClass.build("com.example.A"); + JavaType.FullyQualified b = JavaType.ShallowClass.build("com.example.B"); + JavaSourceSet sourceSet = new JavaSourceSet(Tree.randomId(), "main", + List.of(JavaType.ShallowClass.build("java.lang.String"), a, b), + Map.of("com.example:example:1.0", List.of(a, b))); + + String printed = CSharpRewriteRpc.getOrStart() + .print(document.withMarkers(document.getMarkers().add(sourceSet))); + + assertThat(printed).isEqualTo("test"); + } +} diff --git a/rewrite-java/src/main/java/org/openrewrite/java/marker/JavaSourceSet.java b/rewrite-java/src/main/java/org/openrewrite/java/marker/JavaSourceSet.java index 6a565149184..28a94958998 100644 --- a/rewrite-java/src/main/java/org/openrewrite/java/marker/JavaSourceSet.java +++ b/rewrite-java/src/main/java/org/openrewrite/java/marker/JavaSourceSet.java @@ -33,8 +33,13 @@ import org.openrewrite.SourceFile; import org.openrewrite.java.internal.JavaTypeCache; import org.openrewrite.java.internal.JavaTypeFactory; +import org.openrewrite.java.internal.rpc.JavaTypeReceiver; +import org.openrewrite.java.internal.rpc.JavaTypeSender; import org.openrewrite.java.tree.JavaType; import org.openrewrite.marker.SourceSet; +import org.openrewrite.rpc.RpcCodec; +import org.openrewrite.rpc.RpcReceiveQueue; +import org.openrewrite.rpc.RpcSendQueue; import java.beans.ConstructorProperties; import java.io.IOException; @@ -52,7 +57,7 @@ @Value @EqualsAndHashCode(onlyExplicitlyIncluded = true) @With -public class JavaSourceSet implements SourceSet { +public class JavaSourceSet implements SourceSet, RpcCodec { @EqualsAndHashCode.Include UUID id; @@ -436,6 +441,52 @@ private static List typesFrom(List typeNames) { return name; } + /** + * Field order is the protocol; every peer codec mirrors it. Without a codec this marker is + * inlined as a single {@link org.openrewrite.rpc.RpcObjectData} — for a real classpath, one + * JSON document of well over 100 MB, built and parsed whole, holding a second type universe + * the receiver cannot share with the {@link JavaType}s the tree already sent. + */ + @Override + public void rpcSend(JavaSourceSet after, RpcSendQueue q) { + JavaTypeSender typeSender = new JavaTypeSender(); + q.getAndSend(after, JavaSourceSet::getId); + q.getAndSend(after, JavaSourceSet::getName); + q.getAndSendListAsRef(after, JavaSourceSet::getClasspath, + JavaType.FullyQualified::getFullyQualifiedName, t -> typeSender.visit(t, q)); + + // getValueType returns null for a Map, so sending gavToTypes whole would inline the entire + // type graph a second time. Keys first, then each bucket as refs: the values are the same + // instances that are already on the classpath above (see build), so every one is a ref hit. + q.getAndSendList(after, a -> new ArrayList<>(a.getGavToTypes().keySet()), s -> s, null); + for (String gav : after.getGavToTypes().keySet()) { + q.getAndSendListAsRef(after, a -> a.getGavToTypes().get(gav), + JavaType.FullyQualified::getFullyQualifiedName, t -> typeSender.visit(t, q)); + } + } + + @Override + public JavaSourceSet rpcReceive(JavaSourceSet before, RpcReceiveQueue q) { + JavaTypeReceiver typeReceiver = new JavaTypeReceiver(); + JavaSourceSet after = before + .withId(q.receiveAndGet(before.getId(), UUID::fromString)) + .withName(q.receive(before.getName())) + .withClasspath(q.receiveList(before.getClasspath(), + t -> (JavaType.FullyQualified) typeReceiver.visit(t, q))); + + // On an ADD Objenesis leaves every field null, so the before map is never dereferenced. + Map> beforeGavs = before.getGavToTypes(); + List gavs = q.receiveList(beforeGavs == null ? null : new ArrayList<>(beforeGavs.keySet()), null); + Map> gavToTypes = new LinkedHashMap<>(); + if (gavs != null) { + for (String gav : gavs) { + gavToTypes.put(gav, q.receiveList(beforeGavs == null ? null : beforeGavs.get(gav), + t -> (JavaType.FullyQualified) typeReceiver.visit(t, q))); + } + } + return after.withGavToTypes(gavToTypes); + } + // Purely IO-based classpath scanning below this point diff --git a/rewrite-java/src/test/java/org/openrewrite/java/marker/JavaSourceSetRpcTest.java b/rewrite-java/src/test/java/org/openrewrite/java/marker/JavaSourceSetRpcTest.java new file mode 100644 index 00000000000..0b2fe8b0c54 --- /dev/null +++ b/rewrite-java/src/test/java/org/openrewrite/java/marker/JavaSourceSetRpcTest.java @@ -0,0 +1,114 @@ +/* + * Copyright 2026 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.java.marker; + +import org.junit.jupiter.api.Test; +import org.openrewrite.java.tree.JavaType; +import org.openrewrite.rpc.Reference; +import org.openrewrite.rpc.RpcObjectData; +import org.openrewrite.rpc.RpcReceiveQueue; +import org.openrewrite.rpc.RpcSendQueue; + +import java.util.*; + +import static java.util.Arrays.asList; +import static org.assertj.core.api.Assertions.assertThat; + +class JavaSourceSetRpcTest { + + private static final String GAV = "com.example:example:1.0"; + + @Test + void roundTrip() { + JavaSourceSet before = sourceSet(); + + JavaSourceSet received = sendAndReceive(before); + + assertThat(received.getId()).isEqualTo(before.getId()); + assertThat(received.getName()).isEqualTo("main"); + assertThat(received.getClasspath()).extracting(JavaType.FullyQualified::getFullyQualifiedName) + .containsExactly("java.lang.String", "com.example.A", "com.example.B"); + assertThat(received.getGavToTypes()).containsOnlyKeys(GAV); + assertThat(received.getGavToTypes().get(GAV)).extracting(JavaType.FullyQualified::getFullyQualifiedName) + .containsExactly("com.example.A", "com.example.B"); + assertThat(received.getTypeFactory()) + .as("transient and never sent, so the receiver falls back to a fresh factory") + .isNull(); + } + + @Test + void gavBucketsAreRefsIntoTheClasspath() { + JavaSourceSet received = sendAndReceive(sourceSet()); + + // The whole point of decomposing gavToTypes rather than sending the map inline: its values + // are the same instances that are already on the classpath, so the receiver must end up + // with one type universe rather than two. + List bucket = received.getGavToTypes().get(GAV); + assertThat(bucket.get(0)).isSameAs(received.getClasspath().get(1)); + assertThat(bucket.get(1)).isSameAs(received.getClasspath().get(2)); + } + + @Test + void aRepeatedSendCollapsesToARefOnlyAdd() { + JavaSourceSet sourceSet = sourceSet(); + Deque> batches = new ArrayDeque<>(); + RpcSendQueue sq = new RpcSendQueue(1_000_000, batches::addLast, new IdentityHashMap<>(), null, false); + + sq.send(Reference.asRef(sourceSet), null, null); + sq.flush(); + assertThat(batches.removeFirst()).hasSizeGreaterThan(1); + + // Markers travel asRef, so once the ref map survives eviction the second source file + // carrying this same instance pays one message instead of the whole classpath again. + sq.send(Reference.asRef(sourceSet), null, null); + sq.flush(); + List second = batches.removeFirst(); + assertThat(second).hasSize(1); + assertThat(second.get(0).getRef()).isNotNull(); + assertThat(second.get(0).getValueType()).isNull(); + assertThat((Object) second.get(0).getValue()).isNull(); + } + + private static JavaSourceSet sourceSet() { + JavaType.FullyQualified string = JavaType.ShallowClass.build("java.lang.String"); + JavaType.FullyQualified a = JavaType.ShallowClass.build("com.example.A"); + JavaType.FullyQualified b = JavaType.ShallowClass.build("com.example.B"); + Map> gavToTypes = new LinkedHashMap<>(); + gavToTypes.put(GAV, asList(a, b)); + return new JavaSourceSet(UUID.randomUUID(), "main", asList(string, a, b), gavToTypes); + } + + private static JavaSourceSet sendAndReceive(JavaSourceSet sourceSet) { + Deque> batches = new ArrayDeque<>(); + RpcSendQueue sq = new RpcSendQueue(1_000_000, batches::addLast, new IdentityHashMap<>(), null, false); + sq.send(sourceSet, null, null); + sq.flush(); + + // The wire carries a UUID as a string, as the transport's JSON encoding would. + List all = new ArrayList<>(); + while (!batches.isEmpty()) { + for (RpcObjectData data : batches.removeFirst()) { + all.add(data.getValue() instanceof UUID ? + new RpcObjectData(data.getState(), data.getValueType(), data.getValue().toString(), data.getRef(), false) : + data); + } + } + Deque> drain = new ArrayDeque<>(); + drain.add(all); + + return new RpcReceiveQueue(new HashMap<>(), drain::removeFirst, null, null).receive(null); + } +} diff --git a/rewrite-javascript/rewrite/src/java/markers.ts b/rewrite-javascript/rewrite/src/java/markers.ts index 7dde8eadaf7..e8c95e87ca0 100644 --- a/rewrite-javascript/rewrite/src/java/markers.ts +++ b/rewrite-javascript/rewrite/src/java/markers.ts @@ -15,10 +15,11 @@ */ import {Marker} from "../markers"; import {J} from "./tree"; -import {RpcCodecs, RpcReceiveQueue, RpcSendQueue} from "../rpc"; +import {Type} from "./type"; +import {asRef, RpcCodecs, RpcReceiveQueue, RpcSendQueue} from "../rpc"; import {updateIfChanged} from "../util"; // The `RpcCodec` for `J.Space` is registered in the `rpc` module. -import "./rpc"; +import {TypeReceiver, TypeSender} from "./rpc"; declare module "./tree" { namespace J { @@ -26,6 +27,7 @@ declare module "./tree" { readonly Semicolon: "org.openrewrite.java.marker.Semicolon"; readonly TrailingComma: "org.openrewrite.java.marker.TrailingComma"; readonly OmitParentheses: "org.openrewrite.java.marker.OmitParentheses"; + readonly JavaSourceSet: "org.openrewrite.java.marker.JavaSourceSet"; }; } } @@ -34,7 +36,8 @@ declare module "./tree" { (J as any).Markers = { Semicolon: "org.openrewrite.java.marker.Semicolon", TrailingComma: "org.openrewrite.java.marker.TrailingComma", - OmitParentheses: "org.openrewrite.java.marker.OmitParentheses" + OmitParentheses: "org.openrewrite.java.marker.OmitParentheses", + JavaSourceSet: "org.openrewrite.java.marker.JavaSourceSet" } as const; export interface Semicolon extends Marker { @@ -50,6 +53,17 @@ export interface OmitParentheses extends Marker { readonly kind: typeof J.Markers.OmitParentheses; } +/** + * Rides on every source file in a Java source set, resource files included, so it reaches this + * peer through PlainText/JSON/YAML as well as JavaScript. + */ +export interface JavaSourceSet extends Marker { + readonly kind: typeof J.Markers.JavaSourceSet; + readonly name: string; + readonly classpath: Type.FullyQualified[]; + readonly gavToTypes: { [gav: string]: Type.FullyQualified[] }; +} + // Register codecs for all Java markers with additional properties RpcCodecs.registerCodec(J.Markers.TrailingComma, { async rpcReceive(before: TrailingComma, q: RpcReceiveQueue): Promise { @@ -65,6 +79,39 @@ RpcCodecs.registerCodec(J.Markers.TrailingComma, { } }); +// Field order mirrors org.openrewrite.java.marker.JavaSourceSet#rpcSend, which is the canonical +// protocol. gavToTypes travels as a key list plus one ref-deduplicated bucket per key, so its +// values resolve to the classpath instances sent above rather than a second copy of the graph. +RpcCodecs.registerCodec(J.Markers.JavaSourceSet, { + async rpcReceive(before: JavaSourceSet, q: RpcReceiveQueue): Promise { + const typeReceiver = new TypeReceiver(); + const id = await q.receive(before.id); + const name = await q.receive(before.name); + const classpath = await q.receiveList(before.classpath, + t => typeReceiver.visit(t, q) as Promise); + const gavs = await q.receiveList(before.gavToTypes && Object.keys(before.gavToTypes)); + const gavToTypes: { [gav: string]: Type.FullyQualified[] } = {}; + for (const gav of gavs || []) { + gavToTypes[gav] = (await q.receiveList(before.gavToTypes?.[gav], + t => typeReceiver.visit(t, q) as Promise))!; + } + return updateIfChanged(before, {id, name, classpath, gavToTypes}); + }, + + async rpcSend(after: JavaSourceSet, q: RpcSendQueue): Promise { + const typeSender = new TypeSender(); + await q.getAndSend(after, a => a.id); + await q.getAndSend(after, a => a.name); + await q.getAndSendList(after, a => (a.classpath || []).map(t => asRef(t)), + t => Type.signature(t), t => typeSender.visit(t, q)); + await q.getAndSendList(after, a => Object.keys(a.gavToTypes || {}), gav => gav); + for (const gav of Object.keys(after.gavToTypes || {})) { + await q.getAndSendList(after, a => (a.gavToTypes[gav] || []).map(t => asRef(t)), + t => Type.signature(t), t => typeSender.visit(t, q)); + } + } +}); + /** * Registers an RPC codec for any marker without additional properties. */ diff --git a/rewrite-javascript/rewrite/src/java/rpc.ts b/rewrite-javascript/rewrite/src/java/rpc.ts index 0fd33e552c7..b6471ca411a 100644 --- a/rewrite-javascript/rewrite/src/java/rpc.ts +++ b/rewrite-javascript/rewrite/src/java/rpc.ts @@ -22,7 +22,7 @@ import {TypeVisitor} from "./type-visitor"; import {updateIfChanged} from "../util"; import Space = J.Space; -class TypeSender extends TypeVisitor { +export class TypeSender extends TypeVisitor { protected async visitPrimitive(primitive: Type.Primitive, q: RpcSendQueue): Promise { await q.getAndSend(primitive, p => p.keyword); return primitive; @@ -138,7 +138,7 @@ class TypeSender extends TypeVisitor { } } -class TypeReceiver extends TypeVisitor { +export class TypeReceiver extends TypeVisitor { async preVisit(_type: Type, _q: RpcReceiveQueue): Promise { // Don't call default preVisit to avoid circular references return _type; diff --git a/rewrite-javascript/src/integTest/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpcTest.java b/rewrite-javascript/src/integTest/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpcTest.java index d159236f75f..6cf1ce143ea 100644 --- a/rewrite-javascript/src/integTest/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpcTest.java +++ b/rewrite-javascript/src/integTest/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpcTest.java @@ -28,6 +28,7 @@ import org.openrewrite.internal.RecipeLoader; import org.openrewrite.java.JavaIsoVisitor; import org.openrewrite.java.JavaVisitor; +import org.openrewrite.java.marker.JavaSourceSet; import org.openrewrite.java.tree.J; import org.openrewrite.java.tree.JavaType; import org.openrewrite.javascript.JavaScriptIsoVisitor; @@ -284,6 +285,28 @@ void printText() { ); } + @Test + void javaSourceSetMarkerAcrossRpcBoundary() { + rewriteRun( + text( + "Hello Jon!", + spec -> spec.beforeRecipe(text -> { + JavaType.FullyQualified a = JavaType.ShallowClass.build("com.example.A"); + JavaType.FullyQualified b = JavaType.ShallowClass.build("com.example.B"); + JavaSourceSet sourceSet = new JavaSourceSet(Tree.randomId(), "main", + List.of(JavaType.ShallowClass.build("java.lang.String"), a, b), + Map.of("com.example:example:1.0", List.of(a, b))); + + // Resource files in a Java source set carry this marker, so it crosses to a peer + // that accepts PlainText/JSON/YAML even though that peer never sees a Java CU. + // Without a codec on both sides the queue desynchronizes here. + assertThat(client().print(text.withMarkers(text.getMarkers().add(sourceSet)))) + .isEqualTo("Hello Jon!"); + }) + ) + ); + } + @Test void printFencedMarker() { rewriteRun(