@openrewrite/rewrite 8.92.16 → 8.92.17

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);
     },