testing: improve performance when ending test with large number of results

Takes `markTaskComplete` 11,200ms to 70ms when running a 10k test suite,
by maintaining a count of computed states for children and avoiding
refreshing nodes unnecessarily.

For microsoft/vscode-python#21507

Author: connor4312

src/vs/workbench/contrib/testing/browser/testingProgressUiService.ts

@@ -13,8 +13,8 @@ import { ProgressLocation, UnmanagedProgress } from 'vs/platform/progress/common
 import { IViewsService } from 'vs/workbench/common/views';
 import { AutoOpenTesting, getTestingConfiguration, TestingConfigKeys } from 'vs/workbench/contrib/testing/common/configuration';
 import { Testing } from 'vs/workbench/contrib/testing/common/constants';
-import { isFailedState } from 'vs/workbench/contrib/testing/common/testingStates';
-import { LiveTestResult, TestResultItemChangeReason, TestStateCount } from 'vs/workbench/contrib/testing/common/testResult';
+import { isFailedState, TestStateCount } from 'vs/workbench/contrib/testing/common/testingStates';
+import { LiveTestResult, TestResultItemChangeReason } from 'vs/workbench/contrib/testing/common/testResult';
 import { ITestResultService } from 'vs/workbench/contrib/testing/common/testResultService';
 import { TestResultState } from 'vs/workbench/contrib/testing/common/testTypes';
 

src/vs/workbench/contrib/testing/common/getComputedState.ts

@@ -5,7 +5,7 @@
 
 import { Iterable } from 'vs/base/common/iterator';
 import { TestResultState } from 'vs/workbench/contrib/testing/common/testTypes';
-import { maxPriority, statePriority } from 'vs/workbench/contrib/testing/common/testingStates';
+import { makeEmptyCounts, maxPriority, statePriority } from 'vs/workbench/contrib/testing/common/testingStates';
 
 /**
  * Accessor for nodes in get and refresh computed state.
@@ -32,18 +32,28 @@ const isDurationAccessor = <T>(accessor: IComputedStateAccessor<T>): accessor is
  * if it was previously set.
  */
 
-const getComputedState = <T>(accessor: IComputedStateAccessor<T>, node: T, force = false) => {
+const getComputedState = <T extends object>(accessor: IComputedStateAccessor<T>, node: T, force = false) => {
 	let computed = accessor.getCurrentComputedState(node);
 	if (computed === undefined || force) {
 		computed = accessor.getOwnState(node) ?? TestResultState.Unset;
 
+		let childrenCount = 0;
+		const stateMap = makeEmptyCounts();
+
 		for (const child of accessor.getChildren(node)) {
 			const childComputed = getComputedState(accessor, child);
+			childrenCount++;
+			stateMap[childComputed]++;
+
 			// If all children are skipped, make the current state skipped too if unset (#131537)
 			computed = childComputed === TestResultState.Skipped && computed === TestResultState.Unset
 				? TestResultState.Skipped : maxPriority(computed, childComputed);
 		}
 
+		if (childrenCount > LARGE_NODE_THRESHOLD) {
+			largeNodeChildrenStates.set(node, stateMap);
+		}
+
 		accessor.setComputedState(node, computed);
 	}
 
@@ -72,11 +82,19 @@ const getComputedDuration = <T>(accessor: IComputedStateAndDurationAccessor<T>,
 	return computed;
 };
 
+const LARGE_NODE_THRESHOLD = 64;
+
+/**
+ * Map of how many nodes have in each state. This is used to optimize state
+ * computation in large nodes with children above the `LARGE_NODE_THRESHOLD`.
+ */
+const largeNodeChildrenStates = new WeakMap<object, { [K in TestResultState]: number }>();
+
 /**
  * Refreshes the computed state for the node and its parents. Any changes
  * elements cause `addUpdated` to be called.
  */
-export const refreshComputedState = <T>(
+export const refreshComputedState = <T extends object>(
 	accessor: IComputedStateAccessor<T>,
 	node: T,
 	explicitNewComputedState?: TestResultState,
@@ -92,28 +110,46 @@ export const refreshComputedState = <T>(
 		accessor.setComputedState(node, newState);
 		toUpdate.add(node);
 
-		if (newPriority > oldPriority) {
-			// Update all parents to ensure they're at least this priority.
-			for (const parent of accessor.getParents(node)) {
-				const prev = accessor.getCurrentComputedState(parent);
+		let moveFromState = oldState;
+		let moveToState = newState;
+
+		for (const parent of accessor.getParents(node)) {
+			const lnm = largeNodeChildrenStates.get(parent);
+			if (lnm) {
+				lnm[moveFromState]--;
+				lnm[moveToState]++;
+			}
+
+			const prev = accessor.getCurrentComputedState(parent);
+			if (newPriority > oldPriority) {
+				// Update all parents to ensure they're at least this priority.
 				if (prev !== undefined && statePriority[prev] >= newPriority) {
 					break;
 				}
 
+				if (lnm && lnm[moveToState] > 0) {
+					break;
+				}
+
+				// moveToState remains the same, the new higher priority node state
 				accessor.setComputedState(parent, newState);
 				toUpdate.add(parent);
-			}
-		} else if (newPriority < oldPriority) {
-			// Re-render all parents of this node whose computed priority might have come from this node
-			for (const parent of accessor.getParents(node)) {
-				const prev = accessor.getCurrentComputedState(parent);
+			} else /* newProirity < oldPriority */ {
+				// Update all parts whose statese might have been based on this one
 				if (prev === undefined || statePriority[prev] > oldPriority) {
 					break;
 				}
 
-				accessor.setComputedState(parent, getComputedState(accessor, parent, true));
+				if (lnm && lnm[moveFromState] > 0) {
+					break;
+				}
+
+				moveToState = getComputedState(accessor, parent, true);
+				accessor.setComputedState(parent, moveToState);
 				toUpdate.add(parent);
 			}
+
+			moveFromState = prev;
 		}
 	}
 

src/vs/workbench/contrib/testing/common/testResult.ts

@@ -14,7 +14,7 @@ import { IComputedStateAccessor, refreshComputedState } from 'vs/workbench/contr
 import { IObservableValue, MutableObservableValue, staticObservableValue } from 'vs/workbench/contrib/testing/common/observableValue';
 import { TestCoverage } from 'vs/workbench/contrib/testing/common/testCoverage';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
-import { maxPriority, statesInOrder, terminalStatePriorities } from 'vs/workbench/contrib/testing/common/testingStates';
+import { makeEmptyCounts, maxPriority, statesInOrder, terminalStatePriorities, TestStateCount } from 'vs/workbench/contrib/testing/common/testingStates';
 import { getMarkId, IRichLocation, ISerializedTestResults, ITestItem, ITestMessage, ITestOutputMessage, ITestRunTask, ITestTaskState, ResolvedTestRunRequest, TestItemExpandState, TestMessageType, TestResultItem, TestResultState } from 'vs/workbench/contrib/testing/common/testTypes';
 
 export interface ITestRunTaskResults extends ITestRunTask {
@@ -185,20 +185,6 @@ export const resultItemParents = function* (results: ITestResult, item: TestResu
 	}
 };
 
-/**
- * Count of the number of tests in each run state.
- */
-export type TestStateCount = { [K in TestResultState]: number };
-
-export const makeEmptyCounts = () => {
-	const o: Partial<TestStateCount> = {};
-	for (const state of statesInOrder) {
-		o[state] = 0;
-	}
-
-	return o as TestStateCount;
-};
-
 export const maxCountPriority = (counts: Readonly<TestStateCount>) => {
 	for (const state of statesInOrder) {
 		if (counts[state] > 0) {
@@ -266,7 +252,7 @@ export class LiveTestResult implements ITestResult {
 	/**
 	 * @inheritdoc
 	 */
-	public readonly counts: { [K in TestResultState]: number } = makeEmptyCounts();
+	public readonly counts = makeEmptyCounts();
 
 	/**
 	 * @inheritdoc

src/vs/workbench/contrib/testing/common/testingStates.ts

@@ -69,3 +69,13 @@ export const terminalStatePriorities: { [key in TestResultState]?: number } = {
 	[TestResultState.Failed]: 2,
 	[TestResultState.Errored]: 3,
 };
+
+/**
+ * Count of the number of tests in each run state.
+ */
+export type TestStateCount = { [K in TestResultState]: number };
+
+export const makeEmptyCounts = (): TestStateCount => {
+	// shh! don't tell anyone this is actually an array!
+	return new Uint32Array(statesInOrder.length) as any as { [K in TestResultState]: number };
+};

src/vs/workbench/contrib/testing/test/common/testResultService.test.ts

@@ -10,10 +10,11 @@ import { MockContextKeyService } from 'vs/platform/keybinding/test/common/mockKe
 import { NullLogService } from 'vs/platform/log/common/log';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
 import { TestProfileService } from 'vs/workbench/contrib/testing/common/testProfileService';
-import { HydratedTestResult, LiveTestResult, makeEmptyCounts, resultItemParents, TaskRawOutput, TestResultItemChange, TestResultItemChangeReason } from 'vs/workbench/contrib/testing/common/testResult';
+import { HydratedTestResult, LiveTestResult, resultItemParents, TaskRawOutput, TestResultItemChange, TestResultItemChangeReason } from 'vs/workbench/contrib/testing/common/testResult';
 import { TestResultService } from 'vs/workbench/contrib/testing/common/testResultService';
 import { InMemoryResultStorage, ITestResultStorage } from 'vs/workbench/contrib/testing/common/testResultStorage';
 import { ITestTaskState, ResolvedTestRunRequest, TestResultItem, TestResultState, TestRunProfileBitset } from 'vs/workbench/contrib/testing/common/testTypes';
+import { makeEmptyCounts } from 'vs/workbench/contrib/testing/common/testingStates';
 import { getInitializedMainTestCollection, testStubs, TestTestCollection } from 'vs/workbench/contrib/testing/test/common/testStubs';
 import { TestStorageService } from 'vs/workbench/test/common/workbenchTestServices';
 
@@ -93,27 +94,24 @@ suite('Workbench - Test Results Service', () => {
 		});
 
 		test('initializes with valid counts', () => {
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Unset]: 4,
-			});
+			const c = makeEmptyCounts();
+			c[TestResultState.Unset] = 4;
+			assert.deepStrictEqual(r.counts, c);
 		});
 
 		test('setAllToState', () => {
 			changed.clear();
 			r.setAllToStatePublic(TestResultState.Queued, 't', (_, t) => t.item.label !== 'root');
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Unset]: 1,
-				[TestResultState.Queued]: 3,
-			});
+			const c = makeEmptyCounts();
+			c[TestResultState.Unset] = 1;
+			c[TestResultState.Queued] = 3;
+			assert.deepStrictEqual(r.counts, c);
 
 			r.setAllToStatePublic(TestResultState.Failed, 't', (_, t) => t.item.label !== 'root');
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Unset]: 1,
-				[TestResultState.Failed]: 3,
-			});
+			const c2 = makeEmptyCounts();
+			c2[TestResultState.Unset] = 1;
+			c2[TestResultState.Failed] = 3;
+			assert.deepStrictEqual(r.counts, c2);
 
 			assert.deepStrictEqual(r.getStateById(new TestId(['ctrlId', 'id-a']).toString())?.ownComputedState, TestResultState.Failed);
 			assert.deepStrictEqual(r.getStateById(new TestId(['ctrlId', 'id-a']).toString())?.tasks[0].state, TestResultState.Failed);
@@ -134,11 +132,10 @@ suite('Workbench - Test Results Service', () => {
 			changed.clear();
 			const testId = new TestId(['ctrlId', 'id-a', 'id-aa']).toString();
 			r.updateState(testId, 't', TestResultState.Running);
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Unset]: 3,
-				[TestResultState.Running]: 1,
-			});
+			const c = makeEmptyCounts();
+			c[TestResultState.Running] = 1;
+			c[TestResultState.Unset] = 3;
+			assert.deepStrictEqual(r.counts, c);
 			assert.deepStrictEqual(r.getStateById(testId)?.ownComputedState, TestResultState.Running);
 			// update computed state:
 			assert.deepStrictEqual(r.getStateById(tests.root.id)?.computedState, TestResultState.Running);
@@ -161,10 +158,9 @@ suite('Workbench - Test Results Service', () => {
 		test('ignores outside run', () => {
 			changed.clear();
 			r.updateState(new TestId(['ctrlId', 'id-b']).toString(), 't', TestResultState.Running);
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Unset]: 4,
-			});
+			const c = makeEmptyCounts();
+			c[TestResultState.Unset] = 4;
+			assert.deepStrictEqual(r.counts, c);
 			assert.deepStrictEqual(r.getStateById(new TestId(['ctrlId', 'id-b']).toString()), undefined);
 		});
 
@@ -175,11 +171,10 @@ suite('Workbench - Test Results Service', () => {
 
 			r.markComplete();
 
-			assert.deepStrictEqual(r.counts, {
-				...makeEmptyCounts(),
-				[TestResultState.Passed]: 1,
-				[TestResultState.Unset]: 3,
-			});
+			const c = makeEmptyCounts();
+			c[TestResultState.Unset] = 3;
+			c[TestResultState.Passed] = 1;
+			assert.deepStrictEqual(r.counts, c);
 
 			assert.deepStrictEqual(r.getStateById(tests.root.id)?.ownComputedState, TestResultState.Unset);
 			assert.deepStrictEqual(r.getStateById(new TestId(['ctrlId', 'id-a', 'id-aa']).toString())?.ownComputedState, TestResultState.Passed);