testing: handle default profile changes in watch, nested cancellations (microsoft#203583)

Fixes microsoft#203496
Fixes microsoft#203497

Author: connor4312

src/vs/base/common/prefixTree.ts

@@ -48,33 +48,57 @@ export class WellDefinedPrefixTree<V> {
 
 	/** Deletes a node from the prefix tree, returning the value it contained. */
 	delete(key: Iterable<string>): V | undefined {
-		const path = [{ part: '', node: this.root }];
-		let i = 0;
-		for (const part of key) {
-			const node = path[i].node.children?.get(part);
-			if (!node) {
-				return undefined; // node not in tree
-			}
-
-			path.push({ part, node });
-			i++;
+		const path = this.getPathToKey(key);
+		if (!path) {
+			return;
 		}
 
+		let i = path.length - 1;
 		const value = path[i].node._value;
 		if (value === unset) {
 			return; // not actually a real node
 		}
 
 		this._size--;
+		path[i].node._value = unset;
+
 		for (; i > 0; i--) {
+			const { node, part } = path[i];
+			if (node.children?.size || node._value !== unset) {
+				break;
+			}
+
+			path[i - 1].node.children!.delete(part);
+		}
+
+		return value;
+	}
+
+	/** Deletes a subtree from the prefix tree, returning the values they contained. */
+	*deleteRecursive(key: Iterable<string>): Iterable<V> {
+		const path = this.getPathToKey(key);
+		if (!path) {
+			return;
+		}
+
+		const subtree = path[path.length - 1].node;
+
+		// important: run the deletion before we start to yield results, so that
+		// it still runs even if the caller doesn't consumer the iterator
+		for (let i = path.length - 1; i > 0; i--) {
 			const parent = path[i - 1];
 			parent.node.children!.delete(path[i].part);
 			if (parent.node.children!.size > 0 || parent.node._value !== unset) {
 				break;
 			}
 		}
 
-		return value;
+		for (const node of bfsIterate(subtree)) {
+			if (node._value !== unset) {
+				this._size--;
+				yield node._value;
+			}
+		}
 	}
 
 	/** Gets a value from the tree. */
@@ -140,6 +164,22 @@ export class WellDefinedPrefixTree<V> {
 		return node._value !== unset;
 	}
 
+	private getPathToKey(key: Iterable<string>) {
+		const path = [{ part: '', node: this.root }];
+		let i = 0;
+		for (const part of key) {
+			const node = path[i].node.children?.get(part);
+			if (!node) {
+				return; // node not in tree
+			}
+
+			path.push({ part, node });
+			i++;
+		}
+
+		return path;
+	}
+
 	private opNode(key: Iterable<string>, fn: (node: Node<V>) => void, onDescend?: (node: Node<V>) => void): void {
 		let node = this.root;
 		for (const part of key) {
@@ -157,26 +197,31 @@ export class WellDefinedPrefixTree<V> {
 			onDescend?.(node);
 		}
 
-		if (node._value === unset) {
-			this._size++;
-		}
-
+		const sizeBefore = node._value === unset ? 0 : 1;
 		fn(node);
+		const sizeAfter = node._value === unset ? 0 : 1;
+		this._size += sizeAfter - sizeBefore;
 	}
 
 	/** Returns an iterable of the tree values in no defined order. */
 	*values() {
-		const stack = [this.root];
-		while (stack.length > 0) {
-			const node = stack.pop()!;
-			if (node._value !== unset) {
-				yield node._value;
+		for (const { _value } of bfsIterate(this.root)) {
+			if (_value !== unset) {
+				yield _value;
 			}
+		}
+	}
+}
+
+function* bfsIterate<T>(root: Node<T>): Iterable<Node<T>> {
+	const stack = [root];
+	while (stack.length > 0) {
+		const node = stack.pop()!;
+		yield node;
 
-			if (node.children) {
-				for (const child of node.children.values()) {
-					stack.push(child);
-				}
+		if (node.children) {
+			for (const child of node.children.values()) {
+				stack.push(child);
 			}
 		}
 	}

src/vs/base/test/common/prefixTree.test.ts

@@ -137,4 +137,26 @@ suite('WellDefinedPrefixTree', () => {
 
 		assert.deepStrictEqual([...tree.values()], [43, 42]);
 	});
+
+
+	test('delete recursive', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'bar', 'baz'];
+		const key3 = ['foo', 'bar', 'baz2', 'baz3'];
+		const key4 = ['foo', 'bar2'];
+		tree.insert(key1, 42);
+		tree.insert(key2, 43);
+		tree.insert(key3, 44);
+		tree.insert(key4, 45);
+		assert.strictEqual(tree.size, 4);
+
+		assert.deepStrictEqual([...tree.deleteRecursive(key1)], [42, 44, 43]);
+		assert.strictEqual(tree.size, 1);
+
+		assert.deepStrictEqual([...tree.deleteRecursive(key1)], []);
+		assert.strictEqual(tree.size, 1);
+
+		assert.deepStrictEqual([...tree.deleteRecursive(key4)], [45]);
+		assert.strictEqual(tree.size, 0);
+	});
 });

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

@@ -290,21 +290,14 @@ export class ContinuousRunTestAction extends Action2 {
 
 	public override async run(accessor: ServicesAccessor, ...elements: TestItemTreeElement[]): Promise<any> {
 		const crService = accessor.get(ITestingContinuousRunService);
-		const profileService = accessor.get(ITestProfileService);
 		for (const element of elements) {
 			const id = element.test.item.extId;
 			if (crService.isSpecificallyEnabledFor(id)) {
 				crService.stop(id);
 				continue;
 			}
 
-			const profiles = profileService.getGroupDefaultProfiles(TestRunProfileBitset.Run)
-				.filter(p => p.supportsContinuousRun && p.controllerId === element.test.controllerId);
-			if (!profiles.length) {
-				continue;
-			}
-
-			crService.start(profiles, id);
+			crService.start(TestRunProfileBitset.Run, id);
 		}
 	}
 }

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

@@ -4,18 +4,20 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
-import { Disposable, toDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
 import { StoredValue } from 'vs/workbench/contrib/testing/common/storedValue';
 import { TestingContextKeys } from 'vs/workbench/contrib/testing/common/testingContextKeys';
 import { ITestService } from 'vs/workbench/contrib/testing/common/testService';
 import { TestService } from 'vs/workbench/contrib/testing/common/testServiceImpl';
-import { ITestRunProfile } from 'vs/workbench/contrib/testing/common/testTypes';
+import { ITestRunProfile, TestRunProfileBitset } from 'vs/workbench/contrib/testing/common/testTypes';
 import { Emitter, Event } from 'vs/base/common/event';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
 import { WellDefinedPrefixTree } from 'vs/base/common/prefixTree';
+import { ITestProfileService } from 'vs/workbench/contrib/testing/common/testProfileService';
+import * as arrays from 'vs/base/common/arrays';
 
 export const ITestingContinuousRunService = createDecorator<ITestingContinuousRunService>('testingContinuousRunService');
 
@@ -56,10 +58,11 @@ export interface ITestingContinuousRunService {
 	isEnabled(): boolean;
 
 	/**
-	 * Starts a continuous auto run with a specific profile or set of profiles.
-	 * Globally if no test is given, for a specific test otherwise.
+	 * Starts a continuous auto run with a specific set of profiles, or all
+	 * default profiles in a group. Globally if no test is given,
+	 * for a specific test otherwise.
 	 */
-	start(profile: ITestRunProfile[], testId?: string): void;
+	start(profile: ITestRunProfile[] | TestRunProfileBitset, testId?: string): void;
 
 	/**
 	 * Stops any continuous run
@@ -72,8 +75,8 @@ export class TestingContinuousRunService extends Disposable implements ITestingC
 	declare readonly _serviceBrand: undefined;
 
 	private readonly changeEmitter = new Emitter<string | undefined>();
-	private globallyRunning?: CancellationTokenSource;
-	private readonly running = new WellDefinedPrefixTree<CancellationTokenSource>();
+	private globallyRunning?: IDisposable;
+	private readonly running = new WellDefinedPrefixTree<IDisposable>();
 	private readonly lastRun: StoredValue<Set<number>>;
 	private readonly isGloballyOn: IContextKey<boolean>;
 
@@ -87,6 +90,7 @@ export class TestingContinuousRunService extends Disposable implements ITestingC
 		@ITestService private readonly testService: TestService,
 		@IStorageService storageService: IStorageService,
 		@IContextKeyService contextKeyService: IContextKeyService,
+		@ITestProfileService private readonly testProfileService: ITestProfileService,
 	) {
 		super();
 		this.isGloballyOn = TestingContextKeys.isContinuousModeOn.bindTo(contextKeyService);
@@ -133,45 +137,69 @@ export class TestingContinuousRunService extends Disposable implements ITestingC
 	}
 
 	/** @inheritdoc */
-	public start(profile: ITestRunProfile[], testId?: string): void {
+	public start(profiles: ITestRunProfile[] | TestRunProfileBitset, testId?: string): void {
+		const store = new DisposableStore();
 		const cts = new CancellationTokenSource();
+		store.add(toDisposable(() => cts.dispose(true)));
 
 		if (testId === undefined) {
 			this.isGloballyOn.set(true);
 		}
 
 		if (!testId) {
-			this.globallyRunning?.dispose(true);
-			this.globallyRunning = cts;
+			this.globallyRunning?.dispose();
+			this.globallyRunning = store;
 		} else {
 			this.running.mutate(TestId.fromString(testId).path, c => {
-				c?.dispose(true);
-				return cts;
+				c?.dispose();
+				return store;
 			});
 		}
 
-		this.lastRun.store(new Set(profile.map(p => p.profileId)));
+		let actualProfiles: ITestRunProfile[];
+		if (profiles instanceof Array) {
+			actualProfiles = profiles;
+		} else {
+			// restart the continuous run when default profiles change, if we were
+			// asked to run for a group
+			const getRelevant = () => this.testProfileService.getGroupDefaultProfiles(profiles)
+				.filter(p => p.supportsContinuousRun && (!testId || TestId.root(testId) === p.controllerId));
+			actualProfiles = getRelevant();
+			store.add(this.testProfileService.onDidChange(() => {
+				if (!arrays.equals(getRelevant(), actualProfiles)) {
+					this.start(profiles, testId);
+				}
+			}));
+		}
 
-		this.testService.startContinuousRun({
-			continuous: true,
-			targets: profile.map(p => ({
-				testIds: [testId ?? p.controllerId],
-				controllerId: p.controllerId,
-				profileGroup: p.group,
-				profileId: p.profileId
-			})),
-		}, cts.token);
+		this.lastRun.store(new Set(actualProfiles.map(p => p.profileId)));
+
+		if (actualProfiles.length) {
+			this.testService.startContinuousRun({
+				continuous: true,
+				targets: actualProfiles.map(p => ({
+					testIds: [testId ?? p.controllerId],
+					controllerId: p.controllerId,
+					profileGroup: p.group,
+					profileId: p.profileId
+				})),
+			}, cts.token);
+		}
 
 		this.changeEmitter.fire(testId);
 	}
 
 	/** @inheritdoc */
 	public stop(testId?: string): void {
 		if (!testId) {
-			this.globallyRunning?.dispose(true);
+			this.globallyRunning?.dispose();
 			this.globallyRunning = undefined;
 		} else {
-			this.running.delete(TestId.fromString(testId).path)?.dispose(true);
+			const cancellations = [...this.running.deleteRecursive(TestId.fromString(testId).path)];
+			// deleteRecursive returns a BFS order, reverse it so children are cancelled before parents
+			for (let i = cancellations.length - 1; i >= 0; i--) {
+				cancellations[i].dispose();
+			}
 		}
 
 		if (testId === undefined) {