Weave cancellation through test discovery

Author: plemarquand

src/TestExplorer/TestExplorer.ts

@@ -37,6 +37,7 @@ export class TestExplorer {
     private lspTestDiscovery: LSPTestDiscovery;
     private subscriptions: { dispose(): unknown }[];
     private testFileEdited = true;
+    private tokenSource = new vscode.CancellationTokenSource();
 
     // Emits after the `vscode.TestController` has been updated.
     private onTestItemsDidChangeEmitter = new vscode.EventEmitter<vscode.TestController>();
@@ -56,7 +57,7 @@ export class TestExplorer {
 
         this.controller.resolveHandler = async item => {
             if (!item) {
-                await this.discoverTestsInWorkspace();
+                await this.discoverTestsInWorkspace(this.tokenSource.token);
             }
         };
 
@@ -86,7 +87,7 @@ export class TestExplorer {
                 this.testFileEdited = false;
                 // only run discover tests if the library has tests
                 if (this.folderContext.swiftPackage.getTargets(TargetType.test).length > 0) {
-                    this.discoverTestsInWorkspace();
+                    this.discoverTestsInWorkspace(this.tokenSource.token);
                 }
             }
         });
@@ -99,6 +100,7 @@ export class TestExplorer {
         });
 
         this.subscriptions = [
+            this.tokenSource,
             fileWatcher,
             onDidEndTask,
             this.controller,
@@ -110,6 +112,9 @@ export class TestExplorer {
     }
 
     dispose() {
+        this.controller.refreshHandler = undefined;
+        this.controller.resolveHandler = undefined;
+        this.tokenSource.cancel();
         this.subscriptions.forEach(element => element.dispose());
     }
 
@@ -120,47 +125,64 @@ export class TestExplorer {
      * @returns Observer disposable
      */
     static observeFolders(workspaceContext: WorkspaceContext): vscode.Disposable {
-        return workspaceContext.onDidChangeFolders(({ folder, operation, workspace }) => {
-            switch (operation) {
-                case FolderOperation.add:
-                    if (folder) {
-                        if (folder.swiftPackage.getTargets(TargetType.test).length > 0) {
-                            folder.addTestExplorer();
-                            // discover tests in workspace but only if disableAutoResolve is not on.
-                            // discover tests will kick off a resolve if required
-                            if (!configuration.folder(folder.workspaceFolder).disableAutoResolve) {
-                                folder.testExplorer?.discoverTestsInWorkspace();
+        const tokenSource = new vscode.CancellationTokenSource();
+        const disposable = workspaceContext.onDidChangeFolders(
+            ({ folder, operation, workspace }) => {
+                switch (operation) {
+                    case FolderOperation.add:
+                        if (folder) {
+                            if (folder.swiftPackage.getTargets(TargetType.test).length > 0) {
+                                folder.addTestExplorer();
+                                // discover tests in workspace but only if disableAutoResolve is not on.
+                                // discover tests will kick off a resolve if required
+                                if (
+                                    !configuration.folder(folder.workspaceFolder).disableAutoResolve
+                                ) {
+                                    folder.testExplorer?.discoverTestsInWorkspace(
+                                        tokenSource.token
+                                    );
+                                }
                             }
                         }
-                    }
-                    break;
-                case FolderOperation.packageUpdated:
-                    if (folder) {
-                        const hasTestTargets =
-                            folder.swiftPackage.getTargets(TargetType.test).length > 0;
-                        if (hasTestTargets && !folder.hasTestExplorer()) {
-                            folder.addTestExplorer();
-                            // discover tests in workspace but only if disableAutoResolve is not on.
-                            // discover tests will kick off a resolve if required
-                            if (!configuration.folder(folder.workspaceFolder).disableAutoResolve) {
-                                folder.testExplorer?.discoverTestsInWorkspace();
+                        break;
+                    case FolderOperation.packageUpdated:
+                        if (folder) {
+                            const hasTestTargets =
+                                folder.swiftPackage.getTargets(TargetType.test).length > 0;
+                            if (hasTestTargets && !folder.hasTestExplorer()) {
+                                folder.addTestExplorer();
+                                // discover tests in workspace but only if disableAutoResolve is not on.
+                                // discover tests will kick off a resolve if required
+                                if (
+                                    !configuration.folder(folder.workspaceFolder).disableAutoResolve
+                                ) {
+                                    folder.testExplorer?.discoverTestsInWorkspace(
+                                        tokenSource.token
+                                    );
+                                }
+                            } else if (!hasTestTargets && folder.hasTestExplorer()) {
+                                folder.removeTestExplorer();
+                            } else if (folder.hasTestExplorer()) {
+                                folder.refreshTestExplorer();
                             }
-                        } else if (!hasTestTargets && folder.hasTestExplorer()) {
-                            folder.removeTestExplorer();
-                        } else if (folder.hasTestExplorer()) {
-                            folder.refreshTestExplorer();
                         }
-                    }
-                    break;
-                case FolderOperation.focus:
-                    if (folder) {
-                        workspace.languageClientManager.documentSymbolWatcher = (
-                            document,
-                            symbols
-                        ) => TestExplorer.onDocumentSymbols(folder, document, symbols);
-                    }
+                        break;
+                    case FolderOperation.focus:
+                        if (folder) {
+                            workspace.languageClientManager.documentSymbolWatcher = (
+                                document,
+                                symbols
+                            ) => TestExplorer.onDocumentSymbols(folder, document, symbols);
+                        }
+                }
             }
-        });
+        );
+        return {
+            dispose: () => {
+                tokenSource.dispose();
+                disposable.dispose();
+            },
+        };
     }
 
     /**
@@ -224,25 +246,25 @@ export class TestExplorer {
     /**
      * Discover tests
      */
-    async discoverTestsInWorkspace() {
+    async discoverTestsInWorkspace(token: vscode.CancellationToken) {
         try {
             // If the LSP cannot produce a list of tests it throws and
             // we fall back to discovering tests with SPM.
-            await this.discoverTestsInWorkspaceLSP();
+            await this.discoverTestsInWorkspaceLSP(token);
         } catch {
             this.folderContext.workspaceContext.outputChannel.logDiagnostic(
                 "workspace/tests LSP request not supported, falling back to SPM to discover tests.",
                 "Test Discovery"
             );
-            await this.discoverTestsInWorkspaceSPM();
+            await this.discoverTestsInWorkspaceSPM(token);
         }
     }
 
     /**
      * Discover tests
      * Uses `swift test --list-tests` to get the list of tests
      */
-    async discoverTestsInWorkspaceSPM() {
+    async discoverTestsInWorkspaceSPM(token: vscode.CancellationToken) {
         async function runDiscover(explorer: TestExplorer, firstTry: boolean) {
             try {
                 const toolchain = explorer.folderContext.workspaceContext.toolchain;
@@ -263,6 +285,11 @@ export class TestExplorer {
                         return;
                     }
                 }
+
+                if (token.isCancellationRequested) {
+                    return;
+                }
+
                 // get list of tests from `swift test --list-tests`
                 let listTestArguments: string[];
                 if (toolchain.swiftVersion.isGreaterThanOrEqual(new Version(5, 8, 0))) {
@@ -284,7 +311,7 @@ export class TestExplorer {
                         explorer.updateTests(explorer.controller, tests);
                     }
                 );
-                await explorer.folderContext.taskQueue.queueOperation(listTestsOperation);
+                await explorer.folderContext.taskQueue.queueOperation(listTestsOperation, token);
             } catch (error) {
                 // If a test list fails its possible the tests have not been built.
                 // Build them and try again, and if we still fail then notify the user.
@@ -336,10 +363,14 @@ export class TestExplorer {
     /**
      * Discover tests
      */
-    async discoverTestsInWorkspaceLSP() {
+    async discoverTestsInWorkspaceLSP(token: vscode.CancellationToken) {
         const tests = await this.lspTestDiscovery.getWorkspaceTests(
             this.folderContext.swiftPackage
         );
+        if (token.isCancellationRequested) {
+            return;
+        }
+
         TestDiscovery.updateTestsFromClasses(
             this.controller,
             this.folderContext.swiftPackage,

src/tasks/TaskQueue.ts

@@ -85,6 +85,9 @@ export class TaskOperation implements SwiftOperation {
         workspaceContext: WorkspaceContext,
         token?: vscode.CancellationToken
     ): Promise<number | undefined> {
+        if (token?.isCancellationRequested) {
+            return Promise.resolve(undefined);
+        }
         workspaceContext.outputChannel.log(`Exec Task: ${this.task.detail ?? this.task.name}`);
         return workspaceContext.tasks.executeTaskAndWait(this.task, token);
     }
@@ -230,6 +233,7 @@ export class TaskQueue {
         if (!this.activeOperation) {
             // get task from queue
             const operation = this.queue.shift();
+
             if (operation) {
                 //const task = operation.task;
                 this.activeOperation = operation;
@@ -250,20 +254,14 @@ export class TaskQueue {
                     .run(this.workspaceContext)
                     .then(result => {
                         // log result
-                        if (operation.log) {
+                        if (operation.log && !operation.token?.isCancellationRequested) {
                             switch (result) {
                                 case 0:
                                     this.workspaceContext.outputChannel.log(
                                         `${operation.log}: ... done.`,
                                         this.folderContext.name
                                     );
                                     break;
-                                case undefined:
-                                    this.workspaceContext.outputChannel.log(
-                                        `${operation.log}: ... cancelled.`,
-                                        this.folderContext.name
-                                    );
-                                    break;
                                 default:
                                     this.workspaceContext.outputChannel.log(
                                         `${operation.log}: ... failed.`,

src/ui/SwiftOutputChannel.ts

@@ -30,7 +30,7 @@ export class SwiftOutputChannel implements vscode.OutputChannel {
     ) {
         this.name = name;
         this.logToConsole = process.env["CI"] !== "1" && logToConsole;
-        this.channel = vscode.window.createOutputChannel(name);
+        this.channel = vscode.window.createOutputChannel(name, "Swift");
         this.logStore = new RollingLog(logStoreLinesSize);
     }
 

test/integration-tests/ui/PackageDependencyProvider.test.ts

@@ -20,7 +20,6 @@ import {
 import { executeTaskAndWaitForResult, waitForNoRunningTasks } from "../../utilities";
 import { getBuildAllTask, SwiftTask } from "../../../src/tasks/SwiftTaskProvider";
 import { testAssetPath } from "../../fixtures";
-import { Version } from "../../../src/utilities/version";
 import {
     activateExtension,
     deactivateExtension,
@@ -33,10 +32,6 @@ suite("PackageDependencyProvider Test Suite", function () {
 
     suiteSetup(async function () {
         const workspaceContext = await activateExtension();
-        // workspace-state.json was not introduced until swift 5.7
-        if (workspaceContext.toolchain.swiftVersion.isLessThan(new Version(5, 7, 0))) {
-            this.skip();
-        }
         await waitForNoRunningTasks();
         const folderContext = await folderInRootWorkspace("dependencies", workspaceContext);
         await executeTaskAndWaitForResult((await getBuildAllTask(folderContext)) as SwiftTask);
@@ -45,7 +40,7 @@ suite("PackageDependencyProvider Test Suite", function () {
     });
 
     suiteTeardown(async () => {
-        treeProvider?.dispose();
+        treeProvider.dispose();
         await deactivateExtension();
     });