Fix run all cells to wait for the previous cell (#10219)

* Fix run all cells to wait for the previous cell

* Add support for canceling in the middle and modify the run all test to verify

* Fix finishedPos on cancel.
Allow more than one execute chain to be in progress

* Add trace for execution

* Review feedback

* Sonar code smells

* Fix functional tests - storage was messing up next test

Author: rchiodo

news/2 Fixes/10016.md

@@ -0,0 +1 @@
+Fix run all cells to force each cell to finish before running the next one.

src/client/datascience/interactive-common/interactiveBase.ts

@@ -55,7 +55,7 @@ import {
     IGotoCode,
     IInteractiveWindowMapping,
     InteractiveWindowMessages,
-    IReExecuteCell,
+    IReExecuteCells,
     IRemoteAddCode,
     IRemoteReexecuteCode,
     IShowDataViewer,
@@ -230,8 +230,8 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
                 this.handleMessage(message, payload, this.submitNewCell);
                 break;
 
-            case InteractiveWindowMessages.ReExecuteCell:
-                this.handleMessage(message, payload, this.reexecuteCell);
+            case InteractiveWindowMessages.ReExecuteCells:
+                this.handleMessage(message, payload, this.reexecuteCells);
                 break;
 
             case InteractiveWindowMessages.DeleteAllCells:
@@ -462,8 +462,8 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
     // Submits a new cell to the window
     protected abstract submitNewCell(info: ISubmitNewCell): void;
 
-    // Re-executes a cell already in the window
-    protected reexecuteCell(_info: IReExecuteCell): void {
+    // Re-executes cells already in the window
+    protected reexecuteCells(_info: IReExecuteCells): void {
         // Default is not to do anything. This only works in the native editor
     }
 
@@ -499,7 +499,8 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
         line: number,
         id?: string,
         data?: nbformat.ICodeCell | nbformat.IRawCell | nbformat.IMarkdownCell,
-        debug?: boolean
+        debug?: boolean,
+        cancelToken?: CancellationToken
     ): Promise<boolean> {
         traceInfo(`Submitting code for ${this.id}`);
         const stopWatch =
@@ -567,7 +568,7 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
                         file,
                         line,
                         uuid(),
-                        undefined,
+                        cancelToken,
                         true
                     );
                 }
@@ -1437,7 +1438,7 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
                 this.postMessage(InteractiveWindowMessages.LoadTmLanguageResponse, s).ignoreErrors();
             })
             .catch(_e => {
-                this.postMessage(InteractiveWindowMessages.LoadTmLanguageResponse, undefined).ignoreErrors();
+                this.postMessage(InteractiveWindowMessages.LoadTmLanguageResponse).ignoreErrors();
             });
     }
 
@@ -1458,13 +1459,13 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
                     this.postMessage(InteractiveWindowMessages.LoadOnigasmAssemblyResponse, contents).ignoreErrors();
                 } else {
                     traceWarning('Onigasm file not found. Colorization will not be available.');
-                    this.postMessage(InteractiveWindowMessages.LoadOnigasmAssemblyResponse, undefined).ignoreErrors();
+                    this.postMessage(InteractiveWindowMessages.LoadOnigasmAssemblyResponse).ignoreErrors();
                 }
             }
         } else {
             // This happens during testing. Onigasm not needed as we're not testing colorization.
             traceWarning('File system not found. Colorization will not be available.');
-            this.postMessage(InteractiveWindowMessages.LoadOnigasmAssemblyResponse, undefined).ignoreErrors();
+            this.postMessage(InteractiveWindowMessages.LoadOnigasmAssemblyResponse).ignoreErrors();
         }
     }
 

src/client/datascience/interactive-common/interactiveWindowTypes.ts

@@ -76,7 +76,7 @@ export enum InteractiveWindowMessages {
     LoadAllCells = 'load_all_cells',
     LoadAllCellsComplete = 'load_all_cells_complete',
     ScrollToCell = 'scroll_to_cell',
-    ReExecuteCell = 'reexecute_cell',
+    ReExecuteCells = 'reexecute_cells',
     NotebookIdentity = 'identity',
     NotebookDirty = 'dirty',
     NotebookClean = 'clean',
@@ -177,9 +177,8 @@ export interface ISubmitNewCell {
     id: string;
 }
 
-export interface IReExecuteCell {
-    newCode: string;
-    cell: ICell;
+export interface IReExecuteCells {
+    entries: { cell: ICell; code: string }[];
 }
 
 export interface IProvideCompletionItemsRequest {
@@ -372,7 +371,7 @@ export class IInteractiveWindowMapping {
     public [InteractiveWindowMessages.LoadAllCells]: ILoadAllCells;
     public [InteractiveWindowMessages.LoadAllCellsComplete]: ILoadAllCells;
     public [InteractiveWindowMessages.ScrollToCell]: IScrollToCell;
-    public [InteractiveWindowMessages.ReExecuteCell]: IReExecuteCell;
+    public [InteractiveWindowMessages.ReExecuteCells]: IReExecuteCells;
     public [InteractiveWindowMessages.NotebookIdentity]: INotebookIdentity;
     public [InteractiveWindowMessages.NotebookDirty]: never | undefined;
     public [InteractiveWindowMessages.NotebookClean]: never | undefined;

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

@@ -8,7 +8,7 @@ import * as detectIndent from 'detect-indent';
 import { inject, injectable, multiInject, named } from 'inversify';
 import * as path from 'path';
 import * as uuid from 'uuid/v4';
-import { Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
+import { CancellationToken, CancellationTokenSource, Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
 
 import { concatMultilineStringInput, splitMultilineString } from '../../../datascience-ui/common';
 import { createCodeCell, createErrorOutput } from '../../../datascience-ui/common/cellFactory';
@@ -21,10 +21,11 @@ import {
     IWorkspaceService
 } from '../../common/application/types';
 import { ContextKey } from '../../common/contextKey';
-import { traceError } from '../../common/logger';
+import { traceError, traceInfo } from '../../common/logger';
 import { IFileSystem, TemporaryFile } from '../../common/platform/types';
 import {
     GLOBAL_MEMENTO,
+    IAsyncDisposableRegistry,
     IConfigurationService,
     ICryptoUtils,
     IDisposableRegistry,
@@ -55,7 +56,7 @@ import {
     IInsertCell,
     INativeCommand,
     InteractiveWindowMessages,
-    IReExecuteCell,
+    IReExecuteCells,
     IRemoveCell,
     ISaveAll,
     ISubmitNewCell,
@@ -166,6 +167,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     private indentAmount: string = ' ';
     private notebookJson: Partial<nbformat.INotebookContent> = {};
     private debouncedWriteToStorage = debounce(this.writeToStorage.bind(this), 250);
+    private executeCancelTokens = new Set<CancellationTokenSource>();
+    private _disposed = false;
 
     constructor(
         @multiInject(IInteractiveWindowListener) listeners: IInteractiveWindowListener[],
@@ -195,7 +198,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         @inject(ICryptoUtils) private crypto: ICryptoUtils,
         @inject(IExtensionContext) private context: IExtensionContext,
         @inject(ProgressReporter) progressReporter: ProgressReporter,
-        @inject(IExperimentsManager) experimentsManager: IExperimentsManager
+        @inject(IExperimentsManager) experimentsManager: IExperimentsManager,
+        @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry
     ) {
         super(
             progressReporter,
@@ -231,8 +235,10 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
             ViewColumn.Active,
             experimentsManager
         );
+        asyncRegistry.push(this);
     }
     public dispose(): Promise<void> {
+        this._disposed = true;
         super.dispose();
         return this.close();
     }
@@ -279,7 +285,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     public onMessage(message: string, payload: any) {
         super.onMessage(message, payload);
         switch (message) {
-            case InteractiveWindowMessages.ReExecuteCell:
+            case InteractiveWindowMessages.ReExecuteCells:
                 this.executedEvent.fire(this);
                 break;
 
@@ -324,6 +330,14 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
                 this.handleMessage(message, payload, this.clearAllOutputs);
                 break;
 
+            case InteractiveWindowMessages.RestartKernel:
+                this.interruptExecution();
+                break;
+
+            case InteractiveWindowMessages.Interrupt:
+                this.interruptExecution();
+                break;
+
             default:
                 break;
         }
@@ -403,11 +417,12 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         line: number,
         id?: string,
         data?: nbformat.ICodeCell | nbformat.IRawCell | nbformat.IMarkdownCell,
-        debug?: boolean
+        debug?: boolean,
+        cancelToken?: CancellationToken
     ): Promise<boolean> {
         const stopWatch = new StopWatch();
         const submitCodePromise = super
-            .submitCode(code, file, line, id, data, debug)
+            .submitCode(code, file, line, id, data, debug, cancelToken)
             .finally(() => this.sendPerceivedCellExecute(stopWatch));
         // When code is executed, update the version number in the metadata.
         return submitCodePromise.then(value => {
@@ -452,45 +467,35 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
 
     @captureTelemetry(Telemetry.ExecuteNativeCell, undefined, true)
     // tslint:disable-next-line:no-any
-    protected async reexecuteCell(info: IReExecuteCell): Promise<void> {
+    protected async reexecuteCells(info: IReExecuteCells): Promise<void> {
+        const tokenSource = new CancellationTokenSource();
+        this.executeCancelTokens.add(tokenSource);
+        let finishedPos = info && info.entries ? info.entries.length : -1;
         try {
-            // If there's any payload, it has the code and the id
-            if (info && info.newCode && info.cell.id && info.cell.data.cell_type !== 'messages') {
-                // Clear the result if we've run before
-                await this.clearResult(info.cell.id);
-
-                // Send to ourselves.
-                await this.submitCode(info.newCode, Identifiers.EmptyFileName, 0, info.cell.id, info.cell.data);
-
-                // Activate the other side, and send as if came from a file
-                await this.ipynbProvider.show(this.file);
-                this.shareMessage(InteractiveWindowMessages.RemoteReexecuteCode, {
-                    code: info.newCode,
-                    file: Identifiers.EmptyFileName,
-                    line: 0,
-                    id: info.cell.id,
-                    originator: this.id,
-                    debug: false
-                });
+            if (info && info.entries) {
+                for (let i = 0; i < info.entries.length && !tokenSource.token.isCancellationRequested; i += 1) {
+                    await this.reexecuteCell(info.entries[i], tokenSource.token);
+                    if (!tokenSource.token.isCancellationRequested) {
+                        finishedPos = i;
+                    }
+                }
             }
         } catch (exc) {
-            // Make this error our cell output
-            this.sendCellsToWebView([
-                {
-                    // tslint:disable-next-line: no-any
-                    data: { ...info.cell.data, outputs: [createErrorOutput(exc)] } as any, // nyc compiler issue
-                    id: info.cell.id,
-                    file: Identifiers.EmptyFileName,
-                    line: 0,
-                    state: CellState.error
-                }
-            ]);
-
             // Tell the other side we restarted the kernel. This will stop all executions
             this.postMessage(InteractiveWindowMessages.RestartKernel).ignoreErrors();
 
             // Handle an error
             await this.errorHandler.handleError(exc);
+        } finally {
+            this.executeCancelTokens.delete(tokenSource);
+
+            // Make sure everything is marked as finished or error after the final finished
+            // position
+            if (info && info.entries) {
+                for (let i = finishedPos + 1; i < info.entries.length; i += 1) {
+                    this.finishCell(info.entries[i]);
+                }
+            }
         }
     }
 
@@ -563,6 +568,75 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         // Actually don't close, just let the error bubble out
     }
 
+    private interruptExecution() {
+        this.executeCancelTokens.forEach(t => t.cancel());
+    }
+
+    private finishCell(entry: { cell: ICell; code: string }) {
+        this.sendCellsToWebView([
+            {
+                ...entry.cell,
+                // tslint:disable-next-line: no-any
+                data: { ...entry.cell.data, source: entry.code } as any, // nyc compiler issue
+                state: CellState.finished
+            }
+        ]);
+    }
+
+    private async reexecuteCell(entry: { cell: ICell; code: string }, cancelToken: CancellationToken): Promise<void> {
+        try {
+            // If there's any payload, it has the code and the id
+            if (entry.code && entry.cell.id && entry.cell.data.cell_type !== 'messages') {
+                traceInfo(`Executing cell ${entry.cell.id}`);
+
+                // Clear the result if we've run before
+                await this.clearResult(entry.cell.id);
+
+                // Send to ourselves.
+                await this.submitCode(
+                    entry.code,
+                    Identifiers.EmptyFileName,
+                    0,
+                    entry.cell.id,
+                    entry.cell.data,
+                    false,
+                    cancelToken
+                );
+
+                if (!cancelToken.isCancellationRequested) {
+                    // Activate the other side, and send as if came from a file
+                    await this.ipynbProvider.show(this.file);
+                    this.shareMessage(InteractiveWindowMessages.RemoteReexecuteCode, {
+                        code: entry.code,
+                        file: Identifiers.EmptyFileName,
+                        line: 0,
+                        id: entry.cell.id,
+                        originator: this.id,
+                        debug: false
+                    });
+                }
+            }
+        } catch (exc) {
+            // Make this error our cell output
+            this.sendCellsToWebView([
+                {
+                    // tslint:disable-next-line: no-any
+                    data: { ...entry.cell.data, outputs: [createErrorOutput(exc)] } as any, // nyc compiler issue
+                    id: entry.cell.id,
+                    file: Identifiers.EmptyFileName,
+                    line: 0,
+                    state: CellState.error
+                }
+            ]);
+
+            throw exc;
+        } finally {
+            if (entry && entry.cell.id) {
+                traceInfo(`Finished executing cell ${entry.cell.id}`);
+            }
+        }
+    }
+
     private sendPerceivedCellExecute(runningStopWatch?: StopWatch) {
         if (runningStopWatch) {
             const props = { notebook: true };
@@ -882,11 +956,13 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
 
     private async writeToStorage(filePath: string, contents?: string): Promise<void> {
         try {
-            if (contents) {
-                await this.fileSystem.createDirectory(path.dirname(filePath));
-                return this.fileSystem.writeFile(filePath, contents);
-            } else {
-                return this.fileSystem.deleteFile(filePath);
+            if (!this._disposed) {
+                if (contents) {
+                    await this.fileSystem.createDirectory(path.dirname(filePath));
+                    return this.fileSystem.writeFile(filePath, contents);
+                } else {
+                    return this.fileSystem.deleteFile(filePath);
+                }
             }
         } catch (exc) {
             traceError(`Error writing storage for ${filePath}: `, exc);

src/datascience-ui/interactive-common/redux/postOffice.ts

@@ -78,7 +78,7 @@ export enum IncomingMessageActions {
     LOADALLCELLS = 'action.load_all_cells',
     LOADALLCELLSCOMPLETE = 'action.load_all_cells_complete',
     SCROLLTOCELL = 'action.scroll_to_cell',
-    REEXECUTECELL = 'action.reexecute_cell',
+    REEXECUTECELLS = 'action.reexecute_cells',
     NOTEBOOKIDENTITY = 'action.identity',
     NOTEBOOKDIRTY = 'action.dirty',
     NOTEBOOKCLEAN = 'action.clean',