From c6c50c0ddfcbcf0282131b73841d00d3f4b9811a Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Fri, 26 Apr 2024 19:18:54 +1000 Subject: [PATCH] Replace proposed VSC API with our implementation (#15573) --- package.json | 1 - .../editor-integration/codeGenerator.ts | 13 +++- .../codeGenerator.unit.test.ts | 9 +-- .../editor-integration/codeLensFactory.ts | 25 +++--- .../codewatcher.unit.test.ts | 17 +--- .../editor-integration/hoverProvider.ts | 17 ++-- src/kernels/execution/cellExecution.ts | 23 ++---- src/kernels/execution/cellExecutionCreator.ts | 7 +- .../execution/cellExecutionMessageHandler.ts | 4 +- src/kernels/execution/cellExecutionQueue.ts | 18 ++--- src/kernels/execution/helpers.ts | 19 ++--- src/kernels/execution/types.ts | 9 +++ .../remoteJupyterServerMruUpdate.ts | 45 ++++++----- src/kernels/kernelAutoReConnectMonitor.ts | 77 +++++++++---------- .../kernelAutoReConnectMonitor.unit.test.ts | 36 +++------ src/kernels/kernelCrashMonitor.ts | 9 ++- src/kernels/kernelCrashMonitor.unit.test.ts | 11 ++- src/kernels/kernelExecution.ts | 29 +++---- src/kernels/kernelProvider.node.unit.test.ts | 17 +--- src/kernels/types.ts | 2 - .../debugger/kernelDebugAdapterBase.ts | 10 ++- .../jupyterCellOutputMimeTypeTracker.ts | 22 +++--- src/platform/errors/errors.ts | 2 +- .../notebooks/cellExecutionStateService.ts | 68 ++++++++++++++++ .../api/kernels/api.vscode.common.test.ts | 12 ++- src/standalone/executionAnalysis/symbols.ts | 10 +-- src/standalone/import-export/importTracker.ts | 22 ++---- .../import-export/importTracker.unit.test.ts | 50 ++++-------- .../survey/dataScienceSurveyBanner.node.ts | 26 +++---- .../interactiveWindow.vscode.common.test.ts | 1 - src/test/datascience/notebook/helper.ts | 15 +++- .../notebook/interruptRestart.vscode.test.ts | 3 +- .../variablesView/notebookWatcher.ts | 20 ++--- 33 files changed, 323 insertions(+), 326 deletions(-) create mode 100644 src/platform/notebooks/cellExecutionStateService.ts diff --git a/package.json b/package.json index 8be27f5b54d..81edbc85ea9 100644 --- a/package.json +++ b/package.json @@ -2068,7 +2068,6 @@ "notebookDeprecated", "notebookMessaging", "notebookMime", - "notebookCellExecutionState", "contribNotebookStaticPreloads", "portsAttributes", "quickPickSortByLabel", diff --git a/src/interactive-window/editor-integration/codeGenerator.ts b/src/interactive-window/editor-integration/codeGenerator.ts index 7849d9cb1e8..5b435ff0e3a 100644 --- a/src/interactive-window/editor-integration/codeGenerator.ts +++ b/src/interactive-window/editor-integration/codeGenerator.ts @@ -3,14 +3,12 @@ import { Disposable, - NotebookCellExecutionStateChangeEvent, NotebookDocument, Position, Range, TextDocumentChangeEvent, TextDocumentContentChangeEvent, Uri, - notebooks, workspace } from 'vscode'; @@ -21,6 +19,11 @@ import { uncommentMagicCommands } from './cellFactory'; import { CellMatcher } from './cellMatcher'; import { IGeneratedCode, IInteractiveWindowCodeGenerator, IGeneratedCodeStore, InteractiveCellMetadata } from './types'; import { computeHash } from '../../platform/common/crypto'; +import { + NotebookCellExecutionState, + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; // This class provides generated code for debugging jupyter cells. Call getGeneratedCode just before starting debugging to compute all of the // generated codes for cells & update the source maps in the python debugger. @@ -39,7 +42,7 @@ export class CodeGenerator implements IInteractiveWindowCodeGenerator { disposables.push(this); // Watch document changes so we can update our generated code workspace.onDidChangeTextDocument(this.onChangedDocument, this, this.disposables); - notebooks.onDidChangeNotebookCellExecutionState(this.onDidCellStateChange, this, this.disposables); + notebookCellExecutions.onDidChangeNotebookCellExecutionState(this.onDidCellStateChange, this, this.disposables); } public dispose() { @@ -107,7 +110,11 @@ export class CodeGenerator implements IInteractiveWindowCodeGenerator { } private onDidCellStateChange(e: NotebookCellExecutionStateChangeEvent) { + if (e.state !== NotebookCellExecutionState.Idle) { + return; + } if ( + e.state !== NotebookCellExecutionState.Idle || e.cell.notebook !== this.notebook || !e.cell.executionSummary?.executionOrder || this.cellIndexesCounted[e.cell.index] diff --git a/src/interactive-window/editor-integration/codeGenerator.unit.test.ts b/src/interactive-window/editor-integration/codeGenerator.unit.test.ts index 83141e8a34b..e235c71bc16 100644 --- a/src/interactive-window/editor-integration/codeGenerator.unit.test.ts +++ b/src/interactive-window/editor-integration/codeGenerator.unit.test.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. import { assert } from 'chai'; -import { EventEmitter, NotebookCellExecutionStateChangeEvent, NotebookDocument, Position, Range, Uri } from 'vscode'; +import { NotebookDocument, Position, Range, Uri } from 'vscode'; import { IConfigurationService, IDisposable, IWatchableJupyterSettings } from '../../platform/common/types'; import { CodeGenerator } from './codeGenerator'; @@ -10,7 +10,6 @@ import { MockDocumentManager } from '../../test/datascience/mockDocumentManager' import { anything, instance, mock, when } from 'ts-mockito'; import { IGeneratedCodeStore, InteractiveCellMetadata } from './types'; import { GeneratedCodeStorage } from './generatedCodeStorage'; -import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; import { dispose } from '../../platform/common/utils/lifecycle'; // eslint-disable-next-line @@ -21,7 +20,6 @@ suite.skip('Code Generator Unit Tests', () => { let pythonSettings: IWatchableJupyterSettings; let storage: IGeneratedCodeStore; let notebook: NotebookDocument; - let onDidChangeNotebookCellExecutionState: EventEmitter; let disposables: IDisposable[] = []; setup(() => { configurationService = mock(); @@ -30,11 +28,6 @@ suite.skip('Code Generator Unit Tests', () => { when(configurationService.getSettings(anything())).thenReturn(instance(pythonSettings)); documentManager = new MockDocumentManager(); notebook = mock(); - onDidChangeNotebookCellExecutionState = new EventEmitter(); - disposables.push(onDidChangeNotebookCellExecutionState); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - onDidChangeNotebookCellExecutionState.event - ); when(notebook.uri).thenReturn(); codeGenerator = new CodeGenerator(instance(configurationService), storage, instance(notebook), []); disposables.push(codeGenerator); diff --git a/src/interactive-window/editor-integration/codeLensFactory.ts b/src/interactive-window/editor-integration/codeLensFactory.ts index 8f1b521e575..cad430e6e89 100644 --- a/src/interactive-window/editor-integration/codeLensFactory.ts +++ b/src/interactive-window/editor-integration/codeLensFactory.ts @@ -2,17 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { - CodeLens, - Command, - Event, - EventEmitter, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, - Range, - TextDocument, - workspace -} from 'vscode'; +import { CodeLens, Command, Event, EventEmitter, Range, TextDocument, workspace } from 'vscode'; import { traceWarning, traceInfoIfCI, traceVerbose } from '../../platform/logging'; @@ -24,7 +14,11 @@ import { CodeLensCommands, Commands, InteractiveWindowView } from '../../platfor import { generateCellRangesFromDocument } from './cellFactory'; import { CodeLensPerfMeasures, ICodeLensFactory, IGeneratedCode, IGeneratedCodeStorageFactory } from './types'; import { StopWatch } from '../../platform/common/utils/stopWatch'; -import { notebooks } from 'vscode'; +import { + NotebookCellExecutionState, + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; type CodeLensCacheData = { cachedDocumentVersion: number | undefined; @@ -76,7 +70,11 @@ export class CodeLensFactory implements ICodeLensFactory { workspace.onDidCloseTextDocument(this.onClosedDocument, this, disposables); workspace.onDidGrantWorkspaceTrust(() => this.codeLensCache.clear(), this, disposables); this.configService.getSettings(undefined).onDidChange(this.onChangedSettings, this, disposables); - notebooks.onDidChangeNotebookCellExecutionState(this.onDidChangeNotebookCellExecutionState, this, disposables); + notebookCellExecutions.onDidChangeNotebookCellExecutionState( + this.onDidChangeNotebookCellExecutionState, + this, + disposables + ); kernelProvider.onDidDisposeKernel( (kernel) => { this.notebookData.delete(kernel.notebook.uri.toString()); @@ -387,7 +385,6 @@ export class CodeLensFactory implements ICodeLensFactory { range.start.character ]); } - break; case Commands.RunCellAndAllBelowPalette: case Commands.RunCellAndAllBelow: return this.generateCodeLens(range, Commands.RunCellAndAllBelow, runAllBelowTitle, [ diff --git a/src/interactive-window/editor-integration/codewatcher.unit.test.ts b/src/interactive-window/editor-integration/codewatcher.unit.test.ts index a3017516110..a41f15a92e1 100644 --- a/src/interactive-window/editor-integration/codewatcher.unit.test.ts +++ b/src/interactive-window/editor-integration/codewatcher.unit.test.ts @@ -5,17 +5,7 @@ // Disable whitespace / multiline as we use that to pass in our fake file strings import { expect } from 'chai'; import * as TypeMoq from 'typemoq'; -import { - CancellationTokenSource, - CodeLens, - Disposable, - EventEmitter, - NotebookCellExecutionStateChangeEvent, - Range, - Selection, - TextEditor, - Uri -} from 'vscode'; +import { CancellationTokenSource, CodeLens, Disposable, EventEmitter, Range, Selection, TextEditor, Uri } from 'vscode'; import { anything, instance, mock, when } from 'ts-mockito'; import { IDebugService } from '../../platform/common/application/types'; @@ -121,16 +111,11 @@ suite('Code Watcher Unit Tests', () => { when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); const trustedEvent = new EventEmitter(); when(mockedVSCodeNamespaces.workspace.onDidGrantWorkspaceTrust).thenReturn(trustedEvent.event); - const execStateChangeEvent = new EventEmitter(); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - execStateChangeEvent.event - ); const storageFactory = mock(); const kernelProvider = mock(); const kernelDisposedEvent = new EventEmitter(); when(kernelProvider.onDidDisposeKernel).thenReturn(kernelDisposedEvent.event); disposables.push(trustedEvent); - disposables.push(execStateChangeEvent); disposables.push(kernelDisposedEvent); const codeLensFactory = new CodeLensFactory( configService.object, diff --git a/src/interactive-window/editor-integration/hoverProvider.ts b/src/interactive-window/editor-integration/hoverProvider.ts index 2516758ca66..7e1c5335b85 100644 --- a/src/interactive-window/editor-integration/hoverProvider.ts +++ b/src/interactive-window/editor-integration/hoverProvider.ts @@ -16,6 +16,10 @@ import { IKernel, IKernelProvider } from '../../kernels/types'; import { IJupyterVariables } from '../../kernels/variables/types'; import { IInteractiveWindowProvider } from '../types'; import { getInteractiveCellMetadata } from '../helpers'; +import { + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; /** * Provides hover support in python files based on the state of a jupyter kernel. Files that are @@ -35,10 +39,11 @@ export class HoverProvider implements IExtensionSyncActivationService, vscode.Ho @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider ) {} public activate() { - this.onDidChangeNotebookCellExecutionStateHandler = vscode.notebooks.onDidChangeNotebookCellExecutionState( - (e) => this.delayer.trigger(() => this.onDidChangeNotebookCellExecutionState(e)), - this - ); + this.onDidChangeNotebookCellExecutionStateHandler = + notebookCellExecutions.onDidChangeNotebookCellExecutionState( + (e) => this.delayer.trigger(() => this.onDidChangeNotebookCellExecutionState(e)), + this + ); this.kernelProvider.onDidRestartKernel(() => this.runFiles.clear(), this, this.disposables); } public dispose() { @@ -48,9 +53,7 @@ export class HoverProvider implements IExtensionSyncActivationService, vscode.Ho this.hoverProviderRegistration.dispose(); } } - private async onDidChangeNotebookCellExecutionState( - e: vscode.NotebookCellExecutionStateChangeEvent - ): Promise { + private async onDidChangeNotebookCellExecutionState(e: NotebookCellExecutionStateChangeEvent): Promise { try { if (e.cell.notebook.notebookType !== InteractiveWindowView) { return; diff --git a/src/kernels/execution/cellExecution.ts b/src/kernels/execution/cellExecution.ts index 30b7eb3e766..61a8a83bab0 100644 --- a/src/kernels/execution/cellExecution.ts +++ b/src/kernels/execution/cellExecution.ts @@ -2,15 +2,7 @@ // Licensed under the MIT License. import type * as KernelMessage from '@jupyterlab/services/lib/kernel/messages'; -import { - NotebookCell, - NotebookCellExecution, - workspace, - NotebookCellOutput, - NotebookCellExecutionState, - Event, - EventEmitter -} from 'vscode'; +import { NotebookCell, NotebookCellExecution, workspace, NotebookCellOutput } from 'vscode'; import type { Kernel } from '@jupyterlab/services'; import { CellExecutionCreator } from './cellExecutionCreator'; @@ -39,6 +31,7 @@ import { ICellExecution } from './types'; import { KernelError } from '../errors/kernelError'; import { getCachedSysPrefix } from '../../platform/interpreter/helpers'; import { getCellMetadata } from '../../platform/common/utils'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; /** * Factory for CellExecution objects. @@ -74,9 +67,6 @@ export class CellExecution implements ICellExecution, IDisposable { public get result(): Promise { return this._result.promise; } - public get preExecute(): Event { - return this._preExecuteEmitter.event; - } private readonly _result = createDeferred(); private started?: boolean; @@ -94,10 +84,13 @@ export class CellExecution implements ICellExecution, IDisposable { private disposed?: boolean; private request: Kernel.IShellFuture | undefined; private readonly disposables: IDisposable[] = []; - private _preExecuteEmitter = new EventEmitter(); private cellExecutionHandler?: CellExecutionMessageHandler; private session?: IKernelSession; private cancelRequested?: boolean; + private _executionOrder?: number; + public get executionOrder() { + return this._executionOrder; + } private constructor( public readonly cell: NotebookCell, private readonly codeOverride: string | undefined, @@ -129,7 +122,6 @@ export class CellExecution implements ICellExecution, IDisposable { this, this.disposables ); - NotebookCellStateTracker.setCellState(cell, NotebookCellExecutionState.Idle); if (this.canExecuteCell()) { this.execution = CellExecutionCreator.getOrCreate( cell, @@ -378,6 +370,7 @@ export class CellExecution implements ICellExecution, IDisposable { if (activeNotebookCellExecution.get(this.cell.notebook) === this.execution) { activeNotebookCellExecution.set(this.cell.notebook, undefined); } + this._executionOrder = this.execution?.executionOrder; NotebookCellStateTracker.setCellState(this.cell, NotebookCellExecutionState.Idle); this.execution = undefined; } @@ -423,7 +416,7 @@ export class CellExecution implements ICellExecution, IDisposable { const kernelConnection = session.kernel; try { // At this point we're about to ACTUALLY execute some code. Fire an event to indicate that - this._preExecuteEmitter.fire(this.cell); + notebookCellExecutions.changeCellState(this.cell, NotebookCellExecutionState.Executing); traceVerbose(`Cell Index:${this.cell.index} sent to kernel`); // For Jupyter requests, silent === don't output, while store_history === don't update execution count // https://jupyter-client.readthedocs.io/en/stable/api/client.html#jupyter_client.KernelClient.execute diff --git a/src/kernels/execution/cellExecutionCreator.ts b/src/kernels/execution/cellExecutionCreator.ts index e1126fbd3f5..078386a33eb 100644 --- a/src/kernels/execution/cellExecutionCreator.ts +++ b/src/kernels/execution/cellExecutionCreator.ts @@ -21,7 +21,10 @@ import { getNotebookTelemetryTracker } from '../telemetry/notebookTelemetry'; * - Do something when 'end' is called */ export class NotebookCellExecutionWrapper implements NotebookCellExecution { - public started: boolean = false; + public _started: boolean = false; + public get started() { + return this._started; + } private _startTime?: number; public errorInfo: CellExecutionError; /** @@ -56,7 +59,7 @@ export class NotebookCellExecutionWrapper implements NotebookCellExecution { start(startTime?: number): void { // Allow this to be called more than once (so we can switch out a kernel during running a cell) if (!this.started) { - this.started = true; + this._started = true; this._impl.start(startTime); this._startTime = startTime; // We clear the output as soon as we start, diff --git a/src/kernels/execution/cellExecutionMessageHandler.ts b/src/kernels/execution/cellExecutionMessageHandler.ts index b7d388ebd21..c9400223c7a 100644 --- a/src/kernels/execution/cellExecutionMessageHandler.ts +++ b/src/kernels/execution/cellExecutionMessageHandler.ts @@ -449,8 +449,8 @@ export class CellExecutionMessageHandler implements IDisposable { this.previousResultsToRestore = { ...(this.cell.executionSummary || {}) }; this.temporaryExecution = CellExecutionCreator.getOrCreate(this.cell, this.controller); this.temporaryExecution?.start(); - if (this.previousResultsToRestore?.executionOrder && this.execution) { - this.execution.executionOrder = this.previousResultsToRestore.executionOrder; + if (this.previousResultsToRestore?.executionOrder && this.temporaryExecution) { + this.temporaryExecution.executionOrder = this.previousResultsToRestore.executionOrder; } return this.temporaryExecution; } diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index 3616527a206..22f66a0c101 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -12,6 +12,7 @@ import { ICellExecution, ICodeExecution } from './types'; import { CodeExecution } from './codeExecution'; import { once } from '../../platform/common/utils/events'; import { getCellMetadata } from '../../platform/common/utils'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; /** * A queue responsible for execution of cells. @@ -25,8 +26,6 @@ export class CellExecutionQueue implements Disposable { } private cancelledOrCompletedWithErrors = false; private startedRunningCells = false; - private readonly _onPreExecute = new EventEmitter(); - private readonly _onPostExecute = new EventEmitter(); private disposables: Disposable[] = []; private lastCellExecution?: ICellExecution | ICodeExecution; /** @@ -57,14 +56,6 @@ export class CellExecutionQueue implements Disposable { this.lastCellExecution?.dispose(); } - public get onPreExecute() { - return this._onPreExecute.event; - } - - public get onPostExecute() { - return this._onPostExecute.event; - } - /** * Queue the cell for execution & start processing it immediately. */ @@ -101,7 +92,6 @@ export class CellExecutionQueue implements Disposable { const cellExecution = this.executionFactory.create(cell, codeOverride, this.metadata); executionItem = cellExecution; this.disposables.push(cellExecution); - cellExecution.preExecute((c) => this._onPreExecute.fire(c), this, this.disposables); this.queueOfItemsToExecute.push(cellExecution); traceCellMessage(cell, 'User queued cell for execution'); @@ -245,7 +235,11 @@ export class CellExecutionQueue implements Disposable { } if (itemToExecute.type === 'cell') { - this._onPostExecute.fire(itemToExecute.cell); + notebookCellExecutions.changeCellState( + itemToExecute.cell, + NotebookCellExecutionState.Idle, + itemToExecute.executionOrder + ); } } diff --git a/src/kernels/execution/helpers.ts b/src/kernels/execution/helpers.ts index 60e56976ef0..76bab49efaa 100644 --- a/src/kernels/execution/helpers.ts +++ b/src/kernels/execution/helpers.ts @@ -2,14 +2,7 @@ // Licensed under the MIT License. import type * as nbformat from '@jupyterlab/nbformat'; -import { - NotebookCellOutput, - NotebookCellOutputItem, - NotebookCell, - NotebookCellExecutionState, - Position, - Range -} from 'vscode'; +import { NotebookCellOutput, NotebookCellOutputItem, NotebookCell, Position, Range } from 'vscode'; // eslint-disable-next-line @typescript-eslint/no-require-imports import type { KernelMessage } from '@jupyterlab/services'; import fastDeepEqual from 'fast-deep-equal'; @@ -29,9 +22,10 @@ import { getKernelRegistrationInfo } from '../helpers'; import { StopWatch } from '../../platform/common/utils/stopWatch'; -import { getExtensionSpecifcStack } from '../../platform/errors/errors'; +import { getExtensionSpecificStack } from '../../platform/errors/errors'; import { getCachedEnvironment, getVersion } from '../../platform/interpreter/helpers'; import { base64ToUint8Array, uint8ArrayToBase64 } from '../../platform/common/utils/string'; +import type { NotebookCellExecutionState } from '../../platform/notebooks/cellExecutionStateService'; export enum CellOutputMimeTypes { error = 'application/vnd.code.notebook.error', @@ -117,7 +111,7 @@ export function traceCellMessage(cell: NotebookCell, message: string | (() => st `Cell Index:${cell.index}, of document ${uriPath.basename( cell.notebook.uri )} with state:${NotebookCellStateTracker.getCellStatus(cell)}, exec: ${cell.executionSummary - ?.executionOrder}. ${messageToLog()}. called from ${getExtensionSpecifcStack()}` + ?.executionOrder}. ${messageToLog()}. called from ${getExtensionSpecificStack()}` ); } @@ -775,9 +769,10 @@ export async function endCellAndDisplayErrorsInCell( // Start execution if not already (Cell execution wrapper will ensure it won't start twice) const execution = CellExecutionCreator.getOrCreate(cell, controller); + const originalExecutionOrder = execution.executionOrder; if (!execution.started) { - execution.start(cell.executionSummary?.timing?.endTime); - execution.executionOrder = cell.executionSummary?.executionOrder; + execution.start(cell.executionSummary?.timing?.startTime); + execution.executionOrder = cell.executionSummary?.executionOrder || originalExecutionOrder; } await execution.appendOutput(output); execution.end(isCancelled ? undefined : false, cell.executionSummary?.timing?.endTime); diff --git a/src/kernels/execution/types.ts b/src/kernels/execution/types.ts index 3bd1731f653..2a8a763d1f1 100644 --- a/src/kernels/execution/types.ts +++ b/src/kernels/execution/types.ts @@ -10,6 +10,15 @@ export interface ICellExecution { type: 'cell'; cell: NotebookCell; result: Promise; + /** + * Execution count for the cell after it completed execution. + * Its possible that cell.executionSummary.executionOrder is undefined when this is set. + * Thats because the data has not yet made its way to VS Code and back into the extension host model API. + * + * This gives us access to the execution count as soon as its available. + * Without having to wait for the roundtrip to complete. + */ + executionOrder?: number; start(session: IKernelSession): Promise; cancel(forced?: boolean): Promise; dispose(): void; diff --git a/src/kernels/jupyter/connection/remoteJupyterServerMruUpdate.ts b/src/kernels/jupyter/connection/remoteJupyterServerMruUpdate.ts index 778dbe7bb33..6c96e34981d 100644 --- a/src/kernels/jupyter/connection/remoteJupyterServerMruUpdate.ts +++ b/src/kernels/jupyter/connection/remoteJupyterServerMruUpdate.ts @@ -9,13 +9,18 @@ import { IDisposable, IDisposableRegistry } from '../../../platform/common/types import { noop } from '../../../platform/common/utils/misc'; import { IJupyterServerUriStorage } from '../types'; import { IKernel, IKernelProvider, isRemoteConnection } from '../../types'; +import { + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../../platform/notebooks/cellExecutionStateService'; +import { Delayer } from '../../../platform/common/utils/async'; const INTERVAL_IN_SECONDS_TO_UPDATE_MRU = 1_000; @injectable() export class RemoteJupyterServerMruUpdate implements IExtensionSyncActivationService { private readonly disposables: IDisposable[] = []; - private readonly monitoredKernels = new WeakSet(); private readonly timeouts = new Set(); + private readonly kernelSpecificUpdates = new WeakMap>(); constructor( @inject(IJupyterServerUriStorage) private readonly serverStorage: IJupyterServerUriStorage, @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, @@ -28,34 +33,34 @@ export class RemoteJupyterServerMruUpdate implements IExtensionSyncActivationSer dispose(Array.from(this.timeouts.values())); } activate(): void { - this.kernelProvider.onDidStartKernel(this.onDidStartKernel, this, this.disposables); - this.kernelProvider.onDidRestartKernel(this.onDidStartKernel, this, this.disposables); + this.disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState( + this.onDidChangeNotebookCellExecutionState, + this + ) + ); } - private onDidStartKernel(kernel: IKernel) { + private onDidChangeNotebookCellExecutionState(e: NotebookCellExecutionStateChangeEvent) { + const kernel = this.kernelProvider.get(e.cell.notebook); + if (!kernel) { + return; + } const connection = kernel.kernelConnectionMetadata; - if (!isRemoteConnection(connection) || this.monitoredKernels.has(kernel)) { + if (!isRemoteConnection(connection)) { return; } - this.monitoredKernels.add(kernel); + const delayer = this.kernelSpecificUpdates.get(kernel) || new Delayer(INTERVAL_IN_SECONDS_TO_UPDATE_MRU); + this.kernelSpecificUpdates.set(kernel, delayer); // We do not want 100s of 1000s of these timeouts, // multiply by notebooks, and multiply by number of kernels, this grows unnecessarily. - const disposables: IDisposable[] = []; - this.disposables.push(new Disposable(() => dispose(disposables))); - - const updateConnectionTime = () => { - dispose(disposables); + void delayer.trigger(() => { + // Log this remote URI into our MRU list if (kernel.disposed || kernel.disposing) { return; } - const timeout = setTimeout(() => { - // Log this remote URI into our MRU list - this.serverStorage.update(connection.serverProviderHandle).catch(noop); - }, INTERVAL_IN_SECONDS_TO_UPDATE_MRU); - disposables.push(new Disposable(() => clearTimeout(timeout))); - }; - this.kernelProvider.getKernelExecution(kernel).onPreExecute(updateConnectionTime, this, this.disposables); - // Log this remote URI into our MRU list - updateConnectionTime(); + // Log this remote URI into our MRU list + this.serverStorage.update(connection.serverProviderHandle).catch(noop); + }); } } diff --git a/src/kernels/kernelAutoReConnectMonitor.ts b/src/kernels/kernelAutoReConnectMonitor.ts index 79ecb10c55d..06a3bbfa8ee 100644 --- a/src/kernels/kernelAutoReConnectMonitor.ts +++ b/src/kernels/kernelAutoReConnectMonitor.ts @@ -3,15 +3,7 @@ import type { Kernel } from '@jupyterlab/services'; import { inject, injectable } from 'inversify'; -import { - CancellationTokenSource, - Disposable, - NotebookCell, - NotebookCellExecutionState, - notebooks, - ProgressLocation, - window -} from 'vscode'; +import { CancellationTokenSource, Disposable, NotebookCell, ProgressLocation, window } from 'vscode'; import { IExtensionSyncActivationService } from '../platform/activation/types'; import { IDisposable, IDisposableRegistry } from '../platform/common/types'; import { createDeferred } from '../platform/common/utils/async'; @@ -30,6 +22,8 @@ import { RemoteKernelConnectionMetadata } from './types'; import { IJupyterServerProviderRegistry, IJupyterServerUriStorage } from './jupyter/types'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; +import { DisposableStore } from '../platform/common/utils/lifecycle'; /** * In the case of Jupyter kernels, when a kernel dies Jupyter will automatically restart that kernel. @@ -38,6 +32,7 @@ import { IJupyterServerProviderRegistry, IJupyterServerUriStorage } from './jupy @injectable() export class KernelAutoReconnectMonitor implements IExtensionSyncActivationService { private kernelsStartedSuccessfully = new WeakSet(); + private kernelDisposables = new WeakMap(); private kernelConnectionToKernelMapping = new WeakMap(); private kernelsRestarting = new WeakSet(); private kernelReconnectProgress = new WeakMap(); @@ -54,6 +49,8 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi this.kernelProvider.onDidStartKernel(this.onDidStartKernel, this, this.disposableRegistry); this.disposableRegistry.push( this.kernelProvider.onDidDisposeKernel((kernel) => { + this.kernelDisposables.get(kernel)?.dispose(); + this.kernelDisposables.delete(kernel); this.kernelReconnectProgress.get(kernel)?.dispose(); this.kernelReconnectProgress.delete(kernel); }, this) @@ -65,7 +62,7 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi }, this) ); this.disposableRegistry.push( - notebooks.onDidChangeNotebookCellExecutionState((e) => { + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { if (!isJupyterNotebook(e.cell.notebook)) { return; } @@ -81,30 +78,33 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi }) ); } + private getKernelSpecificDisposables(kernel: IKernel) { + const kernelDisposables = this.kernelDisposables.get(kernel) || new DisposableStore(); + this.disposableRegistry.push(kernelDisposables); + this.kernelDisposables.set(kernel, kernelDisposables); + return kernelDisposables; + } private onDidStartKernel(kernel: IKernel) { - if (!this.kernelsStartedSuccessfully.has(kernel)) { - this.kernelProvider - .getKernelExecution(kernel) - .onPreExecute( - (cell) => this.lastExecutedCellPerKernel.set(kernel, cell), - this, - this.disposableRegistry - ); + if (this.kernelsStartedSuccessfully.has(kernel) || !kernel.session?.kernel) { + return; + } - if (!kernel.session?.kernel) { - return; - } - this.kernelsStartedSuccessfully.add(kernel); - this.kernelConnectionToKernelMapping.set(kernel.session.kernel, kernel); - kernel.session?.kernel?.connectionStatusChanged.connect(this.onKernelStatusChanged, this); - kernel.onDisposed( - () => { - this.kernelReconnectProgress.get(kernel)?.dispose(); - this.kernelReconnectProgress.delete(kernel); - }, - this, - this.disposableRegistry - ); + const kernelDisposables = this.getKernelSpecificDisposables(kernel); + kernelDisposables.add( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if ( + e.cell.notebook === kernel.notebook && + this.kernelProvider.get(e.cell.notebook) === kernel && + e.state === NotebookCellExecutionState.Executing + ) { + this.lastExecutedCellPerKernel.set(kernel, e.cell); + } + }, this) + ); + this.kernelsStartedSuccessfully.add(kernel); + this.kernelConnectionToKernelMapping.set(kernel.session.kernel, kernel); + kernel.session?.kernel?.connectionStatusChanged.connect(this.onKernelStatusChanged, this); + kernelDisposables.add( kernel.addHook( 'willRestart', async () => { @@ -112,11 +112,10 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi this.kernelReconnectProgress.delete(kernel); this.kernelsRestarting.add(kernel); }, - this, - this.disposableRegistry - ); - kernel.onRestarted(() => this.kernelsRestarting.delete(kernel), this, this.disposableRegistry); - } + this + ) + ); + kernelDisposables.add(kernel.onRestarted(() => this.kernelsRestarting.delete(kernel), this)); } private onKernelStatusChanged(connection: Kernel.IKernelConnection, connectionStatus: Kernel.ConnectionStatus) { const kernel = this.kernelConnectionToKernelMapping.get(connection); @@ -150,13 +149,13 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi } } private async onKernelConnecting(kernel: IKernel) { + const kernelDisposables = this.getKernelSpecificDisposables(kernel); const deferred = createDeferred(); const disposable = new Disposable(() => deferred.resolve()); this.kernelReconnectProgress.set(kernel, disposable); if (isRemoteConnection(kernel.kernelConnectionMetadata)) { const handled = await this.handleRemoteServerShutdown(kernel, kernel.kernelConnectionMetadata); - if (handled) { return; } @@ -169,7 +168,7 @@ export class KernelAutoReconnectMonitor implements IExtensionSyncActivationServi .withProgress({ location: ProgressLocation.Notification, title: message }, async () => deferred.promise) .then(noop, noop); - this.disposableRegistry.push(disposable); + kernelDisposables.add(disposable); } private async onKernelDisconnected(kernel: IKernel) { diff --git a/src/kernels/kernelAutoReConnectMonitor.unit.test.ts b/src/kernels/kernelAutoReConnectMonitor.unit.test.ts index 5c640debff9..01f1a8ef4ab 100644 --- a/src/kernels/kernelAutoReConnectMonitor.unit.test.ts +++ b/src/kernels/kernelAutoReConnectMonitor.unit.test.ts @@ -14,16 +14,7 @@ import { RemoteKernelSpecConnectionMetadata } from './types'; import { anything, instance, mock, verify, when } from 'ts-mockito'; -import { - Disposable, - EventEmitter, - NotebookCell, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, - NotebookDocument, - TextDocument, - Uri -} from 'vscode'; +import { Disposable, EventEmitter, NotebookCell, NotebookDocument, TextDocument, Uri } from 'vscode'; import { Signal } from '@lumino/signaling'; import type { Kernel } from '@jupyterlab/services'; import { KernelAutoReconnectMonitor } from './kernelAutoReConnectMonitor'; @@ -33,6 +24,7 @@ import { JupyterNotebookView } from '../platform/common/constants'; import { IJupyterServerProviderRegistry, IJupyterServerUriEntry, IJupyterServerUriStorage } from './jupyter/types'; import { noop } from '../test/core'; import { JupyterServer, JupyterServerCollection, JupyterServerProvider } from '../api'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; suite('Kernel ReConnect Progress Message', () => { let disposables: IDisposable[] = []; @@ -78,8 +70,6 @@ suite('Kernel ReConnect Progress Message', () => { function createKernel() { const kernel = mock(); const onRestarted = new EventEmitter(); - const onPreExecute = new EventEmitter(); - disposables.push(onPreExecute); disposables.push(onRestarted); const session = mock(); const kernelConnection = mock(); @@ -97,7 +87,6 @@ suite('Kernel ReConnect Progress Message', () => { when(kernel.resourceUri).thenReturn(Uri.file('test.ipynb')); when(session.kernel).thenReturn(instance(kernelConnection)); when(kernel.kernelConnectionMetadata).thenReturn(connectionMetadata); - when(kernelExecution.onPreExecute).thenReturn(onPreExecute.event); when(kernel.onRestarted).thenReturn(onRestarted.event); when(kernel.dispose()).thenResolve(); let onWillRestart: (e: 'willRestart') => Promise = () => Promise.resolve(); @@ -151,7 +140,6 @@ suite('Kernel ReConnect Failed Monitor', () => { let onDidRestartKernel: EventEmitter; let clock: fakeTimers.InstalledClock; let cellExecution: NotebookCellExecutionWrapper; - let onDidChangeNotebookCellExecutionState: EventEmitter; let kernelExecution: INotebookKernelExecution; setup(() => { resetVSCodeMocks(); @@ -186,19 +174,12 @@ suite('Kernel ReConnect Failed Monitor', () => { const stub = sinon.stub(CellExecutionCreator, 'getOrCreate').callsFake(() => instance(cellExecution)); disposables.push(new Disposable(() => stub.restore())); disposables.push(new Disposable(() => clock.uninstall())); - onDidChangeNotebookCellExecutionState = new EventEmitter(); - disposables.push(onDidChangeNotebookCellExecutionState); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - onDidChangeNotebookCellExecutionState.event - ); monitor.activate(); }); teardown(() => (disposables = dispose(disposables))); function createKernel(serverProviderHandle = { handle: '1234', id: '1234', extensionId: '' }) { const kernel = mock(); - const onPreExecute = new EventEmitter(); const onRestarted = new EventEmitter(); - disposables.push(onPreExecute); disposables.push(onRestarted); const session = mock(); const kernelConnection = mock(); @@ -218,11 +199,12 @@ suite('Kernel ReConnect Failed Monitor', () => { when(kernel.resourceUri).thenReturn(Uri.file('test.ipynb')); when(session.kernel).thenReturn(instance(kernelConnection)); when(kernel.kernelConnectionMetadata).thenReturn(connectionMetadata); - when(kernelExecution.onPreExecute).thenReturn(onPreExecute.event); when(kernel.onRestarted).thenReturn(onRestarted.event); when(kernel.dispose()).thenResolve(); + when(kernel.notebook).thenReturn(); + when(kernelProvider.get(anything())).thenReturn(instance(kernel)); - return { kernel, onPreExecute, onRestarted, kernelConnectionStatusSignal }; + return { kernel, onRestarted, kernelConnectionStatusSignal }; } function createNotebook() { const nb = mock(); @@ -284,10 +266,11 @@ suite('Kernel ReConnect Failed Monitor', () => { const kernel = createKernel(); const nb = createNotebook(); + when(kernel.kernel.notebook).thenReturn(instance(nb)); const cell = createCell(instance(nb)); when(kernelProvider.get(instance(nb))).thenReturn(instance(kernel.kernel)); onDidStartKernel.fire(instance(kernel.kernel)); - kernel.onPreExecute.fire(instance(cell)); + notebookCellExecutions.changeCellState(instance(cell), NotebookCellExecutionState.Executing); // Send the kernel into connecting state & then disconnected. kernel.kernelConnectionStatusSignal.emit('connecting'); @@ -301,16 +284,17 @@ suite('Kernel ReConnect Failed Monitor', () => { const kernel = createKernel(); const nb = createNotebook(); + when(kernel.kernel.notebook).thenReturn(instance(nb)); const cell = createCell(instance(nb)); when(kernelProvider.get(instance(nb))).thenReturn(instance(kernel.kernel)); onDidStartKernel.fire(instance(kernel.kernel)); - kernel.onPreExecute.fire(instance(cell)); + notebookCellExecutions.changeCellState(instance(cell), NotebookCellExecutionState.Executing); // Send the kernel into connecting state & then disconnected. kernel.kernelConnectionStatusSignal.emit('connecting'); // Mark the cell as completed. - onDidChangeNotebookCellExecutionState.fire({ cell: instance(cell), state: NotebookCellExecutionState.Idle }); + notebookCellExecutions.changeCellState(instance(cell), NotebookCellExecutionState.Idle); kernel.kernelConnectionStatusSignal.emit('disconnected'); await clock.runAllAsync(); diff --git a/src/kernels/kernelCrashMonitor.ts b/src/kernels/kernelCrashMonitor.ts index 4be9c92ac57..4c596ddec15 100644 --- a/src/kernels/kernelCrashMonitor.ts +++ b/src/kernels/kernelCrashMonitor.ts @@ -14,6 +14,7 @@ import { endCellAndDisplayErrorsInCell } from './execution/helpers'; import { getDisplayNameOrNameOfKernelConnection } from './helpers'; import { IKernel, IKernelProvider } from './types'; import { swallowExceptions } from '../platform/common/utils/decorators'; +import { notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; /** * Monitors kernel crashes and on the event of a crash will display the results in the most recent cell. @@ -33,9 +34,11 @@ export class KernelCrashMonitor implements IExtensionSyncActivationService { } private onDidStartKernel(kernel: IKernel) { this.kernelsStartedSuccessfully.add(kernel); - this.kernelProvider - .getKernelExecution(kernel) - .onPreExecute((cell) => this.lastExecutedCellPerKernel.set(kernel, cell), this, this.disposableRegistry); + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell.notebook === kernel.notebook && this.kernelProvider.get(e.cell.notebook) === kernel) { + this.lastExecutedCellPerKernel.set(kernel, e.cell); + } + }); } @swallowExceptions() diff --git a/src/kernels/kernelCrashMonitor.unit.test.ts b/src/kernels/kernelCrashMonitor.unit.test.ts index 9760caf52e4..8becf0b7834 100644 --- a/src/kernels/kernelCrashMonitor.unit.test.ts +++ b/src/kernels/kernelCrashMonitor.unit.test.ts @@ -23,6 +23,7 @@ import { DataScience } from '../platform/common/utils/localize'; import { createOutputWithErrorMessageForDisplay } from '../platform/errors/errorUtils'; import { getDisplayNameOrNameOfKernelConnection } from './helpers'; import { mockedVSCodeNamespaces } from '../test/vscode-mock'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; suite('Kernel Crash Monitor', () => { let kernelProvider: IKernelProvider; @@ -35,7 +36,6 @@ suite('Kernel Crash Monitor', () => { }>; let onDidStartKernel: EventEmitter; let kernelExecution: INotebookKernelExecution; - let onPreExecute: EventEmitter; let cell: NotebookCell; let kernelSession: IKernelSession; let notebook: TestNotebookDocument; @@ -72,15 +72,14 @@ suite('Kernel Crash Monitor', () => { kernel: IKernel; }>(); onDidStartKernel = new EventEmitter(); - onPreExecute = new EventEmitter(); notebook = new TestNotebookDocument(); cell = await notebook.appendCodeCell('1234'); controller = createKernelController('1'); disposables.push(onDidStartKernel); disposables.push(onKernelStatusChanged); - disposables.push(onPreExecute); when(kernel.dispose()).thenResolve(); when(kernel.disposed).thenReturn(false); + when(kernel.notebook).thenReturn(notebook); when(kernel.controller).thenReturn(controller); when(kernel.disposing).thenReturn(false); when(kernel.session).thenReturn(instance(kernelSession)); @@ -91,7 +90,7 @@ suite('Kernel Crash Monitor', () => { when(kernelProvider.onKernelStatusChanged).thenReturn(onKernelStatusChanged.event); when(kernelProvider.getOrCreate(anything(), anything())).thenReturn(instance(kernel)); when(kernelProvider.getKernelExecution(anything())).thenReturn(instance(kernelExecution)); - when(kernelExecution.onPreExecute).thenReturn(onPreExecute.event); + when(kernelProvider.get(anything())).thenReturn(instance(kernel)); kernelCrashMonitor = new KernelCrashMonitor(disposables, instance(kernelProvider)); clock = fakeTimers.install(); @@ -105,7 +104,7 @@ suite('Kernel Crash Monitor', () => { // Ensure we have a kernel and have started a cell. kernelCrashMonitor.activate(); onDidStartKernel.fire(instance(kernel)); - onPreExecute.fire(cell); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); const execution = controller.createNotebookCellExecution(cell); execution.start(); @@ -135,7 +134,7 @@ suite('Kernel Crash Monitor', () => { // Ensure we have a kernel and have started a cell. kernelCrashMonitor.activate(); onDidStartKernel.fire(instance(kernel)); - onPreExecute.fire(cell); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); const execution = controller.createNotebookCellExecution(cell); execution.start(); diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index f1c922a39fb..2104eecb7a9 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -6,8 +6,6 @@ import { NotebookCell, NotebookCellKind, EventEmitter, - notebooks, - NotebookCellExecutionState, NotebookDocument, workspace, CancellationToken, @@ -47,6 +45,7 @@ import { } from './execution/extensionDisplayDataTracker'; import { CodeExecution } from './execution/codeExecution'; import type { ICodeExecution } from './execution/types'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; /** * Everything in this classes gets disposed via the `onWillCancel` hook. @@ -57,10 +56,6 @@ export class NotebookKernelExecution implements INotebookKernelExecution { return this._visibleExecutionCount; } private _visibleExecutionCount = 0; - private readonly _onPreExecute = new EventEmitter(); - public readonly onPreExecute = this._onPreExecute.event; - private readonly _onPostExecute = new EventEmitter(); - public readonly onPostExecute = this._onPostExecute.event; private readonly documentExecutions = new WeakMap(); private readonly executionFactory: CellExecutionFactory; private readonly _onDidReceiveDisplayUpdate = new EventEmitter(); @@ -81,15 +76,16 @@ export class NotebookKernelExecution implements INotebookKernelExecution { ); this.disposables.push(requestListener); this.executionFactory = new CellExecutionFactory(kernel.controller, requestListener); - - notebooks.onDidChangeNotebookCellExecutionState((e) => { - if (e.cell.notebook === kernel.notebook) { - if (e.state === NotebookCellExecutionState.Idle && e.cell.executionSummary?.executionOrder) { - this._visibleExecutionCount = Math.max( - this._visibleExecutionCount, - e.cell.executionSummary.executionOrder - ); - } + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if ( + e.cell.notebook === kernel.notebook && + e.state === NotebookCellExecutionState.Idle && + e.cell.executionSummary?.executionOrder + ) { + this._visibleExecutionCount = Math.max( + this._visibleExecutionCount, + e.cell.executionSummary.executionOrder + ); } }); kernel.onRestarted(() => (this._visibleExecutionCount = 0), this, this.disposables); @@ -100,7 +96,6 @@ export class NotebookKernelExecution implements INotebookKernelExecution { kernel.onStatusChanged(this.hookupIOPubHandler, this, this.disposables); kernel.onRestarted(this.hookupIOPubHandler, this, this.disposables); this.hookupIOPubHandler(); - this.disposables.push(this._onPreExecute); } private hookupIOPubHandler() { const session = this.kernel.session; @@ -372,8 +367,6 @@ export class NotebookKernelExecution implements INotebookKernelExecution { this, this.disposables ); - newCellExecutionQueue.onPreExecute((c) => this._onPreExecute.fire(c), this, this.disposables); - newCellExecutionQueue.onPostExecute((c) => this._onPostExecute.fire(c), this, this.disposables); this.documentExecutions.set(document, newCellExecutionQueue); return newCellExecutionQueue; } diff --git a/src/kernels/kernelProvider.node.unit.test.ts b/src/kernels/kernelProvider.node.unit.test.ts index 81cab514ec4..971acb15716 100644 --- a/src/kernels/kernelProvider.node.unit.test.ts +++ b/src/kernels/kernelProvider.node.unit.test.ts @@ -5,14 +5,7 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; -import { - EventEmitter, - Memento, - NotebookCellExecutionStateChangeEvent, - NotebookController, - NotebookDocument, - Uri -} from 'vscode'; +import { EventEmitter, Memento, NotebookController, NotebookDocument, Uri } from 'vscode'; import { IConfigurationService, IDisposable, @@ -197,17 +190,9 @@ suite('Jupyter Session', () => { jupyterServerUriStorage = mock(); context = mock(); const configSettings = mock(); - const onDidChangeNotebookCellExecutionState = new EventEmitter(); - disposables.push(onDidChangeNotebookCellExecutionState); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - onDidChangeNotebookCellExecutionState.event - ); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( onDidCloseNotebookDocument.event ); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - onDidChangeNotebookCellExecutionState.event - ); when(configService.getSettings(anything())).thenReturn(instance(configSettings)); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([ instance(sampleNotebook1), diff --git a/src/kernels/types.ts b/src/kernels/types.ts index db270c23d39..59b040546e5 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -432,8 +432,6 @@ export interface INotebookKernelExecution { * Total execution count on this kernel */ readonly executionCount: number; - readonly onPreExecute: Event; - readonly onPostExecute: Event; readonly onDidReceiveDisplayUpdate: Event; /** * Cells that are still being executed (or pending). diff --git a/src/notebooks/debugger/kernelDebugAdapterBase.ts b/src/notebooks/debugger/kernelDebugAdapterBase.ts index 8afefb0e1c9..d48ca3cb2ad 100644 --- a/src/notebooks/debugger/kernelDebugAdapterBase.ts +++ b/src/notebooks/debugger/kernelDebugAdapterBase.ts @@ -11,11 +11,8 @@ import { Event, EventEmitter, NotebookCell, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, NotebookCellKind, NotebookDocument, - notebooks, Uri, workspace } from 'vscode'; @@ -40,6 +37,11 @@ import { shortNameMatchesLongName } from './helper'; import { SessionDisposedError } from '../../platform/errors/sessionDisposedError'; +import { + notebookCellExecutions, + NotebookCellExecutionState, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; /** * For info on the custom requests implemented by jupyter see: @@ -97,7 +99,7 @@ export abstract class KernelDebugAdapterBase implements DebugAdapter, IKernelDeb } this.disposables.push( - notebooks.onDidChangeNotebookCellExecutionState( + notebookCellExecutions.onDidChangeNotebookCellExecutionState( (cellStateChange: NotebookCellExecutionStateChangeEvent) => { // If a cell has moved to idle, stop the debug session if ( diff --git a/src/notebooks/outputs/jupyterCellOutputMimeTypeTracker.ts b/src/notebooks/outputs/jupyterCellOutputMimeTypeTracker.ts index 5af3391aa3f..4953db5d4e8 100644 --- a/src/notebooks/outputs/jupyterCellOutputMimeTypeTracker.ts +++ b/src/notebooks/outputs/jupyterCellOutputMimeTypeTracker.ts @@ -2,14 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { - NotebookCell, - NotebookCellExecutionStateChangeEvent, - NotebookCellKind, - NotebookDocument, - notebooks, - workspace -} from 'vscode'; +import { NotebookCell, NotebookCellKind, NotebookDocument, workspace } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { JupyterNotebookView } from '../../platform/common/constants'; import { dispose } from '../../platform/common/utils/lifecycle'; @@ -22,6 +15,11 @@ import { Telemetry } from '../../telemetry'; import { isTelemetryDisabled } from '../../telemetry'; +import { + NotebookCellExecutionState, + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; /** * Sends telemetry about cell output mime types @@ -40,7 +38,7 @@ export class CellOutputMimeTypeTracker implements IExtensionSyncActivationServic workspace.onDidOpenNotebookDocument(this.onDidOpenCloseDocument, this, this.disposables); workspace.onDidCloseNotebookDocument(this.onDidOpenCloseDocument, this, this.disposables); workspace.onDidSaveNotebookDocument(this.onDidOpenCloseDocument, this, this.disposables); - notebooks.onDidChangeNotebookCellExecutionState( + notebookCellExecutions.onDidChangeNotebookCellExecutionState( this.onDidChangeNotebookCellExecutionState, this, this.disposables @@ -52,7 +50,11 @@ export class CellOutputMimeTypeTracker implements IExtensionSyncActivationServic dispose(this.disposables); } public async onDidChangeNotebookCellExecutionState(e: NotebookCellExecutionStateChangeEvent): Promise { - if (!isJupyterNotebook(e.cell.notebook) || this.isTelemetryDisabled) { + if ( + !isJupyterNotebook(e.cell.notebook) || + this.isTelemetryDisabled || + e.state !== NotebookCellExecutionState.Idle + ) { return; } this.checkCell(e.cell, 'onExecution'); diff --git a/src/platform/errors/errors.ts b/src/platform/errors/errors.ts index e4d652c1dee..fc861ef5ba2 100644 --- a/src/platform/errors/errors.ts +++ b/src/platform/errors/errors.ts @@ -40,6 +40,6 @@ function tagWithKernelRestarterFailed(stdErrOrStackTrace: string, tags: string[] } } -export function getExtensionSpecifcStack() { +export function getExtensionSpecificStack() { return (new Error().stack || '').split('\n').filter((l) => l.includes(JVSC_EXTENSION_ID)); } diff --git a/src/platform/notebooks/cellExecutionStateService.ts b/src/platform/notebooks/cellExecutionStateService.ts new file mode 100644 index 00000000000..1c8b7912605 --- /dev/null +++ b/src/platform/notebooks/cellExecutionStateService.ts @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { EventEmitter, workspace, type NotebookCell } from 'vscode'; +import { trackDisposable } from '../common/utils/lifecycle'; + +/** + * The execution state of a notebook cell. + */ +export enum NotebookCellExecutionState { + /** + * The cell is idle. + */ + Idle = 1, + /** + * Execution for the cell is pending. + */ + Pending = 2, + /** + * The cell is currently executing. + */ + Executing = 3 +} + +/** + * An event describing a cell execution state change. + */ +export interface NotebookCellExecutionStateChangeEvent { + /** + * The {@link NotebookCell cell} for which the execution state has changed. + */ + readonly cell: NotebookCell; + /** + * The new execution state of the cell. + */ + readonly state: NotebookCellExecutionState; +} + +export namespace notebookCellExecutions { + const eventEmitter = trackDisposable(new EventEmitter()); + + /** + * An {@link Event} which fires when the execution state of a cell has changed. + */ + // todo@API this is an event that is fired for a property that cells don't have and that makes me wonder + // how a correct consumer works, e.g the consumer could have been late and missed an event? + export const onDidChangeNotebookCellExecutionState = eventEmitter.event; + + export function changeCellState(cell: NotebookCell, state: NotebookCellExecutionState, executionOrder?: number) { + if (state !== NotebookCellExecutionState.Idle || !executionOrder) { + eventEmitter.fire({ cell, state }); + return; + } + // Wait for VS Code to update the cell execution state before firing the event. + const disposable = trackDisposable( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook !== cell.notebook) { + return; + } + const currentCellChange = e.cellChanges.find((c) => c.cell === cell); + if (currentCellChange?.cell?.executionSummary?.executionOrder === executionOrder) { + disposable.dispose(); + eventEmitter.fire({ cell, state: NotebookCellExecutionState.Idle }); + } + }) + ); + } +} diff --git a/src/standalone/api/kernels/api.vscode.common.test.ts b/src/standalone/api/kernels/api.vscode.common.test.ts index caab735bfc0..0e2e3727aa0 100644 --- a/src/standalone/api/kernels/api.vscode.common.test.ts +++ b/src/standalone/api/kernels/api.vscode.common.test.ts @@ -22,7 +22,7 @@ import { waitForExecutionCompletedSuccessfully } from '../../../test/datascience/notebook/helper'; import { getKernelsApi } from '.'; -import { raceTimeoutError } from '../../../platform/common/utils/async'; +import { createDeferred, raceTimeoutError } from '../../../platform/common/utils/async'; import { dispose } from '../../../platform/common/utils/lifecycle'; import { IKernel, IKernelProvider } from '../../../kernels/types'; import { IControllerRegistration, IVSCodeNotebookController } from '../../../notebooks/controllers/types'; @@ -31,6 +31,7 @@ import { JVSC_EXTENSION_ID_FOR_TESTS } from '../../../test/constants'; import { KernelError } from '../../../kernels/errors/kernelError'; import { JVSC_EXTENSION_ID } from '../../../platform/common/constants'; import { escapeStringToEmbedInPythonCode } from '../../../kernels/chat/generator'; +import { notebookCellExecutions } from '../../../platform/notebooks/cellExecutionStateService'; suiteMandatory('Kernel API Tests @typescript', function () { const disposables: IDisposable[] = []; @@ -110,7 +111,14 @@ suiteMandatory('Kernel API Tests @typescript', function () { // Ensure user has executed some code against this kernel. const cell = notebook.cellAt(0)!; - await Promise.all([runCell(cell), waitForExecutionCompletedSuccessfully(cell)]); + const executionOrderSet = createDeferred(); + const eventHandler = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.cell.executionSummary?.executionOrder) { + executionOrderSet.resolve(); + } + }); + disposables.push(eventHandler); + await Promise.all([runCell(cell), waitForExecutionCompletedSuccessfully(cell), executionOrderSet.promise]); const kernel = await kernels.getKernel(notebook.uri); if (!kernel) { diff --git a/src/standalone/executionAnalysis/symbols.ts b/src/standalone/executionAnalysis/symbols.ts index b5d6ade1722..d70a42c205a 100644 --- a/src/standalone/executionAnalysis/symbols.ts +++ b/src/standalone/executionAnalysis/symbols.ts @@ -11,6 +11,7 @@ import { Range, cellRangesToIndexes } from './common'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; const writeDecorationType = vscode.window.createTextEditorDecorationType({ after: { @@ -269,15 +270,12 @@ export class NotebookDocumentSymbolTracker { ); this._disposables.push( - vscode.notebooks.onDidChangeNotebookCellExecutionState((e) => { - if ( - e.state === vscode.NotebookCellExecutionState.Executing && - e.cell.document.languageId === 'python' - ) { + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.state === NotebookCellExecutionState.Executing && e.cell.document.languageId === 'python') { this._updateCellStatus(e.cell, CellExecutionStatus.Executing); } - if (e.state === vscode.NotebookCellExecutionState.Idle && e.cell.document.languageId === 'python') { + if (e.state === NotebookCellExecutionState.Idle && e.cell.document.languageId === 'python') { // just finished execution this._cellExecution.push({ cell: e.cell, diff --git a/src/standalone/import-export/importTracker.ts b/src/standalone/import-export/importTracker.ts index 6b091b8b6f5..2188034aa00 100644 --- a/src/standalone/import-export/importTracker.ts +++ b/src/standalone/import-export/importTracker.ts @@ -2,17 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { - Disposable, - NotebookCell, - NotebookCellExecutionState, - NotebookCellKind, - NotebookDocument, - TextDocument, - Uri, - notebooks, - workspace -} from 'vscode'; +import { Disposable, NotebookCell, NotebookCellKind, NotebookDocument, TextDocument, Uri, workspace } from 'vscode'; import { ResourceTypeTelemetryProperty, onDidChangeTelemetryEnablement, sendTelemetryEvent } from '../../telemetry'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { isCI, isTestExecution, JupyterNotebookView, PYTHON_LANGUAGE } from '../../platform/common/constants'; @@ -25,6 +15,7 @@ import { isJupyterNotebook } from '../../platform/common/utils'; import { isTelemetryDisabled } from '../../telemetry'; import { ResourceMap } from '../../platform/common/utils/map'; import { Delayer } from '../../platform/common/utils/async'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; /* Python has a fairly rich import statement. Originally the matching regexp was kept simple for @@ -63,6 +54,7 @@ export class ImportTracker implements IExtensionSyncActivationService { private pendingChecks = new ResourceMap(); private disposables = new DisposableStore(); private sentMatches = new Set(); + private readonly processedNotebookCells = new WeakMap(); private isTelemetryDisabled: boolean; constructor(@inject(IDisposableRegistry) disposables: IDisposableRegistry, delay = 1_000) { disposables.push(this.disposables); @@ -77,9 +69,9 @@ export class ImportTracker implements IExtensionSyncActivationService { this.disposables.add( workspace.onDidSaveNotebookDocument((t) => this.onOpenedOrClosedNotebookDocument(t, 'onOpenCloseOrSave')) ); - const delayer = new Delayer(delay); + const delayer = this.disposables.add(new Delayer(delay)); this.disposables.add( - notebooks.onDidChangeNotebookCellExecutionState((e) => { + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { void delayer.trigger(() => { if (e.state == NotebookCellExecutionState.Pending && !this.isTelemetryDisabled) { this.checkNotebookCell(e.cell, 'onExecution').catch(noop); @@ -143,11 +135,13 @@ export class ImportTracker implements IExtensionSyncActivationService { if ( !isJupyterNotebook(cell.notebook) || cell.kind !== NotebookCellKind.Code || - cell.document.languageId !== PYTHON_LANGUAGE + cell.document.languageId !== PYTHON_LANGUAGE || + this.processedNotebookCells.get(cell) === cell.document.version ) { return; } try { + this.processedNotebookCells.set(cell, cell.document.version); const resourceType = cell.notebook.notebookType === JupyterNotebookView ? 'notebook' : 'interactive'; await this.sendTelemetryForImports(this.getDocumentLines(cell.document), resourceType, when); } catch (ex) { diff --git a/src/standalone/import-export/importTracker.unit.test.ts b/src/standalone/import-export/importTracker.unit.test.ts index 89214476daf..298bc1c34e9 100644 --- a/src/standalone/import-export/importTracker.unit.test.ts +++ b/src/standalone/import-export/importTracker.unit.test.ts @@ -3,15 +3,10 @@ /* eslint-disable , , @typescript-eslint/no-explicit-any, no-multi-str, no-trailing-spaces */ import * as sinon from 'sinon'; +import * as fakeTimers from '@sinonjs/fake-timers'; import { assert, expect } from 'chai'; import { when } from 'ts-mockito'; -import { - EventEmitter, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, - NotebookCellKind, - NotebookDocument -} from 'vscode'; +import { Disposable, EventEmitter, NotebookCellKind, NotebookDocument } from 'vscode'; import { InteractiveWindowView, @@ -31,13 +26,13 @@ import { waitForCondition } from '../../test/common'; import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; import { EmptyEvent } from '../../platform/common/utils/events'; +import { notebookCellExecutions, NotebookCellExecutionState } from '../../platform/notebooks/cellExecutionStateService'; [true, false].forEach((useCustomMetadata) => { suite(`Import Tracker (${useCustomMetadata ? 'with custom metadata' : 'without custom metadata'})`, async () => { const oldValueOfVSC_JUPYTER_UNIT_TEST = isUnitTestExecution(); const oldValueOfVSC_JUPYTER_CI_TEST = isTestExecution(); let importTracker: ImportTracker; - let onDidChangeNotebookCellExecutionState: EventEmitter; let onDidOpenNbEvent: EventEmitter; let onDidCloseNbEvent: EventEmitter; let onDidSaveNbEvent: EventEmitter; @@ -51,6 +46,8 @@ import { EmptyEvent } from '../../platform/common/utils/events'; let sklearnHash: string; let randomHash: string; let disposables: IDisposable[] = []; + let clock: fakeTimers.InstalledClock; + class Reporter { public static eventNames: string[] = []; public static properties: Record[] = []; @@ -61,6 +58,8 @@ import { EmptyEvent } from '../../platform/common/utils/events'; resourceType: ResourceTypeTelemetryProperty['resourceType'] = undefined, ...hashes: string[] ) { + clock.tick(1); + void clock.runAllAsync(); if (hashes.length > 0) { await waitForCondition( async () => { @@ -130,16 +129,12 @@ import { EmptyEvent } from '../../platform/common/utils/events'; onDidOpenNbEvent = new EventEmitter(); onDidCloseNbEvent = new EventEmitter(); onDidSaveNbEvent = new EventEmitter(); - onDidChangeNotebookCellExecutionState = new EventEmitter(); disposables.push(onDidOpenNbEvent); disposables.push(onDidCloseNbEvent); disposables.push(onDidSaveNbEvent); when(mockedVSCodeNamespaces.workspace.onDidOpenNotebookDocument).thenReturn(onDidOpenNbEvent.event); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn(onDidCloseNbEvent.event); when(mockedVSCodeNamespaces.workspace.onDidSaveNotebookDocument).thenReturn(onDidSaveNbEvent.event); - when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn( - onDidChangeNotebookCellExecutionState.event - ); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); when(mockedVSCodeNamespaces.workspace.onDidChangeConfiguration).thenReturn(EmptyEvent); when(mockedVSCodeNamespaces.workspace.getConfiguration('telemetry')).thenReturn({ @@ -151,6 +146,8 @@ import { EmptyEvent } from '../../platform/common/utils/events'; } } as any); importTracker = new ImportTracker(disposables, 0); + clock = fakeTimers.install(); + disposables.push(new Disposable(() => clock.uninstall())); }); teardown(() => { sinon.restore(); @@ -345,32 +342,17 @@ import { EmptyEvent } from '../../platform/common/utils/events'; const nb = createMockedNotebookDocument(useCustomMetadata, [ { kind: NotebookCellKind.Code, languageId: 'python', value: code } ]); - onDidChangeNotebookCellExecutionState.fire({ - cell: nb.cellAt(0), - state: NotebookCellExecutionState.Pending - }); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Pending); await Reporter.expectHashes('onExecution', 'notebook', numpyHash); // Executing the cell multiple will have no effect, the telemetry is only sent once. - onDidChangeNotebookCellExecutionState.fire({ - cell: nb.cellAt(0), - state: NotebookCellExecutionState.Pending - }); - onDidChangeNotebookCellExecutionState.fire({ - cell: nb.cellAt(0), - state: NotebookCellExecutionState.Executing - }); - onDidChangeNotebookCellExecutionState.fire({ cell: nb.cellAt(0), state: NotebookCellExecutionState.Idle }); - onDidChangeNotebookCellExecutionState.fire({ - cell: nb.cellAt(0), - state: NotebookCellExecutionState.Pending - }); - onDidChangeNotebookCellExecutionState.fire({ - cell: nb.cellAt(0), - state: NotebookCellExecutionState.Executing - }); - onDidChangeNotebookCellExecutionState.fire({ cell: nb.cellAt(0), state: NotebookCellExecutionState.Idle }); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Pending); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Executing); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Idle); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Pending); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Executing); + notebookCellExecutions.changeCellState(nb.cellAt(0), NotebookCellExecutionState.Idle); await Reporter.expectHashes('onExecution', 'notebook', numpyHash); }); diff --git a/src/standalone/survey/dataScienceSurveyBanner.node.ts b/src/standalone/survey/dataScienceSurveyBanner.node.ts index b58941f34df..eacde1685e4 100644 --- a/src/standalone/survey/dataScienceSurveyBanner.node.ts +++ b/src/standalone/survey/dataScienceSurveyBanner.node.ts @@ -2,15 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { - Disposable, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, - UIKind, - env, - notebooks, - window -} from 'vscode'; +import { Disposable, UIKind, env, window } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { traceError } from '../../platform/logging'; import { @@ -27,6 +19,11 @@ import { noop } from '../../platform/common/utils/misc'; import { openInBrowser } from '../../platform/common/net/browser'; import { getVSCodeChannel } from '../../platform/common/application/applicationEnvironment'; import type { IDisposable } from '@c4312/evt'; +import { + NotebookCellExecutionState, + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../platform/notebooks/cellExecutionStateService'; export const ISurveyBanner = Symbol('ISurveyBanner'); export interface ISurveyBanner extends IExtensionSyncActivationService, IJupyterExtensionBanner {} @@ -122,11 +119,12 @@ export class DataScienceSurveyBanner implements IJupyterExtensionBanner, IExtens } public activate() { - this.onDidChangeNotebookCellExecutionStateHandler = notebooks.onDidChangeNotebookCellExecutionState( - this.onDidChangeNotebookCellExecutionState, - this, - this.disposables - ); + this.onDidChangeNotebookCellExecutionStateHandler = + notebookCellExecutions.onDidChangeNotebookCellExecutionState( + this.onDidChangeNotebookCellExecutionState, + this, + this.disposables + ); } public async showBanner(type: BannerType): Promise { diff --git a/src/test/datascience/interactiveWindow.vscode.common.test.ts b/src/test/datascience/interactiveWindow.vscode.common.test.ts index 866ff50e650..509ecb208b8 100644 --- a/src/test/datascience/interactiveWindow.vscode.common.test.ts +++ b/src/test/datascience/interactiveWindow.vscode.common.test.ts @@ -396,7 +396,6 @@ ${actualCode} await waitForExecutionCompletedSuccessfully(secondCell!); await waitForTextOutput(secondCell!, '1'); }); - test('Error stack traces have correct line hrefs with mix of cell sources', async function () { const settings = vscode.workspace.getConfiguration('jupyter', null); await settings.update('interactiveWindow.creationMode', 'single'); diff --git a/src/test/datascience/notebook/helper.ts b/src/test/datascience/notebook/helper.ts index 75e5969740b..bcff167717b 100644 --- a/src/test/datascience/notebook/helper.ts +++ b/src/test/datascience/notebook/helper.ts @@ -23,7 +23,6 @@ import { Memento, NotebookCell, NotebookCellData, - NotebookCellExecutionState, NotebookCellKind, NotebookCellOutputItem, NotebookData, @@ -43,7 +42,6 @@ import { debug, env, languages, - notebooks, window, workspace } from 'vscode'; @@ -102,6 +100,10 @@ import { JupyterConnection } from '../../../kernels/jupyter/connection/jupyterCo import { JupyterLabHelper } from '../../../kernels/jupyter/session/jupyterLabHelper'; import { getRootFolder } from '../../../platform/common/application/workspace.base'; import { activateIPynbExtension, useCustomMetadata } from '../../../platform/common/utils'; +import { + NotebookCellExecutionState, + notebookCellExecutions +} from '../../../platform/notebooks/cellExecutionStateService'; // Running in Conda environments, things can be a little slower. export const defaultNotebookTestTimeout = 60_000; @@ -919,6 +921,7 @@ export async function waitForCellExecutionToComplete(cell: NotebookCell) { defaultNotebookTestTimeout, 'Execution did not complete' ); + await sleep(100); } export async function waitForCellExecutionState( @@ -928,7 +931,7 @@ export async function waitForCellExecutionState( timeout: number = defaultNotebookTestTimeout ) { const deferred = createDeferred(); - const disposable = notebooks.onDidChangeNotebookCellExecutionState((e) => { + const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { if (e.cell !== cell) { return; } @@ -969,10 +972,14 @@ export async function waitForExecutionCompletedSuccessfully( () => `Cell ${cell.index + 1} did not complete successfully, State = ${NotebookCellStateTracker.getCellStatus( cell - )}` + )}, & state = ${NotebookCellStateTracker.getCellState(cell)}` ), waitForCellExecutionToComplete(cell) ]); + // Some of the tests run the exact same checks, + // hence we can have a race condition + // Wait for an additional 100ms to ensure all code has been executed. + await sleep(100); } export async function waitForCompletions( diff --git a/src/test/datascience/notebook/interruptRestart.vscode.test.ts b/src/test/datascience/notebook/interruptRestart.vscode.test.ts index 9f939764ddd..37b40f9059c 100644 --- a/src/test/datascience/notebook/interruptRestart.vscode.test.ts +++ b/src/test/datascience/notebook/interruptRestart.vscode.test.ts @@ -3,7 +3,7 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { NotebookCellExecutionState, window } from 'vscode'; +import { window } from 'vscode'; import { traceError, traceInfo } from '../../../platform/logging'; import { IConfigurationService, IDisposable, IJupyterSettings, ReadWrite } from '../../../platform/common/types'; import { noop } from '../../../platform/common/utils/misc'; @@ -25,6 +25,7 @@ import { import { hasErrorOutput, NotebookCellStateTracker, getTextOutputValue } from '../../../kernels/execution/helpers'; import { TestNotebookDocument, createKernelController } from './executionHelper'; import { captureScreenShot } from '../../common'; +import { NotebookCellExecutionState } from '../../../platform/notebooks/cellExecutionStateService'; /* eslint-disable @typescript-eslint/no-explicit-any, no-invalid-this, */ /* diff --git a/src/webviews/extension-side/variablesView/notebookWatcher.ts b/src/webviews/extension-side/variablesView/notebookWatcher.ts index 23e8f3234cc..a753815deac 100644 --- a/src/webviews/extension-side/variablesView/notebookWatcher.ts +++ b/src/webviews/extension-side/variablesView/notebookWatcher.ts @@ -2,18 +2,7 @@ // Licensed under the MIT License. import { inject, injectable } from 'inversify'; -import { - Event, - EventEmitter, - NotebookCell, - NotebookCellExecutionState, - NotebookCellExecutionStateChangeEvent, - NotebookDocument, - NotebookEditor, - notebooks, - window, - workspace -} from 'vscode'; +import { Event, EventEmitter, NotebookCell, NotebookDocument, NotebookEditor, window, workspace } from 'vscode'; import { IKernel, IKernelProvider } from '../../../kernels/types'; import { IActiveNotebookChangedEvent, INotebookWatcher } from './types'; import { IInteractiveWindowProvider } from '../../../interactive-window/types'; @@ -21,6 +10,11 @@ import { IDisposableRegistry } from '../../../platform/common/types'; import { IDataViewerFactory } from '../dataviewer/types'; import { JupyterNotebookView } from '../../../platform/common/constants'; import { isJupyterNotebook } from '../../../platform/common/utils'; +import { + NotebookCellExecutionState, + notebookCellExecutions, + type NotebookCellExecutionStateChangeEvent +} from '../../../platform/notebooks/cellExecutionStateService'; type KernelStateEventArgs = { notebook: NotebookDocument; @@ -98,7 +92,7 @@ export class NotebookWatcher implements INotebookWatcher { this, this.disposables ); - notebooks.onDidChangeNotebookCellExecutionState( + notebookCellExecutions.onDidChangeNotebookCellExecutionState( this.onDidChangeNotebookCellExecutionState, this, this.disposables