Skip to content

Commit 22eb85c

Browse files
authored
Javascript RPC metrics (#6118)
1 parent 047cb44 commit 22eb85c

14 files changed

Lines changed: 266 additions & 45 deletions

File tree

rewrite-javascript/rewrite/src/rpc/request/generate.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import * as rpc from "vscode-jsonrpc/node";
1717
import {Recipe, ScanningRecipe} from "../../recipe";
1818
import {Cursor, rootCursor} from "../../tree";
1919
import {ExecutionContext} from "../../execution";
20+
import {withMetrics} from "./metrics";
2021

2122
export interface GenerateResponse {
2223
ids: string[]
@@ -31,8 +32,11 @@ export class Generate {
3132
localObjects: Map<string, any>,
3233
preparedRecipes: Map<String, Recipe>,
3334
recipeCursors: WeakMap<Recipe, Cursor>,
34-
getObject: (id: string) => any): void {
35-
connection.onRequest(new rpc.RequestType<Generate, GenerateResponse, Error>("Generate"), async (request) => {
35+
getObject: (id: string) => any,
36+
metricsCsv?: string): void {
37+
const target = { target: '' };
38+
connection.onRequest(new rpc.RequestType<Generate, GenerateResponse, Error>("Generate"), withMetrics<Generate, GenerateResponse>("Generate", target, metricsCsv)(async (request) => {
39+
target.target = request.id;
3640
const recipe = preparedRecipes.get(request.id);
3741
const response = {
3842
ids: [],
@@ -57,6 +61,6 @@ export class Generate {
5761

5862
}
5963
return response;
60-
});
64+
}));
6165
}
6266
}
Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
import * as rpc from "vscode-jsonrpc/node";
2+
import {withMetrics0} from "./metrics";
23

34
export class GetLanguages {
4-
static handle(connection: rpc.MessageConnection): void {
5-
connection.onRequest(new rpc.RequestType0<string[], Error>("GetLanguages"), async () => {
5+
static handle(connection: rpc.MessageConnection, metricsCsv?: string): void {
6+
const target = {target: ''};
7+
connection.onRequest(new rpc.RequestType0<string[], Error>("GetLanguages"), withMetrics0<string[]>("GetLanguages", target, metricsCsv)(async () => {
68
// Include all languages you want this server to support receiving from a remote peer
7-
const languages: string[] = [
9+
return [
810
"org.openrewrite.text.PlainText",
911
"org.openrewrite.json.tree.Json$Document",
1012
"org.openrewrite.java.tree.J$CompilationUnit",
1113
"org.openrewrite.javascript.tree.JS$CompilationUnit",
1214
];
13-
return languages;
14-
});
15+
}));
1516
}
1617
}

rewrite-javascript/rewrite/src/rpc/request/get-object.ts

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import * as rpc from "vscode-jsonrpc/node";
1717
import {RpcObjectData, RpcObjectState, RpcSendQueue} from "../queue";
1818
import {ReferenceMap} from "../../reference";
19+
import {withMetrics, extractSourcePath} from "./metrics";
1920

2021
export class GetObject {
2122
constructor(private readonly id: string,
@@ -28,11 +29,13 @@ export class GetObject {
2829
localObjects: Map<string, any | ((input: string) => any)>,
2930
localRefs: ReferenceMap,
3031
batchSize: number,
31-
trace: boolean
32+
trace: boolean,
33+
metricsCsv?: string
3234
): void {
3335
const pendingData = new Map<string, RpcObjectData[]>();
36+
const target = { target: '' };
3437

35-
connection.onRequest(new rpc.RequestType<GetObject, any, Error>("GetObject"), async request => {
38+
connection.onRequest(new rpc.RequestType<GetObject, any, Error>("GetObject"), withMetrics<GetObject, any>("GetObject", target, metricsCsv)(async request => {
3639
let objId = request.id;
3740
if (!localObjects.has(objId)) {
3841
return [
@@ -47,9 +50,12 @@ export class GetObject {
4750
localObjects.set(objId, obj);
4851
}
4952

53+
const obj = localObjects.get(objId);
54+
target.target = extractSourcePath(obj);
55+
5056
let allData = pendingData.get(objId);
5157
if (!allData) {
52-
const after = localObjects.get(objId);
58+
const after = obj;
5359
const before = remoteObjects.get(objId);
5460

5561
allData = await new RpcSendQueue(localRefs, request.sourceFileType, trace).generate(after, before);
@@ -66,6 +72,6 @@ export class GetObject {
6672
}
6773

6874
return batch;
69-
});
75+
}));
7076
}
7177
}

rewrite-javascript/rewrite/src/rpc/request/get-recipes.ts

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,17 @@
1515
*/
1616
import * as rpc from "vscode-jsonrpc/node";
1717
import {RecipeDescriptor, RecipeRegistry} from "../../recipe";
18+
import {withMetrics0} from "./metrics";
1819

1920
export class GetRecipes {
20-
static handle(connection: rpc.MessageConnection, registry: RecipeRegistry): void {
21-
connection.onRequest(new rpc.RequestType0<({ name: string } & RecipeDescriptor)[], Error>("GetRecipes"), async () => {
21+
static handle(connection: rpc.MessageConnection, registry: RecipeRegistry, metricsCsv?: string): void {
22+
const target = { target: '' };
23+
connection.onRequest(new rpc.RequestType0<({ name: string } & RecipeDescriptor)[], Error>("GetRecipes"), withMetrics0<({ name: string } & RecipeDescriptor)[]>("GetRecipes", target, metricsCsv)(async () => {
2224
const recipes = [];
2325
for (const [_name, recipe] of registry.all.entries()) {
2426
recipes.push(await new recipe().descriptor());
2527
}
2628
return recipes;
27-
});
29+
}));
2830
}
2931
}

rewrite-javascript/rewrite/src/rpc/request/install-recipes.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {RecipeRegistry} from "../../recipe";
1818
import * as path from "path";
1919
import * as fs from "fs";
2020
import {spawn, ChildProcess} from "child_process";
21+
import {withMetrics} from "./metrics";
2122

2223
export interface InstallRecipesResponse {
2324
recipesInstalled: number
@@ -68,8 +69,10 @@ export class InstallRecipes {
6869
}
6970

7071
static handle(connection: rpc.MessageConnection, installDir: string, registry: RecipeRegistry,
71-
logger?: rpc.Logger): void {
72-
connection.onRequest(new rpc.RequestType<InstallRecipes, InstallRecipesResponse, Error>("InstallRecipes"), async (request) => {
72+
logger?: rpc.Logger, metricsCsv?: string): void {
73+
const target = { target: '' };
74+
connection.onRequest(new rpc.RequestType<InstallRecipes, InstallRecipesResponse, Error>("InstallRecipes"), withMetrics<InstallRecipes, InstallRecipesResponse>("InstallRecipes", target, metricsCsv)(async (request) => {
75+
target.target = typeof request.recipes === "object" ? request.recipes.packageName : request.recipes;
7376
const beforeInstall = registry.all.size;
7477
let resolvedPath;
7578
let recipesName = request.recipes;
@@ -129,7 +132,7 @@ export class InstallRecipes {
129132
}
130133

131134
return {recipesInstalled: registry.all.size - beforeInstall};
132-
});
135+
}));
133136
}
134137
}
135138

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
* Copyright 2025 the original author or authors.
3+
* <p>
4+
* Licensed under the Moderne Source Available License (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
* <p>
8+
* https://docs.moderne.io/licensing/moderne-source-available-license
9+
* <p>
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
import * as fs from 'fs';
17+
import * as rpc from "vscode-jsonrpc/node";
18+
import {Cursor, isSourceFile, SourceFile} from "../../tree";
19+
20+
const CSV_HEADER = 'request,target,durationMs,memoryUsedBytes,memoryMaxBytes';
21+
22+
/**
23+
* Extracts the sourcePath from a tree object, either directly from a SourceFile
24+
* or by finding the nearest SourceFile via a cursor.
25+
*
26+
* @param tree The tree object to extract sourcePath from
27+
* @param cursor Optional cursor to find the nearest SourceFile
28+
* @returns The sourcePath or empty string if not found
29+
*/
30+
export function extractSourcePath(tree: any, cursor?: Cursor): string {
31+
if (isSourceFile(tree)) {
32+
return (tree as SourceFile).sourcePath || '';
33+
}
34+
35+
if (cursor) {
36+
const sourceFile = cursor.firstEnclosing(t => isSourceFile(t));
37+
if (sourceFile) {
38+
return (sourceFile as SourceFile).sourcePath || '';
39+
}
40+
}
41+
42+
return '';
43+
}
44+
45+
/**
46+
* Initializes the metrics CSV file with a header if it doesn't already exist or is empty.
47+
* If the file already contains the correct header, this function is a no-op.
48+
*
49+
* @param metricsCsv The path to the CSV file for recording metrics
50+
* @param logger Optional logger for warnings
51+
*/
52+
export function initializeMetricsCsv(metricsCsv?: string, logger?: rpc.Logger): void {
53+
if (!metricsCsv) {
54+
return;
55+
}
56+
57+
try {
58+
// Check if file exists and has content
59+
if (fs.existsSync(metricsCsv)) {
60+
const content = fs.readFileSync(metricsCsv, 'utf8');
61+
const firstLine = content.split('\n')[0];
62+
63+
// If file already has the correct header, skip initialization
64+
if (firstLine.trim() === CSV_HEADER) {
65+
return;
66+
}
67+
68+
// File exists but has incorrect header - warn and reset
69+
if (firstLine.trim()) {
70+
logger?.warn(`Metrics CSV file ${metricsCsv} has incorrect header. Expected '${CSV_HEADER}' but found '${firstLine.trim()}'. Resetting file.`);
71+
}
72+
}
73+
74+
// Write header to new or empty file (overwrites existing file with incorrect header)
75+
fs.writeFileSync(metricsCsv, CSV_HEADER + '\n');
76+
} catch (err) {
77+
console.error('Failed to initialize metrics CSV:', err);
78+
}
79+
}
80+
81+
/**
82+
* Internal function to wrap a handler with metrics recording.
83+
*/
84+
async function wrapWithMetrics<R>(
85+
handler: () => Promise<R>,
86+
target: { target: string },
87+
request: string,
88+
metricsCsv?: string
89+
): Promise<R> {
90+
if (!metricsCsv) {
91+
// No metrics recording requested, just execute the handler
92+
return handler();
93+
}
94+
95+
const startTime = Date.now();
96+
97+
try {
98+
const result = await handler();
99+
recordMetrics(metricsCsv, target.target, request, startTime);
100+
return result;
101+
} catch (error) {
102+
recordMetrics(metricsCsv, target.target, request, startTime);
103+
throw error;
104+
}
105+
}
106+
107+
/**
108+
* Wraps an RPC request handler to record performance metrics to a CSV file.
109+
*
110+
* @param request The request type name (e.g., "Visit", "GetObject")
111+
* @param target A mutable object containing the target identifier for metrics
112+
* @param metricsCsv Optional path to the CSV file for recording metrics
113+
* @returns A function that wraps a request handler with metrics recording
114+
*/
115+
export function withMetrics<P, R>(
116+
request: string,
117+
target: { target: string },
118+
metricsCsv?: string
119+
): (handler: (request: P) => Promise<R>) => (request: P) => Promise<R> {
120+
return (handler: (requestParam: P) => Promise<R>) => {
121+
return async (requestParam: P): Promise<R> => {
122+
return wrapWithMetrics(() => handler(requestParam), target, request, metricsCsv);
123+
};
124+
};
125+
}
126+
127+
/**
128+
* Wraps an RPC request handler without parameters (RequestType0) to record performance metrics.
129+
*
130+
* @param request The request type name (e.g., "GetLanguages", "GetRecipes")
131+
* @param target A mutable object containing the target identifier for metrics
132+
* @param metricsCsv Optional path to the CSV file for recording metrics
133+
* @returns A function that wraps a request handler with metrics recording
134+
*/
135+
export function withMetrics0<R>(
136+
request: string,
137+
target: { target: string },
138+
metricsCsv?: string
139+
): (handler: () => Promise<R>) => (token: any) => Promise<R> {
140+
return (handler: () => Promise<R>) => {
141+
return async (_: any): Promise<R> => {
142+
return wrapWithMetrics(handler, target, request, metricsCsv);
143+
};
144+
};
145+
}
146+
147+
function recordMetrics(
148+
metricsCsv: string,
149+
target: string,
150+
request: string,
151+
startTime: number
152+
): void {
153+
const endTime = Date.now();
154+
const memEnd = process.memoryUsage();
155+
const durationMs = endTime - startTime;
156+
157+
const memoryUsedBytes = memEnd.heapUsed;
158+
const memoryMaxBytes = memEnd.heapTotal;
159+
160+
const csvRow = `${request},${target},${durationMs},${memoryUsedBytes},${memoryMaxBytes}\n`;
161+
162+
try {
163+
fs.appendFileSync(metricsCsv, csvRow);
164+
} catch (err) {
165+
console.error('Failed to write metrics to CSV:', err);
166+
}
167+
}

rewrite-javascript/rewrite/src/rpc/request/parse.ts

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,26 @@
1515
*/
1616
import * as rpc from "vscode-jsonrpc/node";
1717
import {ExecutionContext} from "../../execution";
18-
import {UUID} from "node:crypto";
1918
import {ParserInput, Parsers} from "../../parser";
20-
import {randomId} from "../../uuid";
19+
import {randomId, UUID} from "../../uuid";
2120
import {produce} from "immer";
2221
import {SourceFile} from "../../tree";
22+
import {withMetrics} from "./metrics";
2323

2424
export class Parse {
2525
constructor(private readonly inputs: ParserInput[],
2626
private readonly relativeTo?: string) {
2727
}
2828

2929
static handle(connection: rpc.MessageConnection,
30-
localObjects: Map<string, ((input: string) => any) | any>): void {
31-
connection.onRequest(new rpc.RequestType<Parse, UUID[], Error>("Parse"), async (request) => {
30+
localObjects: Map<string, ((input: string) => any) | any>,
31+
metricsCsv?: string): void {
32+
const target = { target: '' };
33+
connection.onRequest(new rpc.RequestType<Parse, UUID[], Error>("Parse"), withMetrics<Parse, UUID[]>("Parse", target, metricsCsv)(async (request) => {
34+
// Set target to comma-separated list of file paths
35+
target.target = request.inputs.map(input =>
36+
typeof input === 'string' ? input : input.sourcePath
37+
).join(',');
3238
let parser = Parsers.createParser("javascript", {
3339
ctx: new ExecutionContext(),
3440
relativeTo: request.relativeTo
@@ -38,7 +44,7 @@ export class Parse {
3844
return [];
3945
}
4046
const generator = parser.parse(...request.inputs);
41-
const result: string[] = [];
47+
const result: UUID[] = [];
4248

4349
for (let i = 0; i < request.inputs.length; i++) {
4450
const id = randomId();
@@ -50,6 +56,6 @@ export class Parse {
5056
}
5157

5258
return result;
53-
});
59+
}));
5460
}
5561
}

rewrite-javascript/rewrite/src/rpc/request/prepare-recipe.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,20 @@ import {Check} from "../../preconditions";
2121
import {RpcRecipe} from "../recipe";
2222
import {TreeVisitor} from "../../visitor";
2323
import {ExecutionContext} from "../../execution";
24+
import {withMetrics} from "./metrics";
2425

2526
export class PrepareRecipe {
2627
constructor(private readonly id: string, private readonly options?: any) {
2728
}
2829

2930
static handle(connection: MessageConnection,
3031
registry: RecipeRegistry,
31-
preparedRecipes: Map<String, Recipe>) {
32+
preparedRecipes: Map<String, Recipe>,
33+
metricsCsv?: string) {
3234
const snowflake = SnowflakeId();
33-
connection.onRequest(new rpc.RequestType<PrepareRecipe, PrepareRecipeResponse, Error>("PrepareRecipe"), async (request) => {
35+
const target = { target: '' };
36+
connection.onRequest(new rpc.RequestType<PrepareRecipe, PrepareRecipeResponse, Error>("PrepareRecipe"), withMetrics<PrepareRecipe, PrepareRecipeResponse>("PrepareRecipe", target, metricsCsv)(async (request) => {
37+
target.target = request.id;
3438
const id = snowflake.generate();
3539
const recipeCtor = registry.all.get(request.id);
3640
if (!recipeCtor) {
@@ -54,7 +58,7 @@ export class PrepareRecipe {
5458
scanVisitor: recipe instanceof ScanningRecipe ? `scan:${id}` : undefined,
5559
scanPreconditions: scanPreconditions
5660
}
57-
});
61+
}));
5862
}
5963

6064
/**

0 commit comments

Comments
 (0)