Improve performance of batched requests

Herebyfile.mjs

@@ -656,9 +656,9 @@ async function runGenerateAPI() {
         ],
         exclude: ["**/*_test.go", "**/*_generated.go"],
         envInputs: [],
-        outputs: ["packages/typescript/src/api/proto.generated.ts"],
+        outputs: ["packages/typescript/src/api/proto.generated.ts", "tsc/internal/api/batch_decoder_generated.go"],
         commands: [
-            ["go", "-C", "./tools", "run", "./gen-proto", "../tsc/internal/api/proto.go", "../packages/typescript/src/api/proto.generated.ts"],
+            ["go", "-C", "./tools", "run", "./gen-proto", "../tsc/internal/api/proto.go", "../packages/typescript/src/api/proto.generated.ts", "../tsc/internal/api/batch_decoder_generated.go"],
             ["dprint", "fmt", "packages/typescript/src/api/proto.generated.ts"],
         ],
     });

packages/typescript/src/api/async/api.ts

@@ -37,6 +37,10 @@ import {
     unescapeLeadingUnderscores,
 } from "../../ast/index.ts";
 import { assertNever } from "../../internal/utils.ts";
+import {
+    createGroupedBatchRequest,
+    getBatchResults,
+} from "../batch.ts";
 import {
     encodeNode,
     uint8ArrayToBase64,
@@ -354,7 +358,7 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
     // @sync-skip-block-end
     // @sync-only-start
     // batch<T extends readonly AnyAPIRequestGenerator[]>(...requestGenerators: T): ExecutedGeneratorsResults<T> {
-    //     return executeRequestGenerators(requestGenerators, requests => this.client.batchRequests(requests).responses);
+    //     return executeRequestGenerators(requestGenerators, requests => this.client.batchRequests(requests));
     // }
     // @sync-only-end
 
@@ -1922,12 +1926,12 @@ export class Checker {
     getSymbolAtLocation(nodes: readonly Node[]): Promise<(Symbol | undefined)[]>;
     async getSymbolAtLocation(nodeOrNodes: Node | readonly Node[]): Promise<Symbol | (Symbol | undefined)[] | undefined> {
         if (Array.isArray(nodeOrNodes)) {
-            const data = await this.client.apiRequest("getSymbolsAtLocations", {
-                snapshot: this.snapshotId,
-                project: this.project.id,
-                locations: nodeOrNodes.map(node => getNodeId(node)),
-            });
-            return data.map(d => d ? this.objectRegistry.getOrCreateSymbol(d) : undefined);
+            const response = await this.client.batchRequest(createGroupedBatchRequest(
+                "getSymbolAtLocation",
+                { snapshot: this.snapshotId, project: this.project.id },
+                { location: nodeOrNodes.map(node => getNodeId(node)) },
+            ));
+            return getBatchResults<SymbolResponse | null>(response).map(data => data ? this.objectRegistry.getOrCreateSymbol(data) : undefined);
         }
         const data = await this.client.apiRequest("getSymbolAtLocation", {
             snapshot: this.snapshotId,
@@ -1949,25 +1953,24 @@ export class Checker {
             });
             return data ? this.objectRegistry.getOrCreateSymbol(data) : undefined;
         }
-        const data = await this.client.apiRequest("getSymbolsAtPositions", {
-            snapshot: this.snapshotId,
-            project: this.project.id,
-            file,
-            positions: positionOrPositions,
-        });
-        return data.map(d => d ? this.objectRegistry.getOrCreateSymbol(d) : undefined);
+        const response = await this.client.batchRequest(createGroupedBatchRequest(
+            "getSymbolAtPosition",
+            { snapshot: this.snapshotId, project: this.project.id, file },
+            { position: positionOrPositions },
+        ));
+        return getBatchResults<SymbolResponse | null>(response).map(data => data ? this.objectRegistry.getOrCreateSymbol(data) : undefined);
     }
 
     getSymbolOfSourceFile(file: DocumentIdentifier): Promise<Symbol | undefined>;
     getSymbolOfSourceFile(files: readonly DocumentIdentifier[]): Promise<(Symbol | undefined)[]>;
     async getSymbolOfSourceFile(fileOrFiles: DocumentIdentifier | readonly DocumentIdentifier[]): Promise<Symbol | (Symbol | undefined)[] | undefined> {
         if (Array.isArray(fileOrFiles)) {
-            const data = await this.client.apiRequest("getSymbolsOfSourceFiles", {
-                snapshot: this.snapshotId,
-                project: this.project.id,
-                files: fileOrFiles,
-            });
-            return data.map(d => d ? this.objectRegistry.getOrCreateSymbol(d) : undefined);
+            const response = await this.client.batchRequest(createGroupedBatchRequest(
+                "getSymbolOfSourceFile",
+                { snapshot: this.snapshotId, project: this.project.id },
+                { file: fileOrFiles },
+            ));
+            return getBatchResults<SymbolResponse | null>(response).map(data => data ? this.objectRegistry.getOrCreateSymbol(data) : undefined);
         }
         const data = await this.client.apiRequest("getSymbolOfSourceFile", {
             snapshot: this.snapshotId,
@@ -1986,12 +1989,12 @@ export class Checker {
     getTypeOfSymbol(symbols: readonly Symbol[]): Promise<Type[]>;
     async getTypeOfSymbol(symbolOrSymbols: Symbol | readonly Symbol[]): Promise<Type | Type[]> {
         if (Array.isArray(symbolOrSymbols)) {
-            const data = await this.client.apiRequest("getTypesOfSymbols", {
-                snapshot: this.snapshotId,
-                project: this.project.id,
-                symbols: symbolOrSymbols.map(s => s.id),
-            });
-            return data.map(d => this.objectRegistry.getOrCreateType(d));
+            const response = await this.client.batchRequest(createGroupedBatchRequest(
+                "getTypeOfSymbol",
+                { snapshot: this.snapshotId, project: this.project.id },
+                { symbol: symbolOrSymbols.map(symbol => symbol.id) },
+            ));
+            return getBatchResults<TypeResponse>(response).map(data => this.objectRegistry.getOrCreateType(data));
         }
         const data = await this.client.apiRequest("getTypeOfSymbol", {
             snapshot: this.snapshotId,
@@ -2064,12 +2067,12 @@ export class Checker {
     getTypeAtLocation(nodes: readonly Node[]): Promise<Type[]>;
     async getTypeAtLocation(nodeOrNodes: Node | readonly Node[]): Promise<Type | Type[]> {
         if (Array.isArray(nodeOrNodes)) {
-            const data = await this.client.apiRequest("getTypeAtLocations", {
-                snapshot: this.snapshotId,
-                project: this.project.id,
-                locations: nodeOrNodes.map(node => getNodeId(node)),
-            });
-            return data.map(d => this.objectRegistry.getOrCreateType(d));
+            const response = await this.client.batchRequest(createGroupedBatchRequest(
+                "getTypeAtLocation",
+                { snapshot: this.snapshotId, project: this.project.id },
+                { location: nodeOrNodes.map(node => getNodeId(node)) },
+            ));
+            return getBatchResults<TypeResponse>(response).map(data => this.objectRegistry.getOrCreateType(data));
         }
         const data = await this.client.apiRequest("getTypeAtLocation", {
             snapshot: this.snapshotId,
@@ -2109,13 +2112,12 @@ export class Checker {
             });
             return data ? this.objectRegistry.getOrCreateType(data) : undefined;
         }
-        const data = await this.client.apiRequest("getTypesAtPositions", {
-            snapshot: this.snapshotId,
-            project: this.project.id,
-            file,
-            positions: positionOrPositions,
-        });
-        return data.map(d => d ? this.objectRegistry.getOrCreateType(d) : undefined);
+        const response = await this.client.batchRequest(createGroupedBatchRequest(
+            "getTypeAtPosition",
+            { snapshot: this.snapshotId, project: this.project.id, file },
+            { position: positionOrPositions },
+        ));
+        return getBatchResults<TypeResponse | null>(response).map(data => data ? this.objectRegistry.getOrCreateType(data) : undefined);
     }
 
     async resolveName(

packages/typescript/src/api/async/client.ts

@@ -9,6 +9,7 @@ import {
 } from "#vscode-jsonrpc/node";
 import type { ChildProcess } from "node:child_process";
 import type { Socket } from "node:net";
+import { groupBatchRequests } from "../batch.ts";
 import {
     type FileSystem,
     fsCallbackNames,
@@ -228,34 +229,20 @@ export class Client {
                 return;
             }
 
-            const requestType = new RequestType<unknown, BatchRequestsResponse, void>("batchRequests");
-            const params: BatchRequestsParams = { requests: requests.map(request => ({ method: request.method, params: request.params })) };
-            if (this.options.maxResponseBytesPerPage !== undefined) {
-                params.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage;
-            }
-            const response = await this.sendRequestWithTiming(requestType, params);
-            let responses = response.responses;
-            let continuationToken = response.continuationToken;
-            while (continuationToken) {
-                const pageParams: BatchRequestsParams = {
-                    requests: [],
-                    continuationToken,
-                };
-                if (this.options.maxResponseBytesPerPage !== undefined) {
-                    pageParams.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage;
-                }
-                const page = await this.sendRequestWithTiming(requestType, pageParams);
-                responses = responses.concat(page.responses);
-                continuationToken = page.continuationToken;
-            }
+            // Paired benchmarks show grouping pays for its construction cost at four requests.
+            const grouped = requests.length >= 4 ? groupBatchRequests(requests) : undefined;
+            const params: BatchRequestsParams = grouped
+                ? { groups: grouped.groups, groupOrder: grouped.groupOrder, maxResponseBytesPerPage: this.options.maxResponseBytesPerPage }
+                : { requests: requests.map(request => ({ method: request.method, params: request.params })), maxResponseBytesPerPage: this.options.maxResponseBytesPerPage };
+            const response = await this.batchRequest(params);
             for (let i = 0; i < requests.length; i++) {
                 const { resolve, reject } = requests[i];
-                const item = responses[i];
-                if (item.error !== undefined) {
-                    reject(new Error(item.error));
+                const error = response.errors?.[i];
+                if (error !== undefined) {
+                    reject(new Error(error));
                 }
                 else {
-                    resolve(item.result);
+                    resolve(response.results[i]);
                 }
             }
         }
@@ -264,6 +251,33 @@ export class Client {
         }
     }
 
+    async batchRequest(params: BatchRequestsParams): Promise<BatchRequestsResponse> {
+        if (this.closed) throw new Error("Client is closed");
+        if (!this.connected) {
+            await this.connect();
+        }
+        const requestType = new RequestType<unknown, BatchRequestsResponse, void>("batchRequests");
+        const response = await this.sendRequestWithTiming(requestType, params);
+        let results = response.results;
+        let errors = response.errors;
+        let continuationToken = response.continuationToken;
+        while (continuationToken) {
+            const page = await this.sendRequestWithTiming(requestType, {
+                requests: [],
+                continuationToken,
+                maxResponseBytesPerPage: this.options.maxResponseBytesPerPage,
+            });
+            const offset = results.length;
+            results = results.concat(page.results);
+            if (page.errors) {
+                errors ??= {};
+                for (const [index, message] of Object.entries(page.errors)) errors[Number(index) + offset] = message;
+            }
+            continuationToken = page.continuationToken;
+        }
+        return errors ? { results, errors } : { results };
+    }
+
     private scheduleImmediateBatch(): void {
         if (this.nextBatch) return;
         this.nextBatch = setImmediate(this.doBatch.bind(this));

packages/typescript/src/api/batch.ts

@@ -0,0 +1,115 @@
+import type {
+    APIRequest,
+    BatchRequestsParams,
+    BatchRequestsResponse,
+} from "./proto.ts";
+
+interface GroupedBatchRequest {
+    method: APIRequest["method"];
+    base?: Record<string, unknown>;
+    count: number;
+    fields?: Record<string, unknown[]>;
+    requests?: unknown[];
+}
+
+export function groupBatchRequests<T extends { method: APIRequest["method"]; params: APIRequest["params"]; }>(requests: readonly T[]): {
+    groups: GroupedBatchRequest[];
+    groupOrder?: number[];
+} | undefined {
+    const groupIndexes = new Map<APIRequest["method"], number>();
+    const requestsByMethod: T[][] = [];
+    const groupOrder: number[] = [];
+    let hasBatchableGroup = false;
+    for (const request of requests) {
+        let groupIndex = groupIndexes.get(request.method);
+        if (groupIndex === undefined) {
+            groupIndex = requestsByMethod.length;
+            groupIndexes.set(request.method, groupIndex);
+            requestsByMethod.push([]);
+        }
+        requestsByMethod[groupIndex].push(request);
+        if (requestsByMethod[groupIndex].length === 4) hasBatchableGroup = true;
+        groupOrder.push(groupIndex);
+    }
+    if (!hasBatchableGroup) return undefined;
+
+    const groups: GroupedBatchRequest[] = [];
+    for (const [method, groupIndex] of groupIndexes) {
+        const groupedRequests = requestsByMethod[groupIndex];
+        const params = groupedRequests.map(request => request.params);
+        if (groupedRequests.length < 4) {
+            groups.push({ method, count: params.length, requests: params });
+            continue;
+        }
+        if (params.every(param => typeof param === "object" && param !== null && !Array.isArray(param))) {
+            const records = params as Record<string, unknown>[];
+            const base = commonParams(records);
+            const commonKeys = new Set(Object.keys(base));
+            const deltas = records.map(param => Object.fromEntries(Object.entries(param).filter(([key]) => !commonKeys.has(key))));
+            const fields = parameterColumns(deltas);
+            groups.push(
+                fields
+                    ? { method, base, count: records.length, fields }
+                    : { method, base, count: records.length, requests: deltas },
+            );
+        }
+        else {
+            groups.push({ method, count: params.length, requests: params });
+        }
+    }
+    return groups.length === 1 ? { groups } : { groups, groupOrder };
+}
+
+function parameterColumns(params: readonly Record<string, unknown>[]): Record<string, unknown[]> | undefined {
+    const keys = Object.keys(params[0]);
+    if (
+        !params.every(param => {
+            const paramKeys = Object.keys(param);
+            return paramKeys.length === keys.length
+                && keys.every(key => Object.hasOwn(param, key) && param[key] !== undefined);
+        })
+    ) return undefined;
+    return Object.fromEntries(keys.map(key => [key, params.map(param => param[key])])) as Record<string, unknown[]>;
+}
+
+function commonParams(params: readonly Record<string, unknown>[]): Record<string, unknown> {
+    const first = params[0];
+    return Object.fromEntries(
+        Object.entries(first).filter(([key, value]) => {
+            const type = typeof value;
+            return (value === null || type !== "object" && type !== "undefined")
+                && params.every(param => Object.hasOwn(param, key) && param[key] === value);
+        }),
+    );
+}
+
+export function createGroupedBatchRequest(
+    method: APIRequest["method"],
+    base: Record<string, unknown>,
+    fields: Record<string, readonly unknown[]>,
+): BatchRequestsParams {
+    const entries = Object.entries(fields);
+    const count = entries[0]?.[1].length ?? 0;
+    if (!entries.every(([, values]) => values.length === count)) {
+        throw new Error("Grouped batch parameter columns must have equal lengths");
+    }
+    if (count < 4) {
+        return {
+            requests: Array.from({ length: count }, (_, index) => ({
+                method,
+                params: { ...base, ...Object.fromEntries(entries.map(([key, values]) => [key, values[index]])) },
+            })),
+        };
+    }
+    return {
+        groups: [{ method, base, count, fields: Object.fromEntries(entries) }],
+    };
+}
+
+export function getBatchResults<T>(response: BatchRequestsResponse): T[] {
+    return response.results.map((result, index) => {
+        const error = response.errors?.[index];
+        if (error !== undefined) throw new Error(error);
+        return result as T;
+    });
+}