Adds tests and fixes issues by verifying incremental project updates are handling language service ref counting correctly (#54504)

Author: sheetalkamat

src/compiler/debug.ts

@@ -72,6 +72,7 @@ import {
     ObjectFlags,
     ObjectType,
     RelationComparisonResult,
+    ScriptKind,
     Signature,
     SignatureCheckMode,
     SignatureFlags,
@@ -439,6 +440,10 @@ export namespace Debug {
         return formatEnum(kind, (ts as any).SnippetKind, /*isFlags*/ false);
     }
 
+    export function formatScriptKind(kind: ScriptKind | undefined): string {
+        return formatEnum(kind, (ts as any).ScriptKind, /*isFlags*/ false);
+    }
+
     export function formatNodeFlags(flags: NodeFlags | undefined): string {
         return formatEnum(flags, (ts as any).NodeFlags, /*isFlags*/ true);
     }

src/compiler/moduleNameResolver.ts

@@ -880,7 +880,7 @@ function compilerOptionValueToString(value: unknown): string {
 
 /** @internal */
 export function getKeyForCompilerOptions(options: CompilerOptions, affectingOptionDeclarations: readonly CommandLineOption[]) {
-    return affectingOptionDeclarations.map(option => compilerOptionValueToString(getCompilerOptionValue(options, option))).join("|") + (options.pathsBasePath ? `|${options.pathsBasePath}` : undefined);
+    return affectingOptionDeclarations.map(option => compilerOptionValueToString(getCompilerOptionValue(options, option))).join("|") + `|${options.pathsBasePath}`;
 }
 
 /** @internal */

src/compiler/program.ts

@@ -2369,8 +2369,8 @@ export function createProgram(rootNamesOrOptions: readonly string[] | CreateProg
         for (const oldSourceFile of oldSourceFiles) {
             const sourceFileOptions = getCreateSourceFileOptions(oldSourceFile.fileName, moduleResolutionCache, host, options);
             let newSourceFile = host.getSourceFileByPath
-                ? host.getSourceFileByPath(oldSourceFile.fileName, oldSourceFile.resolvedPath, sourceFileOptions, /*onError*/ undefined, shouldCreateNewSourceFile || sourceFileOptions.impliedNodeFormat !== oldSourceFile.impliedNodeFormat)
-                : host.getSourceFile(oldSourceFile.fileName, sourceFileOptions, /*onError*/ undefined, shouldCreateNewSourceFile || sourceFileOptions.impliedNodeFormat !== oldSourceFile.impliedNodeFormat); // TODO: GH#18217
+                ? host.getSourceFileByPath(oldSourceFile.fileName, oldSourceFile.resolvedPath, sourceFileOptions, /*onError*/ undefined, shouldCreateNewSourceFile)
+                : host.getSourceFile(oldSourceFile.fileName, sourceFileOptions, /*onError*/ undefined, shouldCreateNewSourceFile); // TODO: GH#18217
 
             if (!newSourceFile) {
                 return StructureIsReused.Not;
@@ -3615,7 +3615,7 @@ export function createProgram(rootNamesOrOptions: readonly string[] | CreateProg
             fileName,
             sourceFileOptions,
             hostErrorMessage => addFilePreprocessingFileExplainingDiagnostic(/*file*/ undefined, reason, Diagnostics.Cannot_read_file_0_Colon_1, [fileName, hostErrorMessage]),
-            shouldCreateNewSourceFile || (oldProgram?.getSourceFileByPath(toPath(fileName))?.impliedNodeFormat !== sourceFileOptions.impliedNodeFormat)
+            shouldCreateNewSourceFile,
         );
 
         if (packageId) {

src/compiler/watchPublic.ts

@@ -724,7 +724,8 @@ export function createWatchProgram<T extends BuilderProgram>(host: WatchCompiler
         }
 
         // Create new source file if requested or the versions dont match
-        if (hostSourceFile === undefined || shouldCreateNewSourceFile || isFilePresenceUnknownOnHost(hostSourceFile)) {
+        const impliedNodeFormat = typeof languageVersionOrOptions === "object" ? languageVersionOrOptions.impliedNodeFormat : undefined;
+        if (hostSourceFile === undefined || shouldCreateNewSourceFile || isFilePresenceUnknownOnHost(hostSourceFile) || hostSourceFile.sourceFile.impliedNodeFormat !== impliedNodeFormat) {
             const sourceFile = getNewSourceFile(fileName, languageVersionOrOptions, onError);
             if (hostSourceFile) {
                 if (sourceFile) {

src/harness/harnessLanguageService.ts

@@ -9,6 +9,7 @@ import * as ts from "./_namespaces/ts";
 import { getNewLineCharacter } from "./_namespaces/ts";
 import * as vfs from "./_namespaces/vfs";
 import * as vpath from "./_namespaces/vpath";
+import { incrementalVerifier } from "./incrementalUtils";
 
 export function makeDefaultProxy(info: ts.server.PluginCreateInfo): ts.LanguageService {
     const proxy = Object.create(/*o*/ null); // eslint-disable-line no-null/no-null
@@ -1016,7 +1017,8 @@ export class ServerLanguageServiceAdapter implements LanguageServiceAdapter {
             byteLength: Buffer.byteLength,
             hrtime: process.hrtime,
             logger: serverHost,
-            canUseEvents: true
+            canUseEvents: true,
+            incrementalVerifier,
         };
         this.server = new FourslashSession(opts);
 

src/harness/incrementalUtils.ts

@@ -0,0 +1,100 @@
+import * as ts from "./_namespaces/ts";
+
+export function reportDocumentRegistryStats(documentRegistry: ts.DocumentRegistry) {
+    const str: string[] = [];
+    documentRegistry.getBuckets().forEach((bucketEntries, key) => {
+        str.push(`  Key:: ${key}`);
+        bucketEntries.forEach((entry, path) => {
+            if (ts.isDocumentRegistryEntry(entry)) {
+                str.push(`    ${path}: ${ts.Debug.formatScriptKind(entry.sourceFile.scriptKind)} ${entry.languageServiceRefCount}`);
+            }
+            else {
+                entry.forEach((real, kind) => str.push(`    ${path}: ${ts.Debug.formatScriptKind(kind)} ${real.languageServiceRefCount}`));
+            }
+        });
+    });
+    return str;
+}
+
+type DocumentRegistryExpectedStats = Map<ts.DocumentRegistryBucketKeyWithMode, Map<ts.Path, Map<ts.ScriptKind, number>>>;
+function verifyDocumentRegistryStats(
+    documentRegistry: ts.DocumentRegistry,
+    stats: DocumentRegistryExpectedStats,
+) {
+    documentRegistry.getBuckets().forEach((bucketEntries, key) => {
+        const statsByPath = stats.get(key);
+        bucketEntries.forEach((entry, path) => {
+            const expected = statsByPath?.get(path);
+            if (ts.isDocumentRegistryEntry(entry)) {
+                ts.Debug.assert(
+                    expected?.size === 1 && expected.has(entry.sourceFile.scriptKind) && expected.get(entry.sourceFile.scriptKind) === entry.languageServiceRefCount,
+                    `Document registry has unexpected language service ref count for ${key} ${path} ${ts.Debug.formatScriptKind(entry.sourceFile.scriptKind)} ${entry.languageServiceRefCount}`,
+                    reportStats,
+                );
+            }
+            else {
+                entry.forEach((real, kind) => ts.Debug.assert(
+                    real.languageServiceRefCount === expected?.get(kind),
+                    `Document registry has unexpected language service ref count for ${key} ${path} ${ts.Debug.formatScriptKind(kind)} ${real.languageServiceRefCount}`,
+                    reportStats,
+                ));
+                expected?.forEach((value, kind) => ts.Debug.assert(
+                    entry.has(kind),
+                    `Document registry expected language service ref count for ${key} ${path} ${ts.Debug.formatScriptKind(kind)} ${value}`,
+                    reportStats,
+                ));
+            }
+        });
+        statsByPath?.forEach((_value, path) => ts.Debug.assert(
+            bucketEntries.has(path),
+            `Document registry does not contain entry for ${key}, ${path}`,
+            reportStats,
+        ));
+    });
+    stats.forEach((_value, key) => ts.Debug.assert(
+        documentRegistry.getBuckets().has(key),
+        `Document registry does not contain entry for key: ${key}`,
+        reportStats,
+    ));
+
+    function reportStats() {
+        const str: string[] = ["", "Actual::", ...reportDocumentRegistryStats(documentRegistry)];
+        str.push("Expected::");
+        stats?.forEach((statsByPath, key) => {
+            str.push(`  Key:: ${key}`);
+            statsByPath.forEach((entry, path) => entry.forEach((refCount, kind) => str.push(`    ${path}: ${ts.Debug.formatScriptKind(kind)} ${refCount}`)));
+        });
+        return str.join("\n");
+    }
+}
+
+function verifyDocumentRegistry(service: ts.server.ProjectService) {
+    const stats: DocumentRegistryExpectedStats = new Map();
+    const collectStats = (project: ts.server.Project) => {
+        if (project.autoImportProviderHost) collectStats(project.autoImportProviderHost);
+        if (project.noDtsResolutionProject) collectStats(project.noDtsResolutionProject);
+        const program = project.getCurrentProgram();
+        if (!program) return;
+        const key = service.documentRegistry.getKeyForCompilationSettings(program.getCompilerOptions());
+        program.getSourceFiles().forEach(f => {
+            const keyWithMode = service.documentRegistry.getDocumentRegistryBucketKeyWithMode(key, f.impliedNodeFormat);
+            let mapForKeyWithMode = stats.get(keyWithMode);
+            let result: Map<ts.ScriptKind, number> | undefined;
+            if (mapForKeyWithMode === undefined) {
+                stats.set(keyWithMode, mapForKeyWithMode = new Map());
+                mapForKeyWithMode.set(f.resolvedPath, result = new Map());
+            }
+            else {
+                result = mapForKeyWithMode.get(f.resolvedPath);
+                if (!result) mapForKeyWithMode.set(f.resolvedPath, result = new Map());
+            }
+            result.set(f.scriptKind, (result.get(f.scriptKind) || 0) + 1);
+        });
+    };
+    service.forEachProject(collectStats);
+    verifyDocumentRegistryStats(service.documentRegistry, stats);
+}
+
+export function incrementalVerifier(service: ts.server.ProjectService) {
+    service.verifyDocumentRegistry = () => verifyDocumentRegistry(service);
+}

src/server/editorServices.ts

@@ -590,6 +590,7 @@ export interface ProjectServiceOptions {
     typesMapLocation?: string;
     serverMode?: LanguageServiceMode;
     session: Session<unknown> | undefined;
+    /** @internal */ incrementalVerifier?: (service: ProjectService) => void;
 }
 
 interface OriginalFileInfo { fileName: NormalizedPath; path: Path; }
@@ -989,12 +990,13 @@ export class ProjectService {
     /** @internal */
     readonly session: Session<unknown> | undefined;
 
-
     private performanceEventHandler?: PerformanceEventHandler;
 
     private pendingPluginEnablements?: Map<Project, Promise<BeginEnablePluginResult>[]>;
     private currentPluginEnablementPromise?: Promise<void>;
 
+    /** @internal */ verifyDocumentRegistry = noop;
+
     constructor(opts: ProjectServiceOptions) {
         this.host = opts.host;
         this.logger = opts.logger;
@@ -1057,6 +1059,7 @@ export class ProjectService {
                 watchDirectory: returnNoopFileWatcher,
             } :
             getWatchFactory(this.host, watchLogLevel, log, getDetailWatchInfo);
+        opts.incrementalVerifier?.(this);
     }
 
     toPath(fileName: string) {
@@ -1334,7 +1337,7 @@ export class ProjectService {
     }
 
     /** @internal */
-    private forEachProject(cb: (project: Project) => void) {
+    forEachProject(cb: (project: Project) => void) {
         this.externalProjects.forEach(cb);
         this.configuredProjects.forEach(cb);
         this.inferredProjects.forEach(cb);
@@ -2640,6 +2643,7 @@ export class ProjectService {
     private clearSemanticCache(project: Project) {
         project.resolutionCache.clear();
         project.getLanguageService(/*ensureSynchronized*/ false).cleanupSemanticCache();
+        project.cleanupProgram();
         project.markAsDirty();
     }
 

src/server/project.ts

@@ -393,7 +393,7 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
     private packageJsonsForAutoImport: Set<string> | undefined;
 
     /** @internal */
-    private noDtsResolutionProject?: AuxiliaryProject | undefined;
+    noDtsResolutionProject?: AuxiliaryProject | undefined;
 
     /** @internal */
     getResolvedProjectReferenceToRedirect(_fileName: string): ResolvedProjectReference | undefined {
@@ -969,13 +969,27 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
         this.projectService.onUpdateLanguageServiceStateForProject(this, /*languageServiceEnabled*/ true);
     }
 
+    /** @internal */
+    cleanupProgram() {
+        if (this.program) {
+            // Root files are always attached to the project irrespective of program
+            for (const f of this.program.getSourceFiles()) {
+                this.detachScriptInfoIfNotRoot(f.fileName);
+            }
+            this.program.forEachResolvedProjectReference(ref =>
+                this.detachScriptInfoFromProject(ref.sourceFile.fileName));
+            this.program = undefined;
+        }
+    }
+
     disableLanguageService(lastFileExceededProgramSize?: string) {
         if (!this.languageServiceEnabled) {
             return;
         }
         Debug.assert(this.projectService.serverMode !== LanguageServiceMode.Syntactic);
         this.languageService.cleanupSemanticCache();
         this.languageServiceEnabled = false;
+        this.cleanupProgram();
         this.lastFileExceededProgramSize = lastFileExceededProgramSize;
         this.builderState = undefined;
         if (this.autoImportProviderHost) {
@@ -984,6 +998,7 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
         this.autoImportProviderHost = undefined;
         this.resolutionCache.closeTypeRootsWatch();
         this.clearGeneratedFileWatch();
+        this.projectService.verifyDocumentRegistry();
         this.projectService.onUpdateLanguageServiceStateForProject(this, /*languageServiceEnabled*/ false);
     }
 
@@ -1030,17 +1045,10 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
     close() {
         this.projectService.typingsCache.onProjectClosed(this);
         this.closeWatchingTypingLocations();
-        if (this.program) {
-            // if we have a program - release all files that are enlisted in program but arent root
-            // The releasing of the roots happens later
-            // The project could have pending update remaining and hence the info could be in the files but not in program graph
-            for (const f of this.program.getSourceFiles()) {
-                this.detachScriptInfoIfNotRoot(f.fileName);
-            }
-            this.program.forEachResolvedProjectReference(ref =>
-                this.detachScriptInfoFromProject(ref.sourceFile.fileName));
-        }
-
+        // if we have a program - release all files that are enlisted in program but arent root
+        // The releasing of the roots happens later
+        // The project could have pending update remaining and hence the info could be in the files but not in program graph
+        this.cleanupProgram();
         // Release external files
         forEach(this.externalFiles, externalFile => this.detachScriptInfoIfNotRoot(externalFile));
         // Always remove root files from the project
@@ -1492,6 +1500,7 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
 
     private updateGraphWorker() {
         const oldProgram = this.languageService.getCurrentProgram();
+        Debug.assert(oldProgram === this.program);
         Debug.assert(!this.isClosed(), "Called update graph worker of closed project");
         this.writeLog(`Starting updateGraphWorker: Project: ${this.getProjectName()}`);
         const start = timestamp();
@@ -1633,6 +1642,8 @@ export abstract class Project implements LanguageServiceHost, ModuleResolutionHo
         else if (this.program !== oldProgram) {
             this.writeLog(`Different program with same set of files`);
         }
+        // Verify the document registry count
+        this.projectService.verifyDocumentRegistry();
         return hasNewProgram;
     }
 
@@ -2342,7 +2353,8 @@ export class InferredProject extends Project {
     }
 }
 
-class AuxiliaryProject extends Project {
+/** @internal */
+export class AuxiliaryProject extends Project {
     constructor(projectService: ProjectService, documentRegistry: DocumentRegistry, compilerOptions: CompilerOptions, currentDirectory: string) {
         super(projectService.newAuxiliaryProjectName(),
             ProjectKind.Auxiliary,
@@ -2361,7 +2373,6 @@ class AuxiliaryProject extends Project {
         return true;
     }
 
-    /** @internal */
     override scheduleInvalidateResolutionsOfFailedLookupLocations(): void {
         // Invalidation will happen on-demand as part of updateGraph
         return;

src/server/session.ts

@@ -946,6 +946,7 @@ export interface SessionOptions {
     pluginProbeLocations?: readonly string[];
     allowLocalPluginLoads?: boolean;
     typesMapLocation?: string;
+    /** @internal */ incrementalVerifier?: (service: ProjectService) => void;
 }
 
 export class Session<TMessage = string> implements EventSender {
@@ -1010,7 +1011,8 @@ export class Session<TMessage = string> implements EventSender {
             allowLocalPluginLoads: opts.allowLocalPluginLoads,
             typesMapLocation: opts.typesMapLocation,
             serverMode: opts.serverMode,
-            session: this
+            session: this,
+            incrementalVerifier: opts.incrementalVerifier,
         };
         this.projectService = new ProjectService(settings);
         this.projectService.setPerformanceEventHandler(this.performanceEventHandler.bind(this));
@@ -1339,6 +1341,7 @@ export class Session<TMessage = string> implements EventSender {
         this.logger.info(`cleaning ${caption}`);
         for (const p of projects) {
             p.getLanguageService(/*ensureSynchronized*/ false).cleanupSemanticCache();
+            p.cleanupProgram();
         }
     }