[api] Add timing collection and dump via environment variable (#64846)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 9920adc3-3fc7-4d27-8ec1-c9c4490b9cf9

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

@@ -838,8 +838,8 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
      *
      * Fetching the snapshot issues a lightweight request to the server to
      * retrieve its timing collection. Collection must be enabled via the
-     * `collectTiming` option; when it is not, the returned snapshot has
-     * `enabled: false` and zeroed totals.
+     * `collectTiming` option or `TS_API_TIMING_DIR` environment variable; when
+     * it is not, the returned snapshot has `enabled: false` and zeroed totals.
      */
     getTimingInfo(): Promise<TimingInfo> {
         return this.client.getTimingInfo();

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

@@ -43,6 +43,7 @@ import {
     TimingCollector,
     type TimingInfo,
 } from "../timing.ts";
+import { TimingDump } from "../timingDump.ts";
 
 export type { ClientOptions, ClientSocketOptions, ClientSpawnOptions };
 
@@ -59,16 +60,22 @@ export class Client {
     private closed = false;
     private connecting: Promise<void> | undefined;
     private timing: TimingCollector | undefined;
+    private timingDump: TimingDump | undefined;
+    private closing: Promise<void> | undefined;
     private batchedRequests: { method: APIRequest["method"]; params: APIRequest["params"]; resolve: (value: unknown) => void; reject: (reason?: any) => void; }[] = [];
     private nextBatch: NodeJS.Immediate | "manual" | undefined;
 
     constructor(options: ClientOptions) {
+        this.timingDump = TimingDump.fromEnvironment("async", isSpawnOptions(options) ? "spawn" : "socket");
+        if (this.timingDump || isSpawnOptions(options) && options.collectTiming) {
+            this.timing = new TimingCollector();
+        }
+        if (this.timingDump && isSpawnOptions(options)) {
+            options = { ...options, collectTiming: true };
+        }
         this.options = options;
         if (isSpawnOptions(options)) {
             configureFileSystemCallbacks(options.fs);
-            if (options.collectTiming) {
-                this.timing = new TimingCollector();
-            }
         }
     }
 
@@ -401,26 +408,52 @@ export class Client {
         return this.connection.sendRequest(requestType, undefined);
     }
 
-    async close(): Promise<void> {
+    close(): Promise<void> {
+        return this.closing ??= this.closeWorker();
+    }
+
+    private async closeWorker(): Promise<void> {
         await this.connecting?.catch(() => {}); // if connection is still in-progress, wait for it to finish before closing the connection
         this.closed = true;
-        if (this.connection) {
-            this.connection.dispose();
-            this.connection = undefined;
+        try {
+            if (this.timingDump && this.timing) {
+                const local = this.timing.getInfo();
+                let server: ServerTimingInfo | undefined;
+                let serverTimingError: string | undefined;
+                const connection = this.connection;
+                // Closing alone does not reject pending JSON-RPC requests.
+                const onClose = connection?.onClose(() => connection.dispose());
+                try {
+                    if (this.connected) server = await this.fetchServerTiming();
+                }
+                catch (error) {
+                    serverTimingError = String(error);
+                }
+                finally {
+                    onClose?.dispose();
+                }
+                this.timingDump.write(local, server, serverTimingError);
+            }
         }
-        if (this.socket) {
-            this.socket.destroy();
-            this.socket = undefined;
+        finally {
+            if (this.connection) {
+                this.connection.dispose();
+                this.connection = undefined;
+            }
+            if (this.socket) {
+                this.socket.destroy();
+                this.socket = undefined;
+            }
+            if (this.process) {
+                // Close stdin to unblock the server's read loop, allowing it to exit cleanly.
+                // The server is blocked on stdin.Read(), so just sending SIGTERM would deadlock:
+                // - Node won't exit while child is alive
+                // - Child can't process SIGTERM while blocked on read
+                // - Read won't error until stdin is closed
+                this.process.stdin?.end();
+                this.process = undefined;
+            }
+            this.connected = false;
         }
-        if (this.process) {
-            // Close stdin to unblock the server's read loop, allowing it to exit cleanly.
-            // The server is blocked on stdin.Read(), so just sending SIGTERM would deadlock:
-            // - Node won't exit while child is alive
-            // - Child can't process SIGTERM while blocked on read
-            // - Read won't error until stdin is closed
-            this.process.stdin?.end();
-            this.process = undefined;
-        }
-        this.connected = false;
     }
 }

packages/typescript/src/api/options.ts

@@ -33,6 +33,9 @@ export interface ClientSpawnOptions {
      * measures its own per-request processing time; both are combined (along
      * with an estimated transport overhead) in the snapshot returned by
      * {@link API.getTimingInfo}.
+     * Setting `TS_API_TIMING_DIR` to a nonempty output directory also enables
+     * collection, even when this option is false, and writes a JSON snapshot
+     * when the API is closed.
      */
     collectTiming?: boolean | undefined;
 }

packages/typescript/src/api/sync/api.ts

@@ -1343,8 +1343,8 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
      *
      * Fetching the snapshot issues a lightweight request to the server to
      * retrieve its timing collection. Collection must be enabled via the
-     * `collectTiming` option; when it is not, the returned snapshot has
-     * `enabled: false` and zeroed totals.
+     * `collectTiming` option or `TS_API_TIMING_DIR` environment variable; when
+     * it is not, the returned snapshot has `enabled: false` and zeroed totals.
      */
     get getTimingInfo(): {
         (): TimingInfo;

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

@@ -25,6 +25,7 @@ import {
     TimingCollector,
     type TimingInfo,
 } from "../timing.ts";
+import { TimingDump } from "../timingDump.ts";
 
 export type { ClientOptions, ClientSocketOptions, ClientSpawnOptions };
 
@@ -32,6 +33,8 @@ export class Client {
     private channel: SyncRpcChannel;
     private encoder = new TextEncoder();
     private timing: TimingCollector | undefined;
+    private timingDump: TimingDump | undefined;
+    private closed = false;
     private maxResponseBytesPerPage: number | undefined;
 
     constructor(options: ClientOptions) {
@@ -39,6 +42,10 @@ export class Client {
             throw new Error("Socket connections are not yet supported in the sync client");
         }
 
+        this.timingDump = TimingDump.fromEnvironment("sync", "spawn");
+        if (this.timingDump) {
+            options = { ...options, collectTiming: true };
+        }
         const args = getAPIProcessArgs(options, false);
         this.maxResponseBytesPerPage = options.maxResponseBytesPerPage;
 
@@ -184,6 +191,24 @@ export class Client {
     }
 
     close(): void {
-        this.channel.close();
+        if (this.closed) return;
+        this.closed = true;
+        try {
+            if (this.timingDump && this.timing) {
+                const local = this.timing.getInfo();
+                let server: ServerTimingInfo | undefined;
+                let serverTimingError: string | undefined;
+                try {
+                    server = JSON.parse(this.channel.requestSync("getServerTiming", "")) as ServerTimingInfo;
+                }
+                catch (error) {
+                    serverTimingError = String(error);
+                }
+                this.timingDump.write(local, server, serverTimingError);
+            }
+        }
+        finally {
+            this.channel.close();
+        }
     }
 }

packages/typescript/src/api/syncChannel.ts

@@ -113,6 +113,7 @@ export class SyncRpcChannel {
     private writeFd: number;
     private pipeFd: number | undefined;
     private callbacks = new Map<string, (name: string, payload: string) => string>();
+    private callbackFailure: { error: unknown; } | undefined;
 
     private methodBufCache = new Map<string, Buffer>();
 
@@ -280,6 +281,9 @@ export class SyncRpcChannel {
         if (this.readFd < 0) {
             throw new Error("SyncRpcChannel is closed");
         }
+        if (this.callbackFailure) {
+            throw this.callbackFailure.error;
+        }
     }
 
     private getMethodBuf(method: string): Buffer {
@@ -326,7 +330,14 @@ export class SyncRpcChannel {
                     );
                 }
                 case MSG_CALL: {
-                    this.handleCall(this._msgName.toString("utf-8"), this._msgPayload);
+                    try {
+                        this.handleCall(this._msgName.toString("utf-8"), this._msgPayload);
+                    }
+                    catch (error) {
+                        // The pending response is unread, so further I/O cannot be synchronized.
+                        this.callbackFailure = { error };
+                        throw error;
+                    }
                     break;
                 }
                 default:

packages/typescript/src/api/timing.ts

@@ -182,7 +182,7 @@ export function disabledServerTimingInfo(): ServerTimingInfo {
  * requests used to fetch timing) do not misalign the two ring buffers.
  */
 export function combineTimingInfo(client: TimingInfo, server: ServerTimingInfo): TimingInfo {
-    if (!client.enabled) {
+    if (!client.enabled || !server.enabled) {
         return client;
     }
 

packages/typescript/src/api/timingDump.ts

@@ -0,0 +1,55 @@
+import { randomUUID } from "node:crypto";
+import {
+    mkdirSync,
+    writeFileSync,
+} from "node:fs";
+import {
+    join,
+    resolve,
+} from "node:path";
+import {
+    combineTimingInfo,
+    type ServerTimingInfo,
+    type TimingInfo,
+} from "./timing.ts";
+
+/** One final snapshot per client, written before its connection is closed. */
+export class TimingDump {
+    private readonly instanceId = randomUUID();
+    private readonly startedAt = new Date().toISOString();
+    private readonly directory: string;
+    private readonly mode: "sync" | "async";
+    private readonly transport: "spawn" | "socket";
+
+    static fromEnvironment(mode: "sync" | "async", transport: "spawn" | "socket"): TimingDump | undefined {
+        const directory = process.env.TS_API_TIMING_DIR;
+        return directory ? new TimingDump(resolve(directory), mode, transport) : undefined;
+    }
+
+    private constructor(directory: string, mode: "sync" | "async", transport: "spawn" | "socket") {
+        this.directory = directory;
+        this.mode = mode;
+        this.transport = transport;
+    }
+
+    write(client: TimingInfo, server: ServerTimingInfo | undefined, serverTimingError?: string): void {
+        const snapshot = {
+            schemaVersion: 1,
+            instanceId: this.instanceId,
+            pid: process.pid,
+            mode: this.mode,
+            transport: this.transport,
+            startedAt: this.startedAt,
+            closedAt: new Date().toISOString(),
+            timing: server ? combineTimingInfo(client, server) : client,
+            serverTiming: server ?? null,
+            serverTimingError,
+        };
+        mkdirSync(this.directory, { recursive: true });
+        writeFileSync(
+            join(this.directory, `timing-${process.pid}-${this.instanceId}.json`),
+            JSON.stringify(snapshot) + "\n",
+            { flag: "wx" },
+        );
+    }
+}

packages/typescript/test/timingDump.test.ts

@@ -0,0 +1,132 @@
+import {
+    createMessageConnection,
+    SocketMessageReader,
+    SocketMessageWriter,
+} from "#vscode-jsonrpc/node";
+import { API as AsyncAPI } from "@typescript/typescript/async";
+import { API as SyncAPI } from "@typescript/typescript/sync";
+import assert from "node:assert/strict";
+import { once } from "node:events";
+import {
+    mkdtempSync,
+    readdirSync,
+    readFileSync,
+    rmSync,
+} from "node:fs";
+import { createServer } from "node:net";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { test } from "node:test";
+import {
+    combineTimingInfo,
+    disabledServerTimingInfo,
+    TimingCollector,
+} from "../src/api/timing.ts";
+import { createVirtualFileSystem } from "./testUtils.ts";
+
+test("disabled server timing preserves the client-only snapshot", () => {
+    const collector = new TimingCollector();
+    collector.record({ method: "initialize", roundTripMs: 10, bytesSent: 1, bytesReceived: 1 });
+    const local = collector.getInfo();
+    assert.strictEqual(combineTimingInfo(local, disabledServerTimingInfo()), local);
+});
+
+test("timing dumps", async t => {
+    const directory = mkdtempSync(join(tmpdir(), "ts-timing-"));
+    const original = process.env.TS_API_TIMING_DIR;
+    t.after(() => {
+        if (original === undefined) delete process.env.TS_API_TIMING_DIR;
+        else process.env.TS_API_TIMING_DIR = original;
+        rmSync(directory, { recursive: true, force: true });
+    });
+
+    function assertDumpCount(count: number): void {
+        const files = readdirSync(directory).filter(file => file.endsWith(".json"));
+        assert.equal(files.length, count);
+        for (const file of files) {
+            assert.ok(Object.keys(JSON.parse(readFileSync(join(directory, file), "utf8"))).length > 0);
+        }
+    }
+
+    let count = 0;
+    for (const [mode, API] of [["async", AsyncAPI], ["sync", SyncAPI]] as const) {
+        for (const value of [undefined, "", directory]) {
+            await t.test(`${mode}: dumping ${value ? "enabled" : value === "" ? "empty" : "unset"}`, async () => {
+                if (value === undefined) delete process.env.TS_API_TIMING_DIR;
+                else process.env.TS_API_TIMING_DIR = value;
+                const api = new API({ collectTiming: false });
+                try {
+                    await api.parseCommandLine([]);
+                    assert.equal((await api.getTimingInfo()).enabled, !!value);
+                    assertDumpCount(count);
+                }
+                finally {
+                    await api.close();
+                }
+                if (value) count++;
+                assertDumpCount(count);
+                await api.close();
+                assertDumpCount(count);
+            });
+        }
+    }
+
+    process.env.TS_API_TIMING_DIR = directory;
+    await t.test("sync: preserves a dump after a filesystem callback throws", async t => {
+        const warning = t.mock.method(process, "emitWarning", () => {});
+        const api = new SyncAPI({
+            fs: {
+                ...createVirtualFileSystem({}),
+                readFile() {
+                    throw new Error("read failed");
+                },
+            },
+        });
+        try {
+            api.parseCommandLine([]);
+            assert.throws(() => api.readConfigFile("/tsconfig.json"), /read failed/);
+            assert.throws(() => api.parseCommandLine([]), /read failed/);
+            assert.throws(() => api.getTimingInfo(), /read failed/);
+        }
+        finally {
+            api.close();
+        }
+        assertDumpCount(++count);
+        assert.equal(warning.mock.callCount(), 0);
+    });
+
+    for (const scenario of ["error", "disconnect", "disabled"] as const) {
+        await t.test(`async: preserves a dump when server timing is ${scenario}`, { timeout: 5000 }, async t => {
+            const warning = t.mock.method(process, "emitWarning", () => {});
+            const pipe = process.platform === "win32" ? `\\\\.\\pipe\\ts-timing-${process.pid}` : join(directory, "api.sock");
+            let disconnect = () => {};
+            const server = createServer(socket => {
+                const connection = createMessageConnection(new SocketMessageReader(socket), new SocketMessageWriter(socket));
+                disconnect = () => {
+                    connection.dispose();
+                    socket.destroy();
+                };
+                connection.onRequest("initialize", () => ({ currentDirectory: process.cwd(), caseSensitivity: 1 }));
+                connection.onRequest("getServerTiming", () => {
+                    if (scenario === "disconnect") {
+                        disconnect();
+                        return new Promise(() => {});
+                    }
+                    if (scenario === "disabled") return disabledServerTimingInfo();
+                    throw new Error("server timing unavailable");
+                });
+                connection.listen();
+            });
+            t.after(() => {
+                disconnect();
+                return new Promise<void>(resolve => server.close(() => resolve()));
+            });
+            server.listen(pipe);
+            await once(server, "listening");
+            const api = await AsyncAPI.fromLSPConnection({ pipe });
+            await api.close();
+            assertDumpCount(++count);
+            assert.equal(warning.mock.callCount(), 0);
+        });
+    }
+});