[api] Give standalone created files a tracked identity and incorporate with the parse cache

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

@@ -47,6 +47,7 @@ import {
     parseNodeHandle,
     readParseOptionsKey,
     readSourceFileHash,
+    readSourceFileLease,
     RemoteSourceFile,
 } from "../node/node.ts";
 import { Wtf8Decoder } from "../node/wtf8.ts";
@@ -327,6 +328,7 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
     private initialized: boolean = false;
     private initializing: Promise<void> | undefined;
     private activeSnapshots: Map<number, Snapshot> = new Map();
+    private activeSourceFileLeases: Map<number, RetainedSourceFile> = new Map();
     readonly printer: Printer;
     readonly internal: InternalAPI;
 
@@ -422,22 +424,57 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
         return this.client.apiRequest("parseJsonConfigFileContent", { json, ...options });
     }
 
-    async createSourceFile(fileName: string, sourceText: string, options: CreateSourceFileOptions = {}): Promise<SourceFile> {
+    /**
+     * Create and retain a source file independently of a program.
+     * Dispose the returned lease when the source file no longer needs to remain available remotely.
+     */
+    async createSourceFile(fileName: string, sourceText: string, options: CreateSourceFileOptions = {}): Promise<RetainedSourceFile> {
         await this.ensureInitialized();
         const data = await this.client.apiRequestBinary("createSourceFile", { fileName, sourceText, options });
         if (!data) {
             throw new Error("createSourceFile returned no source file");
         }
-        return new RemoteSourceFile(data, this.decoder, this.client.getTimingCollector()) as unknown as SourceFile;
+        return this.retainSourceFileResponse(data);
     }
 
-    async createSourceFileFromFile(file: DocumentIdentifier, options: CreateSourceFileOptions = {}): Promise<SourceFile> {
+    /**
+     * Read, create, and retain a source file independently of a program.
+     * Dispose the returned lease when the source file no longer needs to remain available remotely.
+     */
+    async createSourceFileFromFile(file: DocumentIdentifier, options: CreateSourceFileOptions = {}): Promise<RetainedSourceFile> {
         await this.ensureInitialized();
         const data = await this.client.apiRequestBinary("createSourceFileFromFile", { fileName: resolveFileName(file), options });
         if (!data) {
             throw new Error("createSourceFileFromFile returned no source file");
         }
-        return new RemoteSourceFile(data, this.decoder, this.client.getTimingCollector()) as unknown as SourceFile;
+        return this.retainSourceFileResponse(data);
+    }
+
+    private retainSourceFileResponse(data: Uint8Array): RetainedSourceFile {
+        const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
+        const lease = readSourceFileLease(view);
+        try {
+            const decoded = new RemoteSourceFile(data, this.decoder, this.client.getTimingCollector()) as unknown as SourceFile;
+            const sourceFile = this.sourceFileCache.setForLease(decoded.path, decoded, readParseOptionsKey(view), readSourceFileHash(view), lease);
+            const retained = new RetainedSourceFile(sourceFile, lease, this.client, () => {
+                this.activeSourceFileLeases.delete(lease);
+                this.sourceFileCache.releaseLease(lease);
+            });
+            this.activeSourceFileLeases.set(lease, retained);
+            return retained;
+        }
+        catch (error) {
+            // @sync-skip-block-start
+            void this.client.apiRequest("releaseSourceFile", { lease }).catch(() => {});
+            // @sync-skip-block-end
+            // @sync-only-start
+            // try {
+            //     this.client.apiRequest("releaseSourceFile", { lease });
+            // }
+            // catch {}
+            // @sync-only-end
+            throw error;
+        }
     }
 
     async transpileModule(input: string, options: TranspileOptions = {}): Promise<TranspileOutput> {
@@ -608,15 +645,21 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
 
     async close(): Promise<void> {
         await this.initializing?.catch(() => {}); // @sync-skip
-        // Dispose all active snapshots
         try {
-            for (const snapshot of [...this.activeSnapshots.values()]) {
-                await snapshot.dispose();
+            for (const retained of [...this.activeSourceFileLeases.values()]) {
+                await retained.dispose();
             }
-            this.sourceFileCache.clear();
         }
         finally {
-            await this.client.close(); // always close the underlying connection
+            try {
+                for (const snapshot of [...this.activeSnapshots.values()]) {
+                    await snapshot.dispose();
+                }
+                this.sourceFileCache.clear();
+            }
+            finally {
+                await this.client.close(); // always close the underlying connection
+            }
         }
     }
 
@@ -708,6 +751,42 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
 
 type EnsureInitialized = () => Promise<void>; // @sync: type EnsureInitialized = (() => void) & { gen(): Generator<ProtocolRequest, void, ProtocolResponse["result"]>; };
 
+/** An independently retained source file and its disposable remote-lifetime lease. */
+export class RetainedSourceFile {
+    readonly sourceFile: SourceFile;
+    private readonly lease: number;
+    private readonly client: Client;
+    private readonly onDispose: () => void;
+    private disposed = false;
+    private disposePromise: Promise<void> | undefined;
+
+    constructor(sourceFile: SourceFile, lease: number, client: Client, onDispose: () => void) {
+        this.sourceFile = sourceFile;
+        this.lease = lease;
+        this.client = client;
+        this.onDispose = onDispose;
+    }
+
+    [globalThis.Symbol.asyncDispose](): Promise<void> { // @sync: [globalThis.Symbol.dispose](): void {
+        return this.dispose(); // @sync: this.dispose();
+    }
+
+    dispose(): Promise<void> {
+        return this.disposePromise ??= this.disposeWorker();
+    }
+
+    private async disposeWorker(): Promise<void> {
+        if (this.disposed) return;
+        this.disposed = true;
+        try {
+            await this.client.apiRequest("releaseSourceFile", { lease: this.lease });
+        }
+        finally {
+            this.onDispose();
+        }
+    }
+}
+
 export class InternalAPI {
     private client: Client;
     private ensureInitialized: EnsureInitialized;

packages/typescript/src/api/node/node.infrastructure.ts

@@ -11,6 +11,7 @@ import {
     HEADER_OFFSET_HASH_LO0,
     HEADER_OFFSET_HASH_LO1,
     HEADER_OFFSET_PARSE_OPTIONS,
+    HEADER_OFFSET_SOURCE_FILE_LEASE,
     NODE_DATA_TYPE_CHILDREN,
     NODE_DATA_TYPE_EXTENDED,
     NODE_DATA_TYPE_STRING,
@@ -90,6 +91,17 @@ export function readParseOptionsKey(data: DataView): string {
     return data.getUint32(HEADER_OFFSET_PARSE_OPTIONS, true).toString();
 }
 
+export function readSourceFileLease(data: DataView): number {
+    const lease = data.getBigUint64(HEADER_OFFSET_SOURCE_FILE_LEASE, true);
+    if (lease === 0n) {
+        throw new Error("Source file response has no lease");
+    }
+    if (lease > BigInt(Number.MAX_SAFE_INTEGER)) {
+        throw new Error(`Source file lease ${lease} exceeds the maximum safe integer`);
+    }
+    return Number(lease);
+}
+
 function hex8(n: number): string {
     return (n >>> 0).toString(16).padStart(8, "0");
 }

packages/typescript/src/api/node/node.ts

@@ -38,7 +38,7 @@ import { Wtf8Decoder } from "./wtf8.ts";
 
 // Re-export everything consumers need from the other two files.
 export { RemoteNode, RemoteNodeList } from "./node.generated.ts";
-export { readParseOptionsKey, readSourceFileHash, RemoteNodeBase } from "./node.infrastructure.ts";
+export { readParseOptionsKey, readSourceFileHash, readSourceFileLease, RemoteNodeBase } from "./node.infrastructure.ts";
 
 const sourceFileExtendedDataOffsets = {
     Text: 0,

packages/typescript/src/api/node/protocol.ts

@@ -1,4 +1,4 @@
-export const PROTOCOL_VERSION = 8;
+export const PROTOCOL_VERSION = 9;
 
 export const HEADER_OFFSET_METADATA = 0;
 export const HEADER_OFFSET_HASH_LO0 = 4;
@@ -11,7 +11,10 @@ export const HEADER_OFFSET_STRING_TABLE = 28;
 export const HEADER_OFFSET_EXTENDED_DATA = 32;
 export const HEADER_OFFSET_STRUCTURED_DATA = 36;
 export const HEADER_OFFSET_NODES = 40;
-export const HEADER_SIZE = 44;
+export const HEADER_OFFSET_SOURCE_FILE_ID = 44;
+export const HEADER_OFFSET_SOURCE_FILE_LEASE = 52;
+export const HEADER_OFFSET_BINDER_DATA = 60;
+export const HEADER_SIZE = 64;
 
 export const NODE_LEN = 28;
 

packages/typescript/src/api/proto.generated.ts

@@ -21,6 +21,7 @@ export type APIMethod<TParams, TResult> = { params: TParams; result: TResult; };
 
 export interface APIMethodInfo {
     release: APIMethod<ReleaseParams, void>;
+    releaseSourceFile: APIMethod<ReleaseSourceFileParams, unknown>;
     batchRequests: APIMethod<BatchRequestsParams, BatchRequestsResponse>;
     initialize: APIMethod<null, InitializeResponse>;
     createSnapshot: APIMethod<CreateSnapshotParams, CreateSnapshotResponse>;
@@ -201,6 +202,10 @@ export interface ReleaseParams {
     snapshot: number;
 }
 
+export interface ReleaseSourceFileParams {
+    lease: number;
+}
+
 export interface BatchRequestsParams {
     requests: readonly BatchRequest[] | null;
     continuationToken?: string | undefined;
@@ -1186,6 +1191,7 @@ export interface BatchRequest {
         | "readConfigFile"
         | "release"
         | "releaseModuleResolver"
+        | "releaseSourceFile"
         | "resolveModuleName"
         | "resolveName"
         | "saveHeapProfile"
@@ -1354,6 +1360,7 @@ export interface BatchResponse {
         | "readConfigFile"
         | "release"
         | "releaseModuleResolver"
+        | "releaseSourceFile"
         | "resolveModuleName"
         | "resolveName"
         | "saveHeapProfile"