diff --git a/src/api.proposed.kernelStartHook.d.ts b/src/api.proposed.kernelStartHook.d.ts index 5377ee50f08..3b298851726 100644 --- a/src/api.proposed.kernelStartHook.d.ts +++ b/src/api.proposed.kernelStartHook.d.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import type { Event, Uri } from 'vscode'; +import type { CancellationToken, Event, Uri } from 'vscode'; declare module './api' { export interface Kernels { @@ -11,6 +11,16 @@ declare module './api' { onDidStart: Event<{ uri: Uri; kernel: Kernel; + token: CancellationToken; + /** + * Allows to pause the event loop until the provided thenable resolved. + * This can be useful to ensure startup code is executed before user code. + * + * *Note:* This function can only be called during event dispatch. + * + * @param thenable A thenable that delays kernel startup. + */ + waitUntil(thenable: Thenable): void; }>; } } diff --git a/src/kernels/kernel.ts b/src/kernels/kernel.ts index a63893ad1fb..8b482d45d1e 100644 --- a/src/kernels/kernel.ts +++ b/src/kernels/kernel.ts @@ -15,7 +15,8 @@ import { Memento, CancellationError, window, - workspace + workspace, + CancellationToken } from 'vscode'; import { CodeSnippets, @@ -66,6 +67,7 @@ import { KernelInterruptTimeoutError } from './errors/kernelInterruptTimeoutErro import { dispose } from '../platform/common/utils/lifecycle'; import { getCachedVersion, getEnvironmentType } from '../platform/interpreter/helpers'; import { getNotebookTelemetryTracker } from './telemetry/notebookTelemetry'; +import { AsyncEmitter } from '../platform/common/event'; const widgetVersionOutPrefix = 'e976ee50-99ed-4aba-9b6b-9dcd5634d07d:IPyWidgets:'; /** @@ -129,7 +131,7 @@ abstract class BaseKernel implements IBaseKernel { get onStarted(): Event { return this._onStarted.event; } - get onPostInitialized(): Event { + get onPostInitialized(): Event<{ token: CancellationToken; waitUntil(thenable: Thenable): void }> { return this._onPostInitialized.event; } get onDisposed(): Event { @@ -183,7 +185,10 @@ abstract class BaseKernel implements IBaseKernel { private readonly _onStatusChanged = new EventEmitter(); private readonly _onRestarted = new EventEmitter(); private readonly _onStarted = new EventEmitter(); - private readonly _onPostInitialized = new EventEmitter(); + private readonly _onPostInitialized = new AsyncEmitter<{ + waitUntil: (thenable: Thenable) => void; + token: CancellationToken; + }>(); private readonly _onDisposed = new EventEmitter(); private _jupyterSessionPromise?: Promise; private readonly hookedSessionForEvents = new WeakSet(); @@ -196,6 +201,10 @@ abstract class BaseKernel implements IBaseKernel { public get restarting() { return this._restartPromise || Promise.resolve(); } + private _postInitializingDeferred = createDeferred(); + public get postInitializing() { + return this._postInitializingDeferred.promise; + } constructor( public readonly id: string, public readonly uri: Uri, @@ -257,10 +266,12 @@ abstract class BaseKernel implements IBaseKernel { return this.startJupyterSession(options).then((result) => { // If we started and the UI is no longer disabled (ie., a user executed a cell) // then we can signal that the kernel was created and can be used by third-party extensions. - // We also only want to fire off a single event here. + // We also only want to fire off a single. event here. if (!options?.disableUI && !this._postInitializedOnStart) { - this._onPostInitialized.fire(); this._postInitializedOnStart = true; + void this._onPostInitialized.fireAsync({}, this.startCancellation.token).then(() => { + this._postInitializingDeferred.resolve(); + }); } return result; }); @@ -428,7 +439,9 @@ abstract class BaseKernel implements IBaseKernel { this._onRestarted.fire(); // Also signal that the kernel post initialization completed. - this._onPostInitialized.fire(); + void this._onPostInitialized.fireAsync({}, this.startCancellation.token).then(() => { + this._postInitializingDeferred.resolve(); + }); } catch (ex) { logger.error(`Failed to restart kernel ${getDisplayPath(this.uri)}`, ex); throw ex; diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 2c27fee043c..c7785d31d93 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -166,6 +166,10 @@ export class NotebookKernelExecution implements INotebookKernelExecution { const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false))); traceCellMessage(cell, `NotebookKernelExecution.executeCell (3), ${getDisplayPath(cell.notebook.uri)}`); + + // Wait for the kernel to complete post initialization before queueing the cell in case + // we need to allow extensions to run code before the initial user-triggered execution + await this.kernel.postInitializing; const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise); executionQueue.queueCell(cell, codeOverride); let success = true; diff --git a/src/kernels/kernelProvider.base.ts b/src/kernels/kernelProvider.base.ts index f3dd4609bc3..9853efe57e3 100644 --- a/src/kernels/kernelProvider.base.ts +++ b/src/kernels/kernelProvider.base.ts @@ -2,7 +2,7 @@ // Licensed under the MIT License. import type { KernelMessage } from '@jupyterlab/services'; -import { Event, EventEmitter, NotebookDocument, Uri, workspace } from 'vscode'; +import { CancellationToken, Event, EventEmitter, NotebookDocument, Uri, workspace } from 'vscode'; import { logger } from '../platform/logging'; import { getDisplayPath } from '../platform/common/platform/fs-paths'; import { IAsyncDisposable, IAsyncDisposableRegistry, IDisposableRegistry } from '../platform/common/types'; @@ -36,7 +36,11 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider { protected readonly _onDidCreateKernel = new EventEmitter(); protected readonly _onDidDisposeKernel = new EventEmitter(); protected readonly _onKernelStatusChanged = new EventEmitter<{ status: KernelMessage.Status; kernel: IKernel }>(); - protected readonly _onDidPostInitializeKernel = new EventEmitter(); + protected readonly _onDidPostInitializeKernel = new EventEmitter<{ + kernel: IKernel; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }>(); public readonly onKernelStatusChanged = this._onKernelStatusChanged.event; public get kernels() { const kernels = new Set(); @@ -75,7 +79,11 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider { public get onDidCreateKernel(): Event { return this._onDidCreateKernel.event; } - public get onDidPostInitializeKernel(): Event { + public get onDidPostInitializeKernel(): Event<{ + kernel: IKernel; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }> { return this._onDidPostInitializeKernel.event; } public get(uriOrNotebook: Uri | NotebookDocument | string): IKernel | undefined { @@ -192,7 +200,11 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP protected readonly _onDidStartKernel = new EventEmitter(); protected readonly _onDidCreateKernel = new EventEmitter(); protected readonly _onDidDisposeKernel = new EventEmitter(); - protected readonly _onDidPostInitializeKernel = new EventEmitter(); + protected readonly _onDidPostInitializeKernel = new EventEmitter<{ + kernel: IThirdPartyKernel; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }>(); protected readonly _onKernelStatusChanged = new EventEmitter<{ status: KernelMessage.Status; kernel: IThirdPartyKernel; @@ -234,7 +246,11 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP public get onDidCreateKernel(): Event { return this._onDidCreateKernel.event; } - public get onDidPostInitializeKernel(): Event { + public get onDidPostInitializeKernel(): Event<{ + kernel: IThirdPartyKernel; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }> { return this._onDidPostInitializeKernel.event; } public get(uri: Uri | string): IThirdPartyKernel | undefined { diff --git a/src/kernels/kernelProvider.node.ts b/src/kernels/kernelProvider.node.ts index 158b4b59593..a5bdd5bcbc9 100644 --- a/src/kernels/kernelProvider.node.ts +++ b/src/kernels/kernelProvider.node.ts @@ -98,8 +98,8 @@ export class KernelProvider extends BaseCoreKernelProvider { this.disposables ); kernel.onPostInitialized( - () => { - this._onDidPostInitializeKernel.fire(kernel); + ({ token, waitUntil }) => { + this._onDidPostInitializeKernel.fire({ kernel, token, waitUntil }); }, this, this.disposables @@ -160,8 +160,8 @@ export class ThirdPartyKernelProvider extends BaseThirdPartyKernelProvider { this.disposables ); kernel.onPostInitialized( - () => { - this._onDidPostInitializeKernel.fire(kernel); + ({ token, waitUntil }) => { + this._onDidPostInitializeKernel.fire({ kernel, token, waitUntil }); }, this, this.disposables diff --git a/src/kernels/types.ts b/src/kernels/types.ts index b9bc32cf8a4..d0023065f7a 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -355,8 +355,9 @@ export interface IBaseKernel extends IAsyncDisposable { readonly onDisposed: Event; readonly onStarted: Event; readonly onRestarted: Event; - readonly onPostInitialized: Event; + readonly onPostInitialized: Event<{ token: CancellationToken; waitUntil(thenable: Thenable): void }>; readonly restarting: Promise; + readonly postInitializing: Promise; readonly status: KernelMessage.Status; readonly disposed: boolean; readonly disposing: boolean; @@ -513,7 +514,11 @@ export interface IBaseKernelProvider extends IAsyncDispos onDidRestartKernel: Event; onDidDisposeKernel: Event; onKernelStatusChanged: Event<{ status: KernelMessage.Status; kernel: T }>; - onDidPostInitializeKernel: Event; + onDidPostInitializeKernel: Event<{ + kernel: T; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }>; } /** diff --git a/src/platform/common/event.ts b/src/platform/common/event.ts new file mode 100644 index 00000000000..abb3be5b4f4 --- /dev/null +++ b/src/platform/common/event.ts @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { Event, CancellationToken, Disposable } from 'vscode'; + +export interface IWaitUntil { + token: CancellationToken; + waitUntil(thenable: Promise): void; +} +export type IWaitUntilData = Omit, 'token'>; + +// Based on AsyncEmitter in https://github.com/microsoft/vscode/blob/88e6face01795b161523824ba075444fad55466a/src/vs/base/common/event.ts#L1303 +export class AsyncEmitter { + private _listeners: Set<(e: T) => unknown> = new Set(); + private _asyncDeliveryQueue?: Array<[(ev: T) => void, IWaitUntilData]>; + private _event: Event | undefined; + + public get event(): Event { + if (!this._event) { + this._event = (listener: (e: T) => unknown): Disposable => { + this._listeners.add(listener); + return { + dispose: () => this._listeners.delete(listener) + }; + }; + } + return this._event; + } + + async fireAsync(data: IWaitUntilData, token: CancellationToken): Promise { + if (!this._listeners) { + return; + } + + if (!this._asyncDeliveryQueue) { + this._asyncDeliveryQueue = []; + } + + for (const listener of this._listeners) { + this._asyncDeliveryQueue!.push([listener, data]); + } + + while (this._asyncDeliveryQueue.length > 0 && !token.isCancellationRequested) { + const [listener, data] = this._asyncDeliveryQueue.shift()!; + const thenables: Promise[] = []; + + const event = { + ...data, + token, + waitUntil: (p: Promise): void => { + if (Object.isFrozen(thenables)) { + throw new Error('waitUntil can NOT be called asynchronous'); + } + thenables.push(p); + } + }; + + try { + listener(event); + } catch (e) { + console.error(e); + continue; + } + + // freeze thenables-collection to enforce sync-calls to + // wait until and then wait for all thenables to resolve + Object.freeze(thenables); + + await Promise.allSettled(thenables).then((values) => { + for (const value of values) { + if (value.status === 'rejected') { + console.error(value.reason); + } + } + }); + } + } + + dispose(): void { + this._listeners.clear(); + } +} diff --git a/src/standalone/api/kernels/index.ts b/src/standalone/api/kernels/index.ts index b84065cb414..84bed1e8d1c 100644 --- a/src/standalone/api/kernels/index.ts +++ b/src/standalone/api/kernels/index.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { Uri, workspace, EventEmitter } from 'vscode'; +import { Uri, workspace, EventEmitter, CancellationToken } from 'vscode'; import { Kernel, Kernels } from '../../../api'; import { ServiceContainer } from '../../../platform/ioc/container'; import { IKernel, IKernelProvider } from '../../../kernels/types'; @@ -12,7 +12,9 @@ import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../.. import { IDisposableRegistry } from '../../../platform/common/types'; const kernelCache = new WeakMap(); -let _onDidStart: EventEmitter<{ uri: Uri; kernel: Kernel }> | undefined = undefined; +let _onDidStart: + | EventEmitter<{ uri: Uri; kernel: Kernel; token: CancellationToken; waitUntil(thenable: Thenable): void }> + | undefined = undefined; function getWrappedKernel(kernel: IKernel, extensionId: string) { let wrappedKernel = kernelCache.get(kernel) || createKernelApiForExtension(extensionId, kernel); @@ -55,11 +57,21 @@ export function getKernelsApi(extensionId: string): Kernels { if (!_onDidStart) { const kernelProvider = ServiceContainer.instance.get(IKernelProvider); const disposableRegistry = ServiceContainer.instance.get(IDisposableRegistry); - _onDidStart = new EventEmitter<{ uri: Uri; kernel: Kernel }>(); + _onDidStart = new EventEmitter<{ + uri: Uri; + kernel: Kernel; + token: CancellationToken; + waitUntil(thenable: Thenable): void; + }>(); disposableRegistry.push( - kernelProvider.onDidPostInitializeKernel((kernel) => { - _onDidStart?.fire({ uri: kernel.uri, kernel: getWrappedKernel(kernel, extensionId) }); + kernelProvider.onDidPostInitializeKernel(({ kernel, token, waitUntil }) => { + _onDidStart?.fire({ + uri: kernel.uri, + kernel: getWrappedKernel(kernel, extensionId), + token, + waitUntil + }); }), _onDidStart, { dispose: () => (_onDidStart = undefined) }