perf: Avoid repeated test discovery for resolved test items

src/commands/testExplorerCommands.ts

@@ -45,7 +45,7 @@ export async function refreshExplorer(): Promise<void> {
     // Force re-resolution of all existing project roots
     const loadPromises: Promise<void>[] = [];
     testController?.items.forEach((root: TestItem) => {
-        loadPromises.push(loadChildren(root));
+        loadPromises.push(loadChildren(root, undefined, true));
     });
     await Promise.all(loadPromises);
 
@@ -90,13 +90,13 @@ export async function refreshProject(classpathUri: Uri): Promise<void> {
 
     if (matchedProject) {
         // Re-resolve only the matched project's children
-        await loadChildren(matchedProject);
+        await loadChildren(matchedProject, undefined, true);
     } else if (childProjectMatched) {
         // The classpath URI is an ancestor containing test projects – refresh all children
         const loadPromises: Promise<void>[] = [];
         testController?.items.forEach((root: TestItem) => {
             if (root.uri && ensureTrailingSeparator(root.uri.toString()).startsWith(uriString)) {
-                loadPromises.push(loadChildren(root));
+                loadPromises.push(loadChildren(root, undefined, true));
             }
         });
         await Promise.all(loadPromises);

src/controller/testController.ts

@@ -16,8 +16,8 @@ import { JUnitLaunchProtocol } from '../constants';
 import { IJavaTestItem } from '../types';
 import { loadRunConfig } from '../utils/configUtils';
 import { resolveLaunchConfigurationForRunner } from '../utils/launchUtils';
-import { dataCache, ITestItemData } from './testItemDataCache';
-import { createTestItem, findDirectTestChildrenForClass, findTestPackagesAndTypes, findTestTypesAndMethods, loadJavaProjects, resolvePath, synchronizeItemsRecursively, updateItemForDocumentWithDebounce } from './utils';
+import { dataCache, getResolutionVersion, invalidateResolutionVersion, ITestItemData } from './testItemDataCache';
+import { createTestItem, findDirectTestChildrenForClass, findTestPackagesAndTypes, findTestTypesAndMethods, loadJavaProjects, removeOutdatedTestItemsForDocument, resolvePath, synchronizeItemsRecursively, updateItemForDocumentWithDebounce } from './utils';
 import { JavaTestCoverageProvider } from '../provider/JavaTestCoverageProvider';
 import { testRunnerService } from './testRunnerService';
 import { IRunTestContext, TestRunner, TestFinishEvent, TestItemStatusChangeEvent, TestKind, TestLevel, TestResultState, TestIdParts } from '../java-test-runner.api';
@@ -27,6 +27,7 @@ import { parsePartsFromTestId } from '../utils/testItemUtils';
 export let testController: TestController | undefined;
 export const watchers: Disposable[] = [];
 export const runnableTag: TestTag = new TestTag('runnable');
+const pendingTestItemResolutions: WeakMap<TestItem, Promise<void>> = new WeakMap();
 
 export function createTestController(): void {
     testController?.dispose();
@@ -51,18 +52,68 @@ export function creatTestProfile(name: string, kind: TestRunProfileKind): void {
     testController?.createRunProfile(name, kind, runHandler, false, runnableTag);
 }
 
-export const loadChildren: (item: TestItem, token?: CancellationToken) => any = instrumentOperation('java.test.explorer.loadChildren', async (_operationId: string, item: TestItem, token?: CancellationToken) => {
+export const loadChildren: (item: TestItem, token?: CancellationToken, force?: boolean) => Promise<void> = instrumentOperation('java.test.explorer.loadChildren', async (_operationId: string, item: TestItem, token?: CancellationToken, force: boolean = false) => {
     if (!item) {
         await loadJavaProjects();
         return;
     }
 
-    const data: ITestItemData | undefined = dataCache.get(item);
-    if (!data) {
+    if (!dataCache.get(item)) {
         return;
     }
+
+    if (token?.isCancellationRequested) {
+        return;
+    }
+
+    if (force) {
+        invalidateTestItemResolution(item);
+    } else if (!item.canResolveChildren) {
+        return;
+    }
+
+    while (item.canResolveChildren) {
+        if (token?.isCancellationRequested) {
+            return;
+        }
+
+        const pendingResolution: Promise<void> | undefined = pendingTestItemResolutions.get(item);
+        if (pendingResolution) {
+            await pendingResolution;
+            if (pendingTestItemResolutions.get(item) === pendingResolution) {
+                pendingTestItemResolutions.delete(item);
+            }
+            continue;
+        }
+
+        const data: ITestItemData | undefined = dataCache.get(item);
+        if (!data) {
+            return;
+        }
+        const resolutionVersion: number = getResolutionVersion(item);
+        const resolution: Promise<void> = resolveTestItemChildren(item, data, resolutionVersion, token);
+        pendingTestItemResolutions.set(item, resolution);
+        try {
+            await resolution;
+        } finally {
+            if (pendingTestItemResolutions.get(item) === resolution) {
+                pendingTestItemResolutions.delete(item);
+            }
+        }
+
+        if (token?.isCancellationRequested || resolutionVersion === getResolutionVersion(item)) {
+            return;
+        }
+    }
+});
+
+async function resolveTestItemChildren(item: TestItem, data: ITestItemData, resolutionVersion: number,
+    token?: CancellationToken): Promise<void> {
     if (data.testLevel === TestLevel.Project) {
         const packageAndTypes: IJavaTestItem[] = await findTestPackagesAndTypes(data.jdtHandler, token);
+        if (token?.isCancellationRequested || resolutionVersion !== getResolutionVersion(item)) {
+            return;
+        }
         synchronizeItemsRecursively(item, packageAndTypes);
     } else if (data.testLevel === TestLevel.Package) {
         // unreachable code
@@ -72,9 +123,30 @@ export const loadChildren: (item: TestItem, token?: CancellationToken) => any =
             return;
         }
         const testMethods: IJavaTestItem[] = await findDirectTestChildrenForClass(data.jdtHandler, token);
+        if (token?.isCancellationRequested || resolutionVersion !== getResolutionVersion(item)) {
+            return;
+        }
         synchronizeItemsRecursively(item, testMethods);
     }
-});
+
+    if (resolutionVersion !== getResolutionVersion(item)) {
+        return;
+    }
+    item.canResolveChildren = false;
+}
+
+function invalidateTestItemResolution(item: TestItem): void {
+    const testLevel: TestLevel | undefined = dataCache.get(item)?.testLevel;
+    if (testLevel === TestLevel.Project || testLevel === TestLevel.Class) {
+        invalidateResolutionVersion(item);
+    }
+    if (testLevel !== undefined && testLevel <= TestLevel.Class) {
+        item.canResolveChildren = true;
+    }
+    item.children.forEach((child: TestItem) => {
+        invalidateTestItemResolution(child);
+    });
+}
 
 async function startWatchingWorkspace(): Promise<void> {
     if (!workspace.workspaceFolders) {
@@ -127,11 +199,7 @@ async function startWatchingWorkspace(): Promise<void> {
                         return;
                     }
 
-                    belongingPackage.children.forEach((item: TestItem) => {
-                        if (item.uri?.toString() === uri.toString()) {
-                            belongingPackage.children.delete(item.id);
-                        }
-                    });
+                    removeOutdatedTestItemsForDocument(belongingPackage, uri, new Set<string>());
 
                     if (belongingPackage.children.size === 0) {
                         belongingProject.children.delete(belongingPackage.id);

src/controller/testItemDataCache.ts

@@ -27,6 +27,16 @@ class TestItemDataCache {
 
 export const dataCache: TestItemDataCache = new TestItemDataCache();
 
+const resolutionVersions: WeakMap<TestItem, number> = new WeakMap();
+
+export function getResolutionVersion(item: TestItem): number {
+    return resolutionVersions.get(item) ?? 0;
+}
+
+export function invalidateResolutionVersion(item: TestItem): void {
+    resolutionVersions.set(item, getResolutionVersion(item) + 1);
+}
+
 export interface ITestItemData {
     jdtHandler: string;
     fullName: string;

src/controller/utils.ts

@@ -12,7 +12,7 @@ import { IJavaTestItem, ProjectType } from '../types';
 import { executeJavaLanguageServerCommand } from '../utils/commandUtils';
 import { getRequestDelay, lruCache, MovingAverage } from './debouncing';
 import { runnableTag, testController } from './testController';
-import { dataCache } from './testItemDataCache';
+import { dataCache, invalidateResolutionVersion } from './testItemDataCache';
 import { TestKind, TestLevel } from '../java-test-runner.api';
 
 /**
@@ -90,30 +90,43 @@ export async function getProjectType(item: IJavaTestItem): Promise<ProjectType>
  * - If an existing child is not contained in the childrenData parameter, it will be deleted
  * - If a child does not exist, create it, otherwise, update it as well as its metadata.
  */
-export function synchronizeItemsRecursively(parent: TestItem, childrenData: IJavaTestItem[] | undefined): void {
-    if (childrenData) {
-        // remove the out-of-date children
-        parent.children.forEach((child: TestItem) => {
-            if (dataCache.get(child)?.testLevel === TestLevel.Invocation) {
-                // only remove the invocation items before a new test session starts
-                return;
-            }
-            const existingItem: IJavaTestItem | undefined = childrenData.find((data: IJavaTestItem) => data.id === child.id);
-            if (!existingItem) {
-                parent.children.delete(child.id);
-            }
-        });
-        // update/create children
-        for (const child of childrenData) {
-            const childItem: TestItem = updateOrCreateTestItem(parent, child);
-            if (child.testLevel <= TestLevel.Class) {
-                childItem.canResolveChildren = true;
-            }
-            synchronizeItemsRecursively(childItem, child.children);
+export function synchronizeItemsRecursively(parent: TestItem, childrenData: IJavaTestItem[] | undefined,
+    childrenAreComplete: boolean = false): void {
+    if (!childrenData && !childrenAreComplete) {
+        return;
+    }
+
+    const children: IJavaTestItem[] = childrenData ?? [];
+    // remove the out-of-date children
+    parent.children.forEach((child: TestItem) => {
+        if (dataCache.get(child)?.testLevel === TestLevel.Invocation) {
+            // only remove the invocation items before a new test session starts
+            return;
+        }
+        const existingItem: IJavaTestItem | undefined = children.find((data: IJavaTestItem) => data.id === child.id);
+        if (!existingItem) {
+            parent.children.delete(child.id);
+        }
+    });
+    // update/create children
+    for (const child of children) {
+        const childItem: TestItem = updateOrCreateTestItem(parent, child);
+        if (child.testLevel <= TestLevel.Class) {
+            childItem.canResolveChildren = true;
         }
+        synchronizeItemsRecursively(childItem, child.children, childrenAreComplete);
     }
 }
 
+export function markTestClassesResolvedRecursively(item: TestItem): void {
+    if (dataCache.get(item)?.testLevel === TestLevel.Class) {
+        item.canResolveChildren = false;
+    }
+    item.children.forEach((child: TestItem) => {
+        markTestClassesResolvedRecursively(child);
+    });
+}
+
 export function updateOrCreateTestItem(parent: TestItem, childData: IJavaTestItem): TestItem {
     let childItem: TestItem | undefined = parent.children.get(childData.id);
     if (childItem) {
@@ -125,6 +138,7 @@ export function updateOrCreateTestItem(parent: TestItem, childData: IJavaTestIte
 }
 
 function updateTestItem(testItem: TestItem, metaInfo: IJavaTestItem): void {
+    const previousJdtHandler: string | undefined = dataCache.get(testItem)?.jdtHandler;
     testItem.range = asRange(metaInfo.range);
     testItem.label = metaInfo.label;
     dataCache.set(testItem, {
@@ -134,6 +148,10 @@ function updateTestItem(testItem: TestItem, metaInfo: IJavaTestItem): void {
         testLevel: metaInfo.testLevel,
         testKind: metaInfo.testKind,
     });
+    if (previousJdtHandler !== undefined && previousJdtHandler !== metaInfo.jdtHandler &&
+            (metaInfo.testLevel === TestLevel.Project || metaInfo.testLevel === TestLevel.Class)) {
+        invalidateResolutionVersion(testItem);
+    }
 }
 
 /**
@@ -217,15 +235,11 @@ export async function updateItemForDocument(uri: Uri, testTypes?: IJavaTestItem[
         return [];
     }
 
+    const expectedTypeIds: Set<string> = new Set(testTypes.map((testType: IJavaTestItem) => testType.id));
+    removeOutdatedTestItemsForDocument(belongingPackage, uri, expectedTypeIds);
+
     const tests: TestItem[] = [];
-    if (testTypes.length === 0) {
-        // Remove the children with the same uri when no test items is found
-        belongingPackage.children.forEach((typeItem: TestItem) => {
-            if (path.relative(typeItem.uri?.fsPath || '', uri.fsPath) === '') {
-                belongingPackage!.children.delete(typeItem.id);
-            }
-        });
-    } else {
+    if (testTypes.length > 0) {
         for (const testType of testTypes) {
             // here we do not directly call synchronizeItemsRecursively() because testTypes here are just part of the
             // children of the belonging package, we don't want to delete other children unexpectedly.
@@ -237,7 +251,8 @@ export async function updateItemForDocument(uri: Uri, testTypes?: IJavaTestItem[
                 updateTestItem(testTypeItem, testType);
             }
             tests.push(testTypeItem);
-            synchronizeItemsRecursively(testTypeItem, testType.children);
+            synchronizeItemsRecursively(testTypeItem, testType.children, true);
+            markTestClassesResolvedRecursively(testTypeItem);
         }
     }
 
@@ -248,6 +263,41 @@ export async function updateItemForDocument(uri: Uri, testTypes?: IJavaTestItem[
     return tests;
 }
 
+export function removeOutdatedTestItemsForDocument(belongingPackage: TestItem, uri: Uri,
+    expectedTypeIds: Set<string>): void {
+    const belongingProject: TestItem | undefined = belongingPackage.parent;
+    if (!belongingProject) {
+        return;
+    }
+
+    invalidateResolutionVersion(belongingProject);
+    belongingProject.children.forEach((packageItem: TestItem) => {
+        packageItem.children.forEach((typeItem: TestItem) => {
+            if (path.relative(typeItem.uri?.fsPath || '', uri.fsPath) !== '') {
+                return;
+            }
+
+            invalidateResolutionRecursively(typeItem);
+            if (!expectedTypeIds.has(typeItem.id)) {
+                packageItem.children.delete(typeItem.id);
+            }
+        });
+
+        if (packageItem !== belongingPackage && packageItem.children.size === 0) {
+            belongingProject.children.delete(packageItem.id);
+        }
+    });
+}
+
+function invalidateResolutionRecursively(item: TestItem): void {
+    if (dataCache.get(item)?.testLevel === TestLevel.Class) {
+        invalidateResolutionVersion(item);
+    }
+    item.children.forEach((child: TestItem) => {
+        invalidateResolutionRecursively(child);
+    });
+}
+
 /**
  * Give a test item for a type, find its belonging package item according to its id.
  */

src/types.ts

@@ -5,7 +5,7 @@ import { Range } from 'vscode';
 import { TestKind, TestLevel } from './java-test-runner.api';
 
 export interface IJavaTestItem {
-    children: IJavaTestItem[];
+    children?: IJavaTestItem[];
     uri: string | undefined;
     range: Range | undefined;
     jdtHandler: string;