@openrewrite/rewrite 8.87.3 → 8.87.5

package/src/rpc/rewrite-rpc.ts

@@ -38,7 +38,7 @@ import {
 } from "./request";
 import {DataTableStore} from "../data-table";
 import {RecipeMarketplace} from "../marketplace";
-import {initializeMetricsCsv} from "./request/metrics";
+import {initializeMetricsCsv, setCacheSizeProvider} from "./request/metrics";
 import {RpcObjectData, RpcObjectState, RpcReceiveQueue} from "./queue";
 import {RpcRecipe} from "./recipe";
 import {ExecutionContext} from "../execution";
@@ -74,6 +74,10 @@ export class RewriteRpc {
     readonly remoteRefs: Map<number, any> = 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<string, { send: number, recvMax: number }> = new Map();
+
     private remoteLanguages?: string[];
     private readonly logger?: rpc.Logger;
     private traceGetObject: TraceGetObject = {receive: false, send: false};
@@ -93,6 +97,11 @@ export class RewriteRpc {
                 }) {
         // Initialize metrics CSV file if configured
         initializeMetricsCsv(options.metricsCsv, options.logger);
+        setCacheSizeProvider(() => ({
+            local: this.localObjects.size,
+            remote: this.remoteObjects.size,
+            refs: this.remoteRefs.size + this.localRefs.size,
+        }));
         this.logger = options.logger;
 
         const preparedRecipes: Map<String, Recipe> = new Map();
@@ -101,6 +110,19 @@ 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;
 
@@ -109,8 +131,8 @@ export class RewriteRpc {
         // GetMarketplace builds rows so the host can attribute each recipe to its own bundle.
         const recipeOrigin: Map<string, string> = new Map();
 
-        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);
+        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);
         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,
@@ -140,6 +162,7 @@ export class RewriteRpc {
             this.remoteObjects.clear();
             this.remoteRefs.clear();
             this.localRefs.clear();
+            this.refCheckpoints.clear();
             preparedRecipes.clear();
             this.remoteLanguages = undefined;
         };
@@ -156,6 +179,27 @@ 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.
+        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);
+                }
+            }
+        )
+
         RewriteRpc.set(this);
         this.connection.listen();
     }