Cancelling should invoke a single interrupt (#12106)

For #10496

Runing entire notebook and clicking cancel would result in kernel being interrupted n times (where n = number of cells).
I.e. user clicks cancel once, but we end up interrupting multiple times.

Author: DonJayamanne

src/client/datascience/interactive-ipynb/nativeEditorCommandListener.ts

@@ -9,6 +9,7 @@ import { Uri } from 'vscode';
 
 import { ICommandManager } from '../../common/application/types';
 import { IDisposableRegistry } from '../../common/types';
+import { traceError } from '../../logging';
 import { captureTelemetry } from '../../telemetry';
 import { CommandSource } from '../../testing/common/constants';
 import { Commands, Telemetry } from '../constants';
@@ -104,10 +105,10 @@ export class NativeEditorCommandListener implements IDataScienceCommandListener
         }
     }
 
-    private restartKernel() {
+    private async restartKernel() {
         const activeEditor = this.provider.activeEditor;
         if (activeEditor) {
-            activeEditor.restartKernel().ignoreErrors();
+            await activeEditor.restartKernel().catch(traceError.bind('Failed to restart kernel'));
         }
     }
 

src/client/datascience/notebook/executionService.ts

@@ -36,9 +36,10 @@ const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed'
  */
 @injectable()
 export class NotebookExecutionService implements INotebookExecutionService {
-    private registeredIOPubListeners = new WeakSet<INotebook>();
+    private readonly registeredIOPubListeners = new WeakSet<INotebook>();
     private _notebookProvider?: INotebookProvider;
-    private pendingExecutionCancellations = new Map<string, CancellationTokenSource[]>();
+    private readonly pendingExecutionCancellations = new Map<string, CancellationTokenSource[]>();
+    private readonly tokensInterrupted = new WeakSet<CancellationToken>();
     private get notebookProvider(): INotebookProvider {
         this._notebookProvider =
             this._notebookProvider || this.serviceContainer.get<INotebookProvider>(INotebookProvider);
@@ -133,33 +134,21 @@ export class NotebookExecutionService implements INotebookExecutionService {
         const stopWatch = new StopWatch();
 
         wrappedToken.onCancellationRequested(() => {
-            // tslint:disable-next-line: no-suspicious-comment
-            // TODO: Is this the right thing to do?
-            // I think it is, as we have a stop button.
-            // If we're busy executing, then interrupt the execution.
             if (deferred.completed) {
                 return;
             }
-            deferred.resolve();
-            cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Idle;
 
             // Interrupt kernel only if original cancellation was cancelled.
-            if (token.isCancellationRequested) {
+            // I.e. interrupt kernel only if user attempts to stop the execution by clicking stop button.
+            if (token.isCancellationRequested && !this.tokensInterrupted.has(token)) {
+                this.tokensInterrupted.add(token);
                 this.commandManager.executeCommand(Commands.NotebookEditorInterruptKernel).then(noop, noop);
             }
         });
 
         cell.metadata.runStartTime = new Date().getTime();
         cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Running;
 
-        if (!findMappedNotebookCellModel(cell, model.cells)) {
-            // tslint:disable-next-line: no-suspicious-comment
-            // TODO: Possible it was added as we didn't get to know about it.
-            // We need to handle these.
-            // Basically if there's a new cell, we need to first add it into our model,
-            // Similarly we might want to handle deletions.
-            throw new Error('Unable to find corresonding Cell in Model');
-        }
         let subscription: Subscription | undefined;
         try {
             nb.clear(cell.uri.fsPath); // NOSONAR

src/test/datascience/notebook/helper.ts

@@ -16,7 +16,7 @@ import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../client/common/const
 import { IDisposable } from '../../../client/common/types';
 import { noop, swallowExceptions } from '../../../client/common/utils/misc';
 import { NotebookContentProvider } from '../../../client/datascience/notebook/contentProvider';
-import { ICell, INotebookEditorProvider } from '../../../client/datascience/types';
+import { ICell, INotebookEditorProvider, INotebookProvider } from '../../../client/datascience/types';
 import { waitForCondition } from '../../common';
 import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../../constants';
 import { closeActiveWindows, initialize } from '../../initialize';
@@ -162,11 +162,17 @@ export async function swallowSavingOfNotebooks() {
     sinon.stub(contentProvider, 'saveNotebookAs').callsFake(noop as any);
 }
 
+export async function shutdownAllNotebooks() {
+    const api = await initialize();
+    const notebookProvider = api.serviceContainer.get<INotebookProvider>(INotebookProvider);
+    await Promise.all(notebookProvider.activeNotebooks.map(async (item) => (await item).dispose()));
+}
 export async function closeNotebooksAndCleanUpAfterTests(disposables: IDisposable[] = []) {
     // We cannot close notebooks if there are any uncommitted changes (UI could hang with prompts etc).
     await commands.executeCommand('workbench.action.files.saveAll');
     await closeActiveWindows();
     disposeAllDisposables(disposables);
+    await shutdownAllNotebooks();
     sinon.restore();
 }
 

src/test/datascience/notebook/interrupRestart.ds.test.ts

@@ -19,6 +19,7 @@ import {
     canRunTests,
     closeNotebooksAndCleanUpAfterTests,
     deleteAllCellsAndWait,
+    disposeAllDisposables,
     insertPythonCellAndWait,
     startJupyter,
     swallowSavingOfNotebooks
@@ -39,6 +40,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
     let executionService: INotebookExecutionService;
     let vscEditor: VSCNotebookEditor;
     let vscodeNotebook: IVSCodeNotebook;
+    const suiteDisposables: IDisposable[] = [];
     suiteSetup(async function () {
         this.timeout(15_000);
         api = await initialize();
@@ -60,7 +62,8 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         vscEditor = vscodeNotebook.activeNotebookEditor!;
     });
     setup(deleteAllCellsAndWait);
-    suiteTeardown(() => closeNotebooksAndCleanUpAfterTests(disposables));
+    teardown(() => disposeAllDisposables(suiteDisposables));
+    suiteTeardown(() => closeNotebooksAndCleanUpAfterTests(disposables.concat(suiteDisposables)));
 
     test('Cancelling token will cancel cell execution (slow)', async () => {
         await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
@@ -122,6 +125,39 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
 
         await waitForCondition(async () => assertVSCCellIsIdle(cell), 1_000, 'Execution not cancelled');
     });
+    test('When running entire notebook, clicking VSCode Stop button should trigger a single interrupt, not one per cell', async () => {
+        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        const cell1 = vscEditor.document.cells[0];
+        const cell2 = vscEditor.document.cells[1];
+        const cell3 = vscEditor.document.cells[2];
+
+        const interrupt = sinon.spy(editorProvider.activeEditor!, 'interruptKernel');
+        suiteDisposables.push({ dispose: () => interrupt.restore() });
+
+        await commands.executeCommand('notebook.execute');
+
+        // Wait for cells to get busy.
+        await waitForCondition(
+            async () => assertVSCCellIsRunning(cell1) && assertVSCCellIsRunning(cell2) && assertVSCCellIsRunning(cell3),
+            15_000,
+            'Cells not being executed'
+        );
+
+        // Cancel execution.
+        await sleep(1_000);
+        await commands.executeCommand('notebook.cancelExecution');
+
+        // Wait for ?s, and verify cells are not running.
+        await waitForCondition(
+            async () => assertVSCCellIsIdle(cell1) && assertVSCCellIsIdle(cell2) && assertVSCCellIsIdle(cell3),
+            15_000,
+            'Cells are still running'
+        );
+
+        assert.equal(interrupt.callCount, 1, 'Interrupt should have been invoked only once');
+    });
     test('Restarting kernel will cancel cell execution (slow)', async () => {
         await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
         const cell = vscEditor.document.cells[0];
@@ -136,8 +172,21 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         assertVSCCellIsRunning(cell);
 
         // Restart the kernel.
-        await commands.executeCommand('python.datascience.notebookeditor.restartkernel');
+        const restartPromise = commands.executeCommand('python.datascience.notebookeditor.restartkernel');
 
         await waitForCondition(async () => assertVSCCellIsIdle(cell), 1_000, 'Execution not cancelled');
+
+        // Wait before we execute cells again.
+        await restartPromise;
+        await commands.executeCommand('notebook.execute');
+
+        // Wait for cell to get busy.
+        await waitForCondition(async () => assertVSCCellIsRunning(cell), 15_000, 'Cell not being executed');
+
+        // Cleanup (don't leave them running).
+        await commands.executeCommand('notebook.cancelExecution');
+
+        // Wait for ?s, and verify cells are not running.
+        await waitForCondition(async () => assertVSCCellIsIdle(cell), 15_000, 'Cell is still running');
     });
 });