@openrewrite/rewrite 8.92.15 → 8.92.17

package/src/rpc/queue.ts

@@ -182,6 +182,13 @@ export class RpcSendQueue {
 
     send<T>(after: T | undefined, before: T | undefined, onChange?: (() => Promise<any>)): Promise<void> {
         return saveTrace(this.trace, async () => {
+            // The peer has a single absent value, and an ADD that carries nothing is not one it can decode.
+            if (after === null) {
+                after = undefined;
+            }
+            if (before === null) {
+                before = undefined;
+            }
             if (before === after) {
                 this.put({state: RpcObjectState.NO_CHANGE});
             } else if (before === undefined || (after !== undefined && this.typesAreDifferent(after, before))) {
@@ -372,10 +379,20 @@ export class StringInternTable {
     }
 }
 
+// Java's serializer names the class and numbers the object inside the value; the kind already says the first.
+function withoutSerializerKeys(value: any): any {
+    if (value === null || typeof value !== "object") {
+        return value;
+    }
+    const {"@c": _class, "@ref": _ref, ...fields} = value;
+    return fields;
+}
+
 export class RpcReceiveQueue {
     private batch: RpcObjectData[] = [];
     private batchIndex = 0;
     private sinceYield = 0;
+    private recorded: number[] = [];
 
     constructor(private readonly refs: Map<number, any>,
                 private readonly sourceFileType: string | undefined,
@@ -385,6 +402,23 @@ export class RpcReceiveQueue {
                 private readonly internedStrings: StringInternTable = new StringInternTable()) {
     }
 
+    private record(ref: number, value: any): void {
+        this.refs.set(ref, value);
+        this.recorded.push(ref);
+    }
+
+    /**
+     * Drops the refs this transfer defined, for when it could not be read to its end and the
+     * sender is about to forget them too.
+     * @internal
+     */
+    forgetRefs(): void {
+        for (const ref of this.recorded) {
+            this.refs.delete(ref);
+        }
+        this.recorded = [];
+    }
+
     /**
      * Returns a value rather than a promise for all but the message that refills the
      * batch or yields, so the common path costs neither a promise nor an async frame.
@@ -479,7 +513,7 @@ export class RpcReceiveQueue {
                         // For an object like JavaType that we will mutate in place rather than using
                         // immutable updates because of its cyclic nature, the before instance will ultimately
                         // be the same as the after instance below.
-                        this.refs.set(ref, before);
+                        this.record(ref, before);
                     }
                 }
             // Intentional fall-through...
@@ -492,7 +526,7 @@ export class RpcReceiveQueue {
                     after = await codec.rpcReceive(before, this);
                 } else if (message.value !== undefined) {
                     after = message.valueType ?
-                        {kind: this.internedStrings.internType(message.valueType), ...message.value} :
+                        {kind: this.internedStrings.internType(message.valueType), ...withoutSerializerKeys(message.value)} :
                         typeof message.value === "string" ? this.internedStrings.internValue(message.value) : message.value;
                 } else if (message.state === RpcObjectState.ADD && message.valueType) {
                     throw new Error(
@@ -504,7 +538,7 @@ export class RpcReceiveQueue {
                     after = before;
                 }
                 if (ref !== undefined) {
-                    this.refs.set(ref, after);
+                    this.record(ref, after);
                 }
                 return after;
             default:

package/src/rpc/request/get-object.ts

@@ -18,6 +18,11 @@ import {RpcObjectData, RpcObjectState, RpcSendQueue} from "../queue";
 import {ReferenceMap} from "../../reference";
 import {extractSourcePath, withMetrics} from "./metrics";
 
+export class AbortGetObject {
+    constructor(readonly id: string) {
+    }
+}
+
 export class GetObject {
     constructor(private readonly id: string,
                 private readonly sourceFileType?: string) {
@@ -33,6 +38,24 @@ export class GetObject {
         metricsCsv?: string,
     ): void {
         const pendingData = new Map<string, { data: (RpcObjectData | undefined)[], offset: number }>();
+        let latestTransfer: { id: string, refCount: number } | undefined;
+
+        // The receiver could not read a transfer to its end and has dropped what it took from it.
+        connection.onRequest(
+            new rpc.RequestType<AbortGetObject, boolean, Error>("AbortGetObject"),
+            request => {
+                pendingData.delete(request.id);
+                remoteObjects.delete(request.id);
+                if (latestTransfer?.id === request.id) {
+                    localRefs.rollbackTo(latestTransfer.refCount);
+                } else {
+                    // which refs that transfer assigned is no longer known; everything goes out whole again
+                    localRefs.clear();
+                }
+                latestTransfer = undefined;
+                return true;
+            }
+        );
 
         connection.onRequest(
             new rpc.RequestType<GetObject, any, Error>("GetObject"),
@@ -75,6 +98,7 @@ export class GetObject {
                             };
                             pendingData.set(objId, pending);
                             remoteObjects.set(objId, after);
+                            latestTransfer = {id: objId, refCount: savedRefCount};
                         } catch (e) {
                             remoteObjects.delete(objId);
                             localRefs.rollbackTo(savedRefCount);

package/src/rpc/request/parse.ts

@@ -67,7 +67,10 @@ export class Parse {
                         relativeTo: request.relativeTo
                     })!;
 
-                    const generator = parser.parse(...request.inputs);
+                    // an input without its text names a file for the parser to read
+                    const inputs = request.inputs.map(input =>
+                        typeof input === 'object' && input.text == null ? input.sourcePath : input);
+                    const generator = parser.parse(...inputs);
                     const resultIds: UUID[] = [];
 
                     for (let i = 0; i < request.inputs.length; i++) {

package/src/rpc/request/print.ts

@@ -14,25 +14,31 @@
  * limitations under the License.
  */
 import * as rpc from "vscode-jsonrpc/node";
-import {isSourceFile, Tree} from "../../tree";
+import {Cursor, isSourceFile, Tree} from "../../tree";
 import {MarkerPrinter as PrintMarkerPrinter, printer, PrintOutputCapture} from "../../print";
 import {UUID} from "../../uuid";
 import {extractSourcePath, withMetrics} from "./metrics";
 
 export const enum MarkerPrinter {
     DEFAULT = "DEFAULT",
+    SEARCH_MARKERS_ONLY = "SEARCH_MARKERS_ONLY",
     FENCED = "FENCED",
     SANITIZED = "SANITIZED"
 }
 
 export class Print {
+    /**
+     * @param cursor The ids of what encloses a tree that is not a source file, from its parent outward.
+     */
     constructor(private readonly treeId: UUID,
                 private readonly sourceFileType: string,
-                readonly markerPrinter: MarkerPrinter = MarkerPrinter.DEFAULT) {
+                readonly markerPrinter: MarkerPrinter = MarkerPrinter.DEFAULT,
+                readonly cursor?: string[]) {
     }
 
     static handle(connection: rpc.MessageConnection,
                   getObject: (id: string, sourceFileType: string) => any,
+                  getCursor: (cursorIds: string[] | undefined, sourceFileType?: string) => Promise<Cursor>,
                   logger?: rpc.Logger,
                   metricsCsv?: string): void {
         connection.onRequest(
@@ -46,14 +52,12 @@ export class Print {
                     if (logger) {
                         logger.log("Printing " + (isSourceFile(tree) ? tree.sourcePath : `tree of type ${tree.kind} in ${context.target}`));
                     }
-                    const out = new PrintOutputCapture(PrintMarkerPrinter[request.markerPrinter]);
-                    let result: string;
+                    const out = new PrintOutputCapture(PrintMarkerPrinter[request.markerPrinter ?? MarkerPrinter.DEFAULT]);
                     if (isSourceFile(tree)) {
-                        result = await printer(tree).print(tree, out);
-                    } else {
-                        result = await printer(request.sourceFileType).print(tree, out);
+                        return printer(tree).print(tree, out);
                     }
-                    return result;
+                    const cursor = request.cursor ? await getCursor(request.cursor, request.sourceFileType) : undefined;
+                    return printer(request.sourceFileType).print(tree, out, cursor);
                 }
             )
         );

package/src/rpc/rewrite-rpc.ts

@@ -22,6 +22,7 @@ import {
     DependencyTypes,
     Generate,
     GenerateResponse,
+    AbortGetObject,
     GetObject,
     GetMarketplace,
     GetMarketplaceResponseRow,
@@ -30,6 +31,7 @@ import {
     ParseProject,
     PrepareRecipe,
     PrepareRecipeResponse,
+    MarkerPrinter,
     Print,
     TraceGetObject,
     Visit,
@@ -142,7 +144,7 @@ export class RewriteRpc {
         Parse.handle(this.connection, this.localObjects, options.metricsCsv);
         ParseProject.handle(this.connection, this.localObjects, options.metricsCsv);
         DependencyTypes.handle(this.connection, options?.batchSize || 1000, options.metricsCsv);
-        Print.handle(this.connection, getObject, options.logger, options.metricsCsv);
+        Print.handle(this.connection, getObject, getCursor, options.logger, options.metricsCsv);
         InstallRecipes.handle(this.connection, options.recipeInstallDir ?? ".rewrite", marketplace, recipeOrigin, options.logger, options.metricsCsv);
 
         this.connection.onRequest(
@@ -289,6 +291,15 @@ export class RewriteRpc {
                 await nextPage.catch(() => {
                 });
             }
+            // The sender still counts the object and its refs as delivered until it is told otherwise.
+            const rolledBack = await this.connection.sendRequest(
+                new rpc.RequestType<AbortGetObject, boolean, Error>("AbortGetObject"),
+                new AbortGetObject(id)
+            ).catch(() => false);
+            if (rolledBack) {
+                // a peer that could not roll back goes on sending these refs bare, so they are kept for it
+                q.forgetRefs();
+            }
             throw e;
         }
 
@@ -321,9 +332,9 @@ export class RewriteRpc {
         return parsed;
     }
 
-    async print(tree: SourceFile): Promise<string>;
-    async print(tree: Tree, cursor: Cursor): Promise<string>;
-    async print(tree: Tree, cursor?: Cursor): Promise<string> {
+    async print(tree: SourceFile, cursor?: undefined, markerPrinter?: MarkerPrinter): Promise<string>;
+    async print(tree: Tree, cursor: Cursor, markerPrinter?: MarkerPrinter): Promise<string>;
+    async print(tree: Tree, cursor?: Cursor, markerPrinter?: MarkerPrinter): Promise<string> {
         if (!cursor && !isSourceFile(tree)) {
             throw new Error("Cursor is required for non-SourceFile trees");
         }
@@ -331,7 +342,7 @@ export class RewriteRpc {
         const sourceFile = isSourceFile(tree) ? tree : cursor!.firstEnclosing(t => isSourceFile(t))!;
         return await this.connection.sendRequest(
             new rpc.RequestType<Print, string, Error>("Print"),
-            new Print(tree.id, sourceFile.kind)
+            new Print(tree.id, sourceFile.kind, markerPrinter, isSourceFile(tree) ? undefined : this.getCursorIds(cursor))
         );
     }
 

package/src/rpc/server.ts

@@ -22,6 +22,7 @@ import * as fs from "fs";
 import {Command} from 'commander';
 import {dir} from 'tmp-promise';
 import {DependencyWorkspace} from "../javascript/dependency-workspace";
+import {isMainThread, parentPort, Worker} from "worker_threads";
 
 // Include all languages you want this server to support.
 import "../text";
@@ -30,9 +31,6 @@ import "../yaml";
 import "../java";
 import "../javascript";
 
-// Not possible to set the stack size when executing from npx for security reasons
-require('v8').setFlagsFromString('--stack-size=8000');
-
 function initPyroscope(logger: rpc.Logger): any {
     // Strip trailing slashes: the SDK builds the ingest URL as `${serverAddress}/ingest`,
     // so a trailing slash produces `//ingest`, which the server normalizes via redirect —
@@ -126,8 +124,8 @@ async function main() {
         }
     };
 
-    process.on('SIGINT', shutdown);
-    process.on('SIGTERM', shutdown);
+    // signals are delivered to the main thread, which passes them on
+    parentPort!.on('message', shutdown);
 
     const log = options.logFile ? fs.createWriteStream(options.logFile, {flags: 'a'}) : undefined;
     const logger: rpc.Logger = {
@@ -191,4 +189,18 @@ async function main() {
     });
 }
 
-main().catch(console.error);
+if (isMainThread) {
+    // A thread's stack is sized when it starts, and the default is too shallow for the type checker
+    // on long chains of inferred return types, so the server runs on a thread started with more.
+    const server = new Worker(__filename, {
+        argv: process.argv.slice(2),
+        stdin: true,
+        resourceLimits: {stackSizeMb: 8}
+    });
+    process.stdin.pipe(server.stdin!);
+    process.on('SIGINT', () => server.postMessage('SIGINT'));
+    process.on('SIGTERM', () => server.postMessage('SIGTERM'));
+    server.on('exit', code => process.exit(code));
+} else {
+    main().catch(console.error);
+}

package/src/test/java-rpc.ts

@@ -39,6 +39,7 @@
  */
 import {describe, type SuiteAPI, test as base} from "vitest";
 import {findTestClasspath, JavaRpcTestServer} from "../rpc/java-rpc-client";
+import {RewriteRpc} from "../rpc/rewrite-rpc";
 
 export {findTestClasspath, JavaRpcTestServer} from "../rpc/java-rpc-client";
 export type {JavaRpcOptions} from "../rpc/java-rpc-client";
@@ -76,6 +77,8 @@ export const testJavaRpc = base.extend<JavaRpcFixtures>({
         {scope: "worker"},
     ],
     javaRpc: async ({_javaRpcServer}, use) => {
+        // any RewriteRpc constructed in this worker since has replaced it as the active connection
+        RewriteRpc.set(_javaRpcServer.rpc);
         await _javaRpcServer.reset();
         await use(_javaRpcServer);
     },