Cancelling should invoke a single interrupt (microsoft#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');
  });
 });