From 11431255cd164c9488eb3e00b4255da02fa3f89a Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Thu, 12 May 2022 12:01:24 +1000 Subject: [PATCH] Revision to displaying cached remote kernel specs and live kernels (#9985) * Misc * Misc * Misc * Misc * oops * Misc * Add tests * Misc * Misc * Address code review comments * Misc * oops * Fix tests --- package.nls.json | 1 - .../liveRemoteKernelConnectionTracker.ts | 100 ++++++++ .../preferredRemoteKernelIdProvider.ts | 2 +- .../jupyter/remoteKernelConnectionHandler.ts | 81 ++++++ src/kernels/jupyter/serviceRegistry.node.ts | 18 +- src/kernels/jupyter/serviceRegistry.web.ts | 18 +- src/kernels/jupyter/types.ts | 19 +- src/kernels/kernelFinder.node.ts | 15 +- src/kernels/kernelFinder.web.ts | 15 +- src/kernels/types.ts | 2 +- .../controllers/notebookControllerManager.ts | 16 +- .../controllers/vscodeNotebookController.ts | 55 +---- src/notebooks/types.ts | 6 +- src/platform/common/utils/localize.ts | 4 - .../localKernelFinder.unit.test.ts | 8 +- .../remoteKernelFinder.unit.test.ts | 144 ++++++++++- ...RemoteKernelConnectionTracker.unit.test.ts | 225 +++++++++++++++++ ...remoteKernelConnectionHandler.unit.test.ts | 231 ++++++++++++++++++ src/test/mocks/vsc/extHostedTypes.ts | 3 + 19 files changed, 877 insertions(+), 86 deletions(-) create mode 100644 src/kernels/jupyter/liveRemoteKernelConnectionTracker.ts create mode 100644 src/kernels/jupyter/remoteKernelConnectionHandler.ts create mode 100644 src/test/kernels/jupyter/liveRemoteKernelConnectionTracker.unit.test.ts create mode 100644 src/test/kernels/jupyter/remoteKernelConnectionHandler.unit.test.ts diff --git a/package.nls.json b/package.nls.json index 8c9f4884338..57e1ad84b78 100644 --- a/package.nls.json +++ b/package.nls.json @@ -294,7 +294,6 @@ "DataScience.remoteJupyterConnectionFailedWithoutServerWithError": "Connection failure. Verify the server is running and reachable. ({0}).", "DataScience.jupyterNotebookRemoteConnectFailedWeb": "Failed to connect to remote Jupyter server.\r\nCheck that the Jupyter Server URI can be reached from a browser.\r\n{0}. Click [here](https://aka.ms/vscjremoteweb) for more information.", "DataScience.jupyterNotebookRemoteConnectSelfCertsFailed": "Failed to connect to remote Jupyter notebook.\r\nSpecified server is using self signed certs. Enable Allow Unauthorized Remote Connection setting to connect anyways\r\n{0}\r\n{1}", - "DataScience.jupyterRemoteConnectFailedModalMessage": "Failed to connect to the remote Jupyter Server. View Jupyter log for further details.", "DataScience.changeJupyterRemoteConnection": "Change Jupyter Server connection.", "DataScience.notebookVersionFormat": "Jupyter Notebook Version: {0}", "DataScience.jupyterKernelSpecNotFound": "Cannot create a Jupyter kernel spec and none are available for use", diff --git a/src/kernels/jupyter/liveRemoteKernelConnectionTracker.ts b/src/kernels/jupyter/liveRemoteKernelConnectionTracker.ts new file mode 100644 index 00000000000..ad8c7ec9bb8 --- /dev/null +++ b/src/kernels/jupyter/liveRemoteKernelConnectionTracker.ts @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { inject, injectable, named } from 'inversify'; +import { Memento, Uri } from 'vscode'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { GLOBAL_MEMENTO, IDisposableRegistry, IMemento } from '../../platform/common/types'; +import { noop } from '../../platform/common/utils/misc'; +import { LiveRemoteKernelConnectionMetadata } from '../types'; +import { computeServerId } from './jupyterUtils'; +import { IJupyterServerUriStorage, ILiveRemoteKernelConnectionUsageTracker } from './types'; + +export const mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources = 'removeKernelUrisAndSessionsUsedByResources'; + +type ServerId = string; +type KernelId = string; +type UriString = string; +type UriSessionUsedByResources = Record>; + +/** + * Class to track the remote kernels that have been used by notebooks. + */ +@injectable() +export class LiveRemoteKernelConnectionUsageTracker + implements IExtensionSyncActivationService, ILiveRemoteKernelConnectionUsageTracker +{ + private usedRemoteKernelServerIdsAndSessions: UriSessionUsedByResources = {}; + constructor( + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, + @inject(IJupyterServerUriStorage) private readonly uriStorage: IJupyterServerUriStorage, + @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly memento: Memento + ) {} + public activate(): void { + this.usedRemoteKernelServerIdsAndSessions = this.memento.get( + mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, + {} + ); + this.uriStorage.onDidRemoveUri(this.onDidRemoveUri, this, this.disposables); + } + + public wasKernelUsed(connection: LiveRemoteKernelConnectionMetadata) { + return ( + connection.serverId in this.usedRemoteKernelServerIdsAndSessions && + typeof connection.kernelModel.id === 'string' && + connection.kernelModel.id in this.usedRemoteKernelServerIdsAndSessions[connection.serverId] + ); + } + public trackKernelIdAsUsed(resource: Uri, serverId: string, kernelId: string) { + this.usedRemoteKernelServerIdsAndSessions[serverId] = this.usedRemoteKernelServerIdsAndSessions[serverId] || {}; + this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId] = + this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId] || []; + const uris = this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; + if (uris.includes(resource.toString())) { + return; + } + uris.push(resource.toString()); + this.memento + .update( + mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, + this.usedRemoteKernelServerIdsAndSessions + ) + .then(noop, noop); + } + public trackKernelIdAsNotUsed(resource: Uri, serverId: string, kernelId: string) { + if (!(serverId in this.usedRemoteKernelServerIdsAndSessions)) { + return; + } + if (!(kernelId in this.usedRemoteKernelServerIdsAndSessions[serverId])) { + return; + } + const uris = this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; + if (!Array.isArray(uris) || !uris.includes(resource.toString())) { + return; + } + uris.splice(uris.indexOf(resource.toString()), 1); + if (uris.length === 0) { + delete this.usedRemoteKernelServerIdsAndSessions[serverId][kernelId]; + } + if (Object.keys(this.usedRemoteKernelServerIdsAndSessions[serverId]).length === 0) { + delete this.usedRemoteKernelServerIdsAndSessions[serverId]; + } + + this.memento + .update( + mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, + this.usedRemoteKernelServerIdsAndSessions + ) + .then(noop, noop); + } + private onDidRemoveUri(uri: string) { + const serverId = computeServerId(uri); + delete this.usedRemoteKernelServerIdsAndSessions[serverId]; + this.memento + .update( + mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, + this.usedRemoteKernelServerIdsAndSessions + ) + .then(noop, noop); + } +} diff --git a/src/kernels/jupyter/preferredRemoteKernelIdProvider.ts b/src/kernels/jupyter/preferredRemoteKernelIdProvider.ts index eab51cf3b22..72805914889 100644 --- a/src/kernels/jupyter/preferredRemoteKernelIdProvider.ts +++ b/src/kernels/jupyter/preferredRemoteKernelIdProvider.ts @@ -44,7 +44,7 @@ export class PreferredRemoteKernelIdProvider { public async storePreferredRemoteKernelId(uri: Uri, id: string): Promise { await this.updatePreferredRemoteKernelIdInternal(uri, id); } - public async updatePreferredRemoteKernelIdInternal(uri: Uri, id?: string): Promise { + private async updatePreferredRemoteKernelIdInternal(uri: Uri, id?: string): Promise { let requiresUpdate = false; // Don't update in memory representation. diff --git a/src/kernels/jupyter/remoteKernelConnectionHandler.ts b/src/kernels/jupyter/remoteKernelConnectionHandler.ts new file mode 100644 index 00000000000..83b603a1d17 --- /dev/null +++ b/src/kernels/jupyter/remoteKernelConnectionHandler.ts @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { inject, injectable } from 'inversify'; +import { Disposable, NotebookDocument } from 'vscode'; +import { IVSCodeNotebookController } from '../../notebooks/controllers/types'; +import { INotebookControllerManager } from '../../notebooks/types'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IDisposableRegistry } from '../../platform/common/types'; +import { noop } from '../../platform/common/utils/misc'; +import { traceInfo } from '../../platform/logging'; +import { IKernel, IKernelProvider, isLocalConnection } from '../types'; +import { PreferredRemoteKernelIdProvider } from './preferredRemoteKernelIdProvider'; +import { ILiveRemoteKernelConnectionUsageTracker } from './types'; + +@injectable() +export class RemoteKernelConnectionHandler implements IExtensionSyncActivationService { + constructor( + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, + @inject(INotebookControllerManager) private readonly controllers: INotebookControllerManager, + @inject(ILiveRemoteKernelConnectionUsageTracker) + private readonly liveKernelTracker: ILiveRemoteKernelConnectionUsageTracker, + @inject(PreferredRemoteKernelIdProvider) + private readonly preferredRemoteKernelIdProvider: PreferredRemoteKernelIdProvider + ) {} + activate(): void { + this.kernelProvider.onDidStartKernel(this.onDidStartKernel, this, this.disposables); + this.controllers.onNotebookControllerSelectionChanged( + this.onNotebookControllerSelectionChanged, + this, + this.disposables + ); + } + private onNotebookControllerSelectionChanged({ + selected, + notebook, + controller + }: { + selected: boolean; + notebook: NotebookDocument; + controller: IVSCodeNotebookController; + }) { + if (controller.connection.kind === 'connectToLiveRemoteKernel' && controller.connection.kernelModel.id) { + if (selected) { + this.liveKernelTracker.trackKernelIdAsUsed( + notebook.uri, + controller.connection.serverId, + controller.connection.kernelModel.id + ); + } else { + this.liveKernelTracker.trackKernelIdAsNotUsed( + notebook.uri, + controller.connection.serverId, + controller.connection.kernelModel.id + ); + } + } + if (isLocalConnection(controller.connection)) { + this.preferredRemoteKernelIdProvider.clearPreferredRemoteKernelId(notebook.uri).catch(noop); + } + } + private onDidStartKernel(kernel: IKernel) { + if (kernel.creator !== 'jupyterExtension' || !kernel.resourceUri) { + return; + } + const resource = kernel.resourceUri; + if (kernel.kernelConnectionMetadata.kind === 'startUsingRemoteKernelSpec') { + const serverId = kernel.kernelConnectionMetadata.serverId; + const subscription = kernel.kernelSocket.subscribe((info) => { + const kernelId = info?.options.id; + if (!kernel.disposed && !kernel.disposing && kernelId) { + traceInfo(`Updating preferred kernel for remote notebook ${kernelId}`); + this.preferredRemoteKernelIdProvider.storePreferredRemoteKernelId(resource, kernelId).catch(noop); + this.liveKernelTracker.trackKernelIdAsUsed(resource, serverId, kernelId); + } + }); + this.disposables.push(new Disposable(() => subscription.unsubscribe())); + } + } +} diff --git a/src/kernels/jupyter/serviceRegistry.node.ts b/src/kernels/jupyter/serviceRegistry.node.ts index ef020d7012c..c3a2375fafa 100644 --- a/src/kernels/jupyter/serviceRegistry.node.ts +++ b/src/kernels/jupyter/serviceRegistry.node.ts @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import { IExtensionSingleActivationService } from '../../platform/activation/types'; +import { IExtensionSingleActivationService, IExtensionSyncActivationService } from '../../platform/activation/types'; import { IServiceManager } from '../../platform/ioc/types'; import { IRemoteKernelFinder } from '../raw/types'; import { INotebookProvider } from '../types'; @@ -37,6 +37,8 @@ import { NotebookStarter } from './launcher/notebookStarter.node'; import { ServerConnectionType } from './launcher/serverConnectionType'; import { ServerPreload } from './launcher/serverPreload.node'; import { JupyterServerUriStorage } from './launcher/serverUriStorage'; +import { LiveRemoteKernelConnectionUsageTracker } from './liveRemoteKernelConnectionTracker'; +import { RemoteKernelConnectionHandler } from './remoteKernelConnectionHandler'; import { RemoteKernelFinder } from './remoteKernelFinder'; import { JupyterServerSelector } from './serverSelector'; import { BackingFileCreator } from './session/backingFileCreator.node'; @@ -61,7 +63,8 @@ import { INotebookStarter, IJupyterRequestCreator, IJupyterRequestAgentCreator, - INotebookServerFactory + INotebookServerFactory, + ILiveRemoteKernelConnectionUsageTracker } from './types'; import { IJupyterCommandFactory, IJupyterSubCommandExecutionService } from './types.node'; @@ -146,5 +149,14 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole serviceManager.addSingleton(IJupyterRequestAgentCreator, RequestAgentCreator); serviceManager.addSingleton(ServerConnectionType, ServerConnectionType); serviceManager.addSingleton(JupyterConnection, JupyterConnection); - serviceManager.addBinding(JupyterConnection, IExtensionSingleActivationService); + serviceManager.addBinding(JupyterConnection, IExtensionSyncActivationService); + serviceManager.addSingleton( + ILiveRemoteKernelConnectionUsageTracker, + LiveRemoteKernelConnectionUsageTracker + ); + serviceManager.addBinding(ILiveRemoteKernelConnectionUsageTracker, IExtensionSyncActivationService); + serviceManager.addSingleton( + IExtensionSyncActivationService, + RemoteKernelConnectionHandler + ); } diff --git a/src/kernels/jupyter/serviceRegistry.web.ts b/src/kernels/jupyter/serviceRegistry.web.ts index aa15c9b3161..f9ed07d5c0b 100644 --- a/src/kernels/jupyter/serviceRegistry.web.ts +++ b/src/kernels/jupyter/serviceRegistry.web.ts @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import { IExtensionSingleActivationService } from '../../platform/activation/types'; +import { IExtensionSingleActivationService, IExtensionSyncActivationService } from '../../platform/activation/types'; import { IServiceManager } from '../../platform/ioc/types'; import { IRemoteKernelFinder } from '../raw/types'; import { INotebookProvider } from '../types'; @@ -19,6 +19,8 @@ import { NotebookProvider } from './launcher/notebookProvider'; import { NotebookServerProvider } from './launcher/notebookServerProvider'; import { ServerConnectionType } from './launcher/serverConnectionType'; import { JupyterServerUriStorage } from './launcher/serverUriStorage'; +import { LiveRemoteKernelConnectionUsageTracker } from './liveRemoteKernelConnectionTracker'; +import { RemoteKernelConnectionHandler } from './remoteKernelConnectionHandler'; import { RemoteKernelFinder } from './remoteKernelFinder'; import { JupyterServerSelector } from './serverSelector'; import { BackingFileCreator } from './session/backingFileCreator.web'; @@ -35,7 +37,8 @@ import { IJupyterServerProvider, IJupyterExecution, IJupyterRequestCreator, - INotebookServerFactory + INotebookServerFactory, + ILiveRemoteKernelConnectionUsageTracker } from './types'; export function registerTypes(serviceManager: IServiceManager, _isDevMode: boolean) { @@ -71,5 +74,14 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole serviceManager.addSingleton(IJupyterRequestCreator, JupyterRequestCreator); serviceManager.addSingleton(ServerConnectionType, ServerConnectionType); serviceManager.addSingleton(JupyterConnection, JupyterConnection); - serviceManager.addBinding(JupyterConnection, IExtensionSingleActivationService); + serviceManager.addBinding(JupyterConnection, IExtensionSyncActivationService); + serviceManager.addSingleton( + ILiveRemoteKernelConnectionUsageTracker, + LiveRemoteKernelConnectionUsageTracker + ); + serviceManager.addBinding(ILiveRemoteKernelConnectionUsageTracker, IExtensionSyncActivationService); + serviceManager.addSingleton( + IExtensionSyncActivationService, + RemoteKernelConnectionHandler + ); } diff --git a/src/kernels/jupyter/types.ts b/src/kernels/jupyter/types.ts index 940a735a96c..6388fd46c46 100644 --- a/src/kernels/jupyter/types.ts +++ b/src/kernels/jupyter/types.ts @@ -20,7 +20,8 @@ import { IJupyterKernelSpec, GetServerOptions, IKernelSocket, - KernelActionSource + KernelActionSource, + LiveRemoteKernelConnectionMetadata } from '../types'; import { ClassType } from '../../platform/ioc/types'; @@ -294,3 +295,19 @@ export interface INotebookStarter extends IDisposable { cancelToken: CancellationToken ): Promise; } + +export const ILiveRemoteKernelConnectionUsageTracker = Symbol('ILiveRemoteKernelConnectionUsageTracker'); +export interface ILiveRemoteKernelConnectionUsageTracker { + /** + * Whether the provided remote kernel was ever used by any notebook within the extension. + */ + wasKernelUsed(connection: LiveRemoteKernelConnectionMetadata): boolean; + /** + * Tracks the fact that the provided remote kernel for a given server was used by a notebook defined by the uri. + */ + trackKernelIdAsUsed(resource: Uri, serverId: string, kernelId: string): void; + /** + * Tracks the fact that the provided remote kernel for a given server is no longer used by a notebook defined by the uri. + */ + trackKernelIdAsNotUsed(resource: Uri, serverId: string, kernelId: string): void; +} diff --git a/src/kernels/kernelFinder.node.ts b/src/kernels/kernelFinder.node.ts index c6fb2d4367f..3a1a8a4b1c5 100644 --- a/src/kernels/kernelFinder.node.ts +++ b/src/kernels/kernelFinder.node.ts @@ -5,7 +5,7 @@ import { Memento } from 'vscode'; import { IFileSystem } from '../platform/common/platform/types.node'; import { GLOBAL_MEMENTO, IMemento } from '../platform/common/types'; import { ServerConnectionType } from './jupyter/launcher/serverConnectionType'; -import { IJupyterServerUriStorage } from './jupyter/types'; +import { IJupyterServerUriStorage, ILiveRemoteKernelConnectionUsageTracker } from './jupyter/types'; import { BaseKernelFinder } from './kernelFinder.base'; import { PreferredRemoteKernelIdProvider } from './jupyter/preferredRemoteKernelIdProvider'; import { ILocalKernelFinder, IRemoteKernelFinder } from './raw/types'; @@ -21,7 +21,9 @@ export class KernelFinder extends BaseKernelFinder { @inject(IMemento) @named(GLOBAL_MEMENTO) globalState: Memento, @inject(IFileSystem) private readonly fs: IFileSystem, @inject(IJupyterServerUriStorage) serverUriStorage: IJupyterServerUriStorage, - @inject(ServerConnectionType) serverConnectionType: ServerConnectionType + @inject(ServerConnectionType) serverConnectionType: ServerConnectionType, + @inject(ILiveRemoteKernelConnectionUsageTracker) + private readonly liveKernelConnectionTracker: ILiveRemoteKernelConnectionUsageTracker ) { super( preferredRemoteFinder, @@ -50,6 +52,15 @@ export class KernelFinder extends BaseKernelFinder { ? this.fs.localFileExists(kernel.interpreter.uri.fsPath) : Promise.resolve(true); }); + case 'startUsingRemoteKernelSpec': + // Always fetch the latest kernels from remotes, no need to display cached remote kernels. + return false; + case 'connectToLiveRemoteKernel': + // Only list live kernels that was used by the user, + // Even if such a kernel no longer exists on the sever. + // This way things don't just disappear from the list & + // user will get notified when they attempt to re-use this kernel. + return this.liveKernelConnectionTracker.wasKernelUsed(kernel); } return true; diff --git a/src/kernels/kernelFinder.web.ts b/src/kernels/kernelFinder.web.ts index a13a59b4499..08fe8fcda8e 100644 --- a/src/kernels/kernelFinder.web.ts +++ b/src/kernels/kernelFinder.web.ts @@ -4,7 +4,7 @@ import { injectable, inject, named } from 'inversify'; import { Memento } from 'vscode'; import { GLOBAL_MEMENTO, IMemento } from '../platform/common/types'; import { ServerConnectionType } from './jupyter/launcher/serverConnectionType'; -import { IJupyterServerUriStorage } from './jupyter/types'; +import { IJupyterServerUriStorage, ILiveRemoteKernelConnectionUsageTracker } from './jupyter/types'; import { BaseKernelFinder } from './kernelFinder.base'; import { PreferredRemoteKernelIdProvider } from './jupyter/preferredRemoteKernelIdProvider'; import { IRemoteKernelFinder } from './raw/types'; @@ -18,7 +18,9 @@ export class KernelFinder extends BaseKernelFinder { @inject(INotebookProvider) notebookProvider: INotebookProvider, @inject(IMemento) @named(GLOBAL_MEMENTO) globalState: Memento, @inject(IJupyterServerUriStorage) serverUriStorage: IJupyterServerUriStorage, - @inject(ServerConnectionType) serverConnectionType: ServerConnectionType + @inject(ServerConnectionType) serverConnectionType: ServerConnectionType, + @inject(ILiveRemoteKernelConnectionUsageTracker) + private readonly liveKernelConnectionTracker: ILiveRemoteKernelConnectionUsageTracker ) { super( preferredRemoteFinder, @@ -35,6 +37,15 @@ export class KernelFinder extends BaseKernelFinder { case 'startUsingPythonInterpreter': case 'startUsingLocalKernelSpec': return false; + case 'startUsingRemoteKernelSpec': + // Always fetch the latest kernels from remotes, no need to display cached remote kernels. + return false; + case 'connectToLiveRemoteKernel': + // Only list live kernels that was used by the user, + // Even if such a kernel no longer exists on the sever. + // This way things don't just disappear from the list & + // user will get notified when they attempt to re-use this kernel. + return this.liveKernelConnectionTracker.wasKernelUsed(kernel); } return true; diff --git a/src/kernels/types.ts b/src/kernels/types.ts index 28b96c2d2ec..1005b709659 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -124,7 +124,7 @@ export function isLocalConnection( export interface IKernel extends IAsyncDisposable { readonly id: Uri; - /**; + /** * In the case of Notebooks, this is the same as the Notebook Uri. * But in the case of Interactive Window, this is the Uri of the file (such as the Python file). * However if we create an intearctive window without a file, then this is undefined. diff --git a/src/notebooks/controllers/notebookControllerManager.ts b/src/notebooks/controllers/notebookControllerManager.ts index ce84c02dad6..22b9eaa121d 100644 --- a/src/notebooks/controllers/notebookControllerManager.ts +++ b/src/notebooks/controllers/notebookControllerManager.ts @@ -42,7 +42,6 @@ import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { sendTelemetryEvent } from '../../telemetry'; import { Telemetry } from '../../webviews/webview-side/common/constants'; import { NotebookCellLanguageService } from '../../intellisense/cellLanguageService'; -import { PreferredRemoteKernelIdProvider } from '../../kernels/jupyter/preferredRemoteKernelIdProvider'; import { LiveRemoteKernelConnectionMetadata, IKernelProvider, @@ -96,7 +95,11 @@ export class NotebookControllerManager implements INotebookControllerManager, IE notebook: NotebookDocument; controller: VSCodeNotebookController; }>; - private readonly _onNotebookControllerSelectionChanged = new EventEmitter(); + private readonly _onNotebookControllerSelectionChanged = new EventEmitter<{ + notebook: NotebookDocument; + controller: IVSCodeNotebookController; + selected: boolean; + }>(); // Promise to resolve when we have loaded our controllers private controllersPromise?: Promise; @@ -137,8 +140,6 @@ export class NotebookControllerManager implements INotebookControllerManager, IE @inject(ICommandManager) private readonly commandManager: ICommandManager, @inject(IExtensionContext) private readonly context: IExtensionContext, @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, - @inject(PreferredRemoteKernelIdProvider) - private readonly preferredRemoteKernelIdProvider: PreferredRemoteKernelIdProvider, @inject(NotebookCellLanguageService) private readonly languageService: NotebookCellLanguageService, @inject(IWorkspaceService) private readonly workspace: IWorkspaceService, @inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker, @@ -347,11 +348,9 @@ export class NotebookControllerManager implements INotebookControllerManager, IE // Unless the user switches to using local kernels (i.e. doesn't have a remote kernel setup). if ( connectionIsNoLongerValid && - (controller.connection.kind === 'connectToLiveRemoteKernel' || - controller.connection.kind === 'startUsingRemoteKernelSpec') && + controller.connection.kind === 'connectToLiveRemoteKernel' && !this.isLocalLaunch ) { - controller.flagRemoteKernelAsOutdated(); return true; } return connectionIsNoLongerValid; @@ -858,7 +857,6 @@ export class NotebookControllerManager implements INotebookControllerManager, IE this.notebook, this.commandManager, this.kernelProvider, - this.preferredRemoteKernelIdProvider, this.context, this.disposables, this.languageService, @@ -877,7 +875,7 @@ export class NotebookControllerManager implements INotebookControllerManager, IE this.disposables ); controller.onNotebookControllerSelectionChanged( - () => this._onNotebookControllerSelectionChanged.fire(), + (e) => this._onNotebookControllerSelectionChanged.fire({ ...e, controller }), this, this.disposables ); diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 1444987a22e..38683c0b1f9 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -67,13 +67,11 @@ import { areKernelConnectionsEqual, getKernelRegistrationInfo } from '../../kernels/helpers'; -import { PreferredRemoteKernelIdProvider } from '../../kernels/jupyter/preferredRemoteKernelIdProvider'; import { IKernel, IKernelProvider, isLocalConnection, KernelConnectionMetadata, - KernelSocketInformation, LiveRemoteKernelConnectionMetadata, LocalKernelSpecConnectionMetadata, PythonKernelConnectionMetadata @@ -86,19 +84,21 @@ import { DisplayOptions } from '../../kernels/displayOptions'; import { sendNotebookOrKernelLanguageTelemetry } from '../../platform/common/utils'; import { ConsoleForegroundColors, TraceOptions } from '../../platform/logging/types'; import { KernelConnector } from '../../kernels/kernelConnector'; +import { IVSCodeNotebookController } from './types'; -export class VSCodeNotebookController implements Disposable { +export class VSCodeNotebookController implements Disposable, IVSCodeNotebookController { private readonly _onNotebookControllerSelected: EventEmitter<{ notebook: NotebookDocument; controller: VSCodeNotebookController; }>; - private readonly _onNotebookControllerSelectionChanged = new EventEmitter(); + private readonly _onNotebookControllerSelectionChanged = new EventEmitter<{ + selected: boolean; + notebook: NotebookDocument; + }>(); private readonly _onDidDispose = new EventEmitter(); private readonly disposables: IDisposable[] = []; private notebookKernels = new WeakMap(); public readonly controller: NotebookController; - private isConnectionOutdated?: boolean; - private outdatedMessageDisplayed?: boolean; /** * Used purely for testing purposes. */ @@ -142,7 +142,6 @@ export class VSCodeNotebookController implements Disposable { private readonly notebookApi: IVSCodeNotebook, private readonly commandManager: ICommandManager, private readonly kernelProvider: IKernelProvider, - private readonly preferredRemoteKernelIdProvider: PreferredRemoteKernelIdProvider, private readonly context: IExtensionContext, disposableRegistry: IDisposableRegistry, private readonly languageService: NotebookCellLanguageService, @@ -185,9 +184,6 @@ export class VSCodeNotebookController implements Disposable { public updateRemoteKernelDetails(kernelConnection: LiveRemoteKernelConnectionMetadata) { this.controller.detail = getRemoteKernelSessionInformation(kernelConnection); } - public flagRemoteKernelAsOutdated() { - this.isConnectionOutdated = true; - } public updateInterpreterDetails( kernelConnection: LocalKernelSpecConnectionMetadata | PythonKernelConnectionMetadata ) { @@ -236,11 +232,6 @@ export class VSCodeNotebookController implements Disposable { if (cells.length < 1) { return; } - if (this.isConnectionOutdated && !this.outdatedMessageDisplayed) { - void this.appShell.showErrorMessage(DataScience.jupyterRemoteConnectFailedModalMessage(), { modal: true }); - this.outdatedMessageDisplayed = true; - return; - } // Found on CI that sometimes VS Code calls this with old deleted cells. // See here https://github.com/microsoft/vscode-jupyter/runs/5581627878?check_suite_focus=true cells = cells.filter((cell) => { @@ -313,7 +304,8 @@ export class VSCodeNotebookController implements Disposable { void kernel.dispose(); } this.associatedDocuments.delete(event.notebook); - this._onNotebookControllerSelectionChanged.fire(); + this._onNotebookControllerSelectionChanged.fire(event); + return; } // We're only interested in our Notebooks. @@ -331,7 +323,7 @@ export class VSCodeNotebookController implements Disposable { // If this NotebookController was selected, fire off the event this._onNotebookControllerSelected.fire({ notebook: event.notebook, controller: this }); - this._onNotebookControllerSelectionChanged.fire(); + this._onNotebookControllerSelectionChanged.fire(event); traceVerbose(`Controller selection change completed`); deferred.resolve(); } @@ -522,7 +514,6 @@ export class VSCodeNotebookController implements Disposable { return; } this.notebookKernels.set(doc, kernel); - let kernelSocket: KernelSocketInformation | undefined; const handlerDisposables: IDisposable[] = []; // If the notebook is closed, dispose everything. this.notebookApi.onDidCloseNotebookDocument( @@ -534,20 +525,7 @@ export class VSCodeNotebookController implements Disposable { this, handlerDisposables ); - const saveKernelInfo = () => { - const kernelId = kernelSocket?.options.id; - if (!kernelId || isLocalConnection(this.kernelConnection)) { - return; - } - traceInfo(`Updating preferred kernel for remote notebook ${kernelId}`); - this.preferredRemoteKernelIdProvider.storePreferredRemoteKernelId(doc.uri, kernelId).catch(noop); - }; - const kernelDisposedDisposable = kernel.onDisposed(() => disposeAllDisposables(handlerDisposables)); - const subscriptionDisposables = kernel.kernelSocket.subscribe((item) => { - kernelSocket = item; - saveKernelInfo(); - }); const statusChangeDisposable = kernel.onStatusChanged(async () => { if (kernel.disposed || !kernel.info) { return; @@ -563,21 +541,11 @@ export class VSCodeNotebookController implements Disposable { kernel.kernelConnectionMetadata, kernel.info ); - if ( - this.kernelConnection.kind === 'startUsingLocalKernelSpec' || - this.kernelConnection.kind === 'startUsingRemoteKernelSpec' - ) { - if (kernel.info.status === 'ok') { - saveKernelInfo(); - } else { - disposeAllDisposables(handlerDisposables); - } - } else { + if (kernel.info.status === 'ok') { disposeAllDisposables(handlerDisposables); } }); - handlerDisposables.push({ dispose: () => subscriptionDisposables.unsubscribe() }); handlerDisposables.push({ dispose: () => statusChangeDisposable.dispose() }); handlerDisposables.push({ dispose: () => kernelDisposedDisposable?.dispose() }); } @@ -632,9 +600,6 @@ export class VSCodeNotebookController implements Disposable { // Before we start the notebook, make sure the metadata is set to this new kernel. await updateNotebookDocumentMetadata(document, this.documentManager, selectedKernelConnectionMetadata); - if (isLocalConnection(this.connection)) { - this.preferredRemoteKernelIdProvider.clearPreferredRemoteKernelId(document.uri).catch(noop); - } if (document.notebookType === InteractiveWindowView) { // Possible its an interactive window, in that case we'll create the kernel manually. return; diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 5f1bf95f232..376489090ca 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -21,7 +21,11 @@ export const INotebookKernelResolver = Symbol('INotebookKernelResolver'); export const INotebookControllerManager = Symbol('INotebookControllerManager'); export interface INotebookControllerManager { readonly onNotebookControllerSelected: Event<{ notebook: NotebookDocument; controller: IVSCodeNotebookController }>; - readonly onNotebookControllerSelectionChanged: Event; + readonly onNotebookControllerSelectionChanged: Event<{ + notebook: NotebookDocument; + controller: IVSCodeNotebookController; + selected: boolean; + }>; readonly kernelConnections: Promise[]>; readonly remoteRefreshed: Event; /** diff --git a/src/platform/common/utils/localize.ts b/src/platform/common/utils/localize.ts index 68ece7e4b24..99fe54b65c7 100644 --- a/src/platform/common/utils/localize.ts +++ b/src/platform/common/utils/localize.ts @@ -582,10 +582,6 @@ export namespace DataScience { 'Manage Connections' ); - export const jupyterRemoteConnectFailedModalMessage = localize( - 'DataScience.jupyterRemoteConnectFailedModalMessage', - 'Failed to connect to the remote Jupyter Server. View Jupyter log for further details.' - ); export const changeJupyterRemoteConnection = localize( 'DataScience.changeJupyterRemoteConnection', 'Change Jupyter Server connection.' diff --git a/src/test/datascience/kernel-launcher/localKernelFinder.unit.test.ts b/src/test/datascience/kernel-launcher/localKernelFinder.unit.test.ts index b83d5a92e1f..7d8413e8da4 100644 --- a/src/test/datascience/kernel-launcher/localKernelFinder.unit.test.ts +++ b/src/test/datascience/kernel-launcher/localKernelFinder.unit.test.ts @@ -54,6 +54,7 @@ import { NotebookProvider } from '../../../kernels/jupyter/launcher/notebookProv import { RemoteKernelFinder } from '../../../kernels/jupyter/remoteKernelFinder'; import { JupyterServerUriStorage } from '../../../kernels/jupyter/launcher/serverUriStorage'; import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverConnectionType'; +import { ILiveRemoteKernelConnectionUsageTracker } from '../../../kernels/jupyter/types'; [false, true].forEach((isWindows) => { suite(`Local Kernel Finder ${isWindows ? 'Windows' : 'Unix'}`, () => { @@ -69,6 +70,7 @@ import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverCo let tempDirForKernelSpecs: Uri; let jupyterPaths: JupyterPaths; let preferredRemote: PreferredRemoteKernelIdProvider; + let liveKernelUsageTracker: ILiveRemoteKernelConnectionUsageTracker; type TestData = { interpreters?: ( | PythonEnvironment @@ -230,7 +232,7 @@ import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverCo instance(memento) ) ); - + liveKernelUsageTracker = mock(); preferredRemote = mock(PreferredRemoteKernelIdProvider); const notebookProvider = mock(NotebookProvider); const serverUriStorage = mock(JupyterServerUriStorage); @@ -240,6 +242,7 @@ import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverCo const onDidChangeEvent = new EventEmitter(); disposables.push(onDidChangeEvent); when(connectionType.onDidChange).thenReturn(onDidChangeEvent.event); + kernelFinder = new KernelFinder( localKernelFinder, instance(remoteKernelFinder), @@ -248,7 +251,8 @@ import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverCo instance(memento), instance(fs), instance(serverUriStorage), - instance(connectionType) + instance(connectionType), + instance(liveKernelUsageTracker) ); } teardown(() => { diff --git a/src/test/datascience/kernel-launcher/remoteKernelFinder.unit.test.ts b/src/test/datascience/kernel-launcher/remoteKernelFinder.unit.test.ts index 03548f8955f..5adc3f3f37e 100644 --- a/src/test/datascience/kernel-launcher/remoteKernelFinder.unit.test.ts +++ b/src/test/datascience/kernel-launcher/remoteKernelFinder.unit.test.ts @@ -5,17 +5,17 @@ import type { Session } from '@jupyterlab/services'; import { assert } from 'chai'; -import { anything, instance, mock, when } from 'ts-mockito'; +import { anything, instance, mock, when, verify } from 'ts-mockito'; import { getDisplayNameOrNameOfKernelConnection } from '../../../kernels/helpers'; import { PYTHON_LANGUAGE } from '../../../platform/common/constants'; -import { Disposable, EventEmitter, Uri } from 'vscode'; -import { MockMemento } from '../../mocks/mementos'; +import { Disposable, EventEmitter, Memento, Uri } from 'vscode'; import { CryptoUtils } from '../../../platform/common/crypto'; import { noop } from '../../core'; import { IJupyterConnection, IJupyterKernelSpec, IKernelFinder, + KernelConnectionMetadata, LiveRemoteKernelConnectionMetadata } from '../../../kernels/types'; import { IInterpreterService } from '../../../platform/interpreter/contracts'; @@ -23,8 +23,15 @@ import { JupyterSessionManager } from '../../../kernels/jupyter/session/jupyterS import { JupyterSessionManagerFactory } from '../../../kernels/jupyter/session/jupyterSessionManagerFactory'; import { RemoteKernelFinder } from '../../../kernels/jupyter/remoteKernelFinder'; import { ILocalKernelFinder, IRemoteKernelFinder } from '../../../kernels/raw/types'; -import { PreferredRemoteKernelIdProvider } from '../../../kernels/jupyter/preferredRemoteKernelIdProvider'; -import { IJupyterKernel, IJupyterSessionManager } from '../../../kernels/jupyter/types'; +import { + ActiveKernelIdList, + PreferredRemoteKernelIdProvider +} from '../../../kernels/jupyter/preferredRemoteKernelIdProvider'; +import { + IJupyterKernel, + IJupyterSessionManager, + ILiveRemoteKernelConnectionUsageTracker +} from '../../../kernels/jupyter/types'; import { KernelFinder } from '../../../kernels/kernelFinder.node'; import { NotebookProvider } from '../../../kernels/jupyter/launcher/notebookProvider'; import { PythonExtensionChecker } from '../../../platform/api/pythonApi'; @@ -34,6 +41,7 @@ import { JupyterServerUriStorage } from '../../../kernels/jupyter/launcher/serve import { FileSystem } from '../../../platform/common/platform/fileSystem.node'; import { takeTopRankKernel } from './localKernelFinder.unit.test'; import { ServerConnectionType } from '../../../kernels/jupyter/launcher/serverConnectionType'; +import { LocalKernelSpecsCacheKey, RemoteKernelSpecsCacheKey } from '../../../kernels/kernelFinder.base'; suite(`Remote Kernel Finder`, () => { let disposables: Disposable[] = []; @@ -42,10 +50,11 @@ suite(`Remote Kernel Finder`, () => { let localKernelFinder: ILocalKernelFinder; let kernelFinder: IKernelFinder; let fs: IFileSystem; - let memento: MockMemento; + let memento: Memento; let jupyterSessionManager: IJupyterSessionManager; const dummyEvent = new EventEmitter(); let interpreterService: IInterpreterService; + let liveKernelUsageTracker: ILiveRemoteKernelConnectionUsageTracker; const connInfo: IJupyterConnection = { url: 'http://foobar', type: 'jupyter', @@ -110,11 +119,13 @@ suite(`Remote Kernel Finder`, () => { }); setup(() => { + memento = mock(); + when(memento.get(ActiveKernelIdList, anything())).thenReturn([]); const crypto = mock(CryptoUtils); when(crypto.createHash(anything(), anything())).thenCall((d, _c) => { return d.toLowerCase(); }); - preferredRemoteKernelIdProvider = new PreferredRemoteKernelIdProvider(new MockMemento(), instance(crypto)); + preferredRemoteKernelIdProvider = new PreferredRemoteKernelIdProvider(instance(memento), instance(crypto)); jupyterSessionManager = mock(JupyterSessionManager); const jupyterSessionManagerFactory = mock(JupyterSessionManagerFactory); when(jupyterSessionManagerFactory.create(anything())).thenResolve(instance(jupyterSessionManager)); @@ -139,23 +150,24 @@ suite(`Remote Kernel Finder`, () => { const serverUriStorage = mock(JupyterServerUriStorage); when(serverUriStorage.getUri()).thenResolve(connInfo.baseUrl); when(serverUriStorage.getRemoteUri()).thenResolve(connInfo.baseUrl); - memento = new MockMemento(); const connectionType = mock(); when(connectionType.isLocalLaunch).thenReturn(false); when(connectionType.setIsLocalLaunch(anything())).thenResolve(); const onDidChangeEvent = new EventEmitter(); disposables.push(onDidChangeEvent); when(connectionType.onDidChange).thenReturn(onDidChangeEvent.event); - + liveKernelUsageTracker = mock(); + when(liveKernelUsageTracker.wasKernelUsed(anything())).thenReturn(true); kernelFinder = new KernelFinder( instance(localKernelFinder), remoteKernelFinder, preferredRemoteKernelIdProvider, instance(notebookProvider), - memento, + instance(memento), instance(fs), instance(serverUriStorage), - instance(connectionType) + instance(connectionType), + instance(liveKernelUsageTracker) ); }); teardown(() => { @@ -245,6 +257,14 @@ suite(`Remote Kernel Finder`, () => { juliaSpec, interpreterSpec ]); + let activeKernelIdList: unknown = []; + when(memento.update(anything(), anything())).thenCall((key, value) => { + if (key === ActiveKernelIdList) { + activeKernelIdList = value as any; + } + return Promise.resolve(); + }); + when(memento.get(ActiveKernelIdList, anything())).thenCall(() => activeKernelIdList); const uri = Uri.file('/usr/foobar/foo.ipynb'); await preferredRemoteKernelIdProvider.storePreferredRemoteKernelId(uri, '2'); @@ -256,5 +276,107 @@ suite(`Remote Kernel Finder`, () => { python3Kernels[1].name, 'Wrong live kernel returned' ); + + verify(memento.update(ActiveKernelIdList, anything())).once(); + }); + test('Do not return cached remote kernelspecs or live kernels', async () => { + const liveRemoteKernel: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'baseUrl1', + id: '1', + kernelModel: { + lastActivityTime: new Date(), + model: { + id: '1', + name: '', + path: '', + type: '', + kernel: { + id: '1', + name: '' + } + }, + name: '', + numberOfConnections: 0 + }, + kind: 'connectToLiveRemoteKernel', + serverId: 'serverId1' + }; + const cachedKernels: KernelConnectionMetadata[] = [ + { + baseUrl: 'baseUrl1', + id: '1', + kernelSpec: { + argv: [], + display_name: '', + name: '', + uri: Uri.file('') + }, + kind: 'startUsingRemoteKernelSpec', + serverId: 'serverId1' + }, + liveRemoteKernel + ]; + when(liveKernelUsageTracker.wasKernelUsed(anything())).thenReturn(false); + when(memento.get(LocalKernelSpecsCacheKey, anything())).thenReturn([]); + when(memento.get(RemoteKernelSpecsCacheKey, anything())).thenReturn(cachedKernels); + when(jupyterSessionManager.getRunningKernels()).thenResolve([]); + when(jupyterSessionManager.getRunningSessions()).thenResolve([]); + when(jupyterSessionManager.getKernelSpecs()).thenResolve([]); + + const kernels = await kernelFinder.listKernels(Uri.file('a.ipynb'), undefined, 'useCache'); + assert.lengthOf(kernels, 0); + + verify(liveKernelUsageTracker.wasKernelUsed(liveRemoteKernel)).once(); + }); + test('Return cached remote live kernel if used', async () => { + const liveRemoteKernel: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'baseUrl1', + id: '1', + kernelModel: { + lastActivityTime: new Date(), + model: { + id: '1', + name: '', + path: '', + type: '', + kernel: { + id: '1', + name: '' + } + }, + name: '', + numberOfConnections: 0 + }, + kind: 'connectToLiveRemoteKernel', + serverId: 'serverId1' + }; + const cachedKernels: KernelConnectionMetadata[] = [ + { + baseUrl: 'baseUrl1', + id: '1', + kernelSpec: { + argv: [], + display_name: '', + name: '', + uri: Uri.file('') + }, + kind: 'startUsingRemoteKernelSpec', + serverId: 'serverId1' + }, + liveRemoteKernel + ]; + when(liveKernelUsageTracker.wasKernelUsed(anything())).thenReturn(false); + when(liveKernelUsageTracker.wasKernelUsed(liveRemoteKernel)).thenReturn(true); + when(memento.get(LocalKernelSpecsCacheKey, anything())).thenReturn([]); + when(memento.get(RemoteKernelSpecsCacheKey, anything())).thenReturn(cachedKernels); + when(jupyterSessionManager.getRunningKernels()).thenResolve([]); + when(jupyterSessionManager.getRunningSessions()).thenResolve([]); + when(jupyterSessionManager.getKernelSpecs()).thenResolve([]); + + const kernels = await kernelFinder.listKernels(Uri.file('a.ipynb'), undefined, 'useCache'); + assert.lengthOf(kernels, 1); + assert.deepEqual(kernels, [liveRemoteKernel]); + + verify(liveKernelUsageTracker.wasKernelUsed(liveRemoteKernel)).once(); }); }); diff --git a/src/test/kernels/jupyter/liveRemoteKernelConnectionTracker.unit.test.ts b/src/test/kernels/jupyter/liveRemoteKernelConnectionTracker.unit.test.ts new file mode 100644 index 00000000000..283c35524cb --- /dev/null +++ b/src/test/kernels/jupyter/liveRemoteKernelConnectionTracker.unit.test.ts @@ -0,0 +1,225 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { assert, use } from 'chai'; + +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { EventEmitter, Memento, Uri } from 'vscode'; +import { IJupyterServerUriStorage } from '../../../kernels/jupyter/types'; +import { disposeAllDisposables } from '../../../platform/common/helpers'; +import { IDisposable } from '../../../platform/common/types'; +import * as chaiAsPromised from 'chai-as-promised'; +import { + LiveRemoteKernelConnectionUsageTracker, + mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources +} from '../../../kernels/jupyter/liveRemoteKernelConnectionTracker'; +import { LiveRemoteKernelConnectionMetadata } from '../../../kernels/types'; +import { computeServerId } from '../../../kernels/jupyter/jupyterUtils'; + +use(chaiAsPromised); +suite('Live kernel Connection Tracker', async () => { + let serverUriStorage: IJupyterServerUriStorage; + let memento: Memento; + let tracker: LiveRemoteKernelConnectionUsageTracker; + let onDidRemoveUri: EventEmitter; + const disposables: IDisposable[] = []; + const server2Uri = 'http://one:1234/hello?token=1234'; + const remoteLiveKernel1: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'baseUrl', + id: 'connectionId', + kind: 'connectToLiveRemoteKernel', + serverId: 'server1', + kernelModel: { + lastActivityTime: new Date(), + id: 'model1', + model: { + id: 'modelId', + kernel: { + id: 'kernelId', + name: 'kernelName' + }, + name: 'modelName', + path: '', + type: '' + }, + name: '', + numberOfConnections: 0 + } + }; + const remoteLiveKernel2: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'http://one:1234/', + id: 'connectionId2', + kind: 'connectToLiveRemoteKernel', + serverId: computeServerId(server2Uri), + kernelModel: { + id: 'modelId2', + lastActivityTime: new Date(), + model: { + id: 'modelId2', + kernel: { + id: 'kernelI2', + name: 'kernelName2' + }, + name: 'modelName2', + path: '', + type: '' + }, + name: '', + numberOfConnections: 0 + } + }; + const remoteLiveKernel3: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'http://one:1234/', + id: 'connectionId3', + kind: 'connectToLiveRemoteKernel', + serverId: computeServerId(server2Uri), + kernelModel: { + lastActivityTime: new Date(), + id: 'modelId3', + model: { + id: 'modelId3', + kernel: { + id: 'kernelI2', + name: 'kernelName2' + }, + name: 'modelName2', + path: '', + type: '' + }, + name: '', + numberOfConnections: 0 + } + }; + setup(() => { + serverUriStorage = mock(); + memento = mock(); + onDidRemoveUri = new EventEmitter(); + disposables.push(onDidRemoveUri); + when(serverUriStorage.onDidRemoveUri).thenReturn(onDidRemoveUri.event); + tracker = new LiveRemoteKernelConnectionUsageTracker( + disposables, + instance(serverUriStorage), + instance(memento) + ); + }); + teardown(() => { + disposeAllDisposables(disposables); + }); + + test('Ensure event handler is added', () => { + tracker.activate(); + verify(serverUriStorage.onDidRemoveUri).once(); + }); + test('Kernel connection is not used if memento is empty', async () => { + when(memento.get(anything(), anything())).thenCall((_, defaultValue) => defaultValue); + + tracker.activate(); + + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel1)); + }); + test('Kernel connection is not used if memento is not empty but does not contain the same connection info', async () => { + const cachedItems = { + [remoteLiveKernel2.serverId]: { + [remoteLiveKernel2.kernelModel.id!]: [Uri.file('a.ipynb').toString()] + } + }; + when(memento.get(mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, anything())).thenCall( + () => cachedItems + ); + + tracker.activate(); + + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel1)); + }); + test('Kernel connection is used if connection is tracked in memento', async () => { + const cachedItems = { + [remoteLiveKernel2.serverId]: { + [remoteLiveKernel2.kernelModel.id!]: [Uri.file('a.ipynb').toString()] + }, + [remoteLiveKernel1.serverId]: { + [remoteLiveKernel1.kernelModel.id!]: [Uri.file('a.ipynb').toString()] + } + }; + when(memento.get(mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, anything())).thenCall( + () => cachedItems + ); + + tracker.activate(); + + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel1)); + }); + test('Memento is updated to track usage of a kernel connection', async () => { + const cachedItems: any = {}; + when(memento.get(mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, anything())).thenReturn( + cachedItems + ); + when(memento.update(mementoKeyToTrackRemoveKernelUrisAndSessionsUsedByResources, anything())).thenCall( + (_, value) => { + Object.assign(cachedItems, value); + return Promise.resolve(); + } + ); + + tracker.activate(); + tracker.trackKernelIdAsUsed(Uri.file('a.ipynb'), remoteLiveKernel1.serverId, remoteLiveKernel1.kernelModel.id!); + + assert.deepEqual(cachedItems[remoteLiveKernel1.serverId][remoteLiveKernel1.kernelModel.id!], [ + Uri.file('a.ipynb').toString() + ]); + + tracker.trackKernelIdAsUsed(Uri.file('a.ipynb'), remoteLiveKernel2.serverId, remoteLiveKernel2.kernelModel.id!); + + assert.deepEqual(cachedItems[remoteLiveKernel2.serverId][remoteLiveKernel2.kernelModel.id!], [ + Uri.file('a.ipynb').toString() + ]); + + tracker.trackKernelIdAsUsed(Uri.file('a.ipynb'), remoteLiveKernel3.serverId, remoteLiveKernel3.kernelModel.id!); + + assert.deepEqual(cachedItems[remoteLiveKernel3.serverId][remoteLiveKernel3.kernelModel.id!], [ + Uri.file('a.ipynb').toString() + ]); + + tracker.trackKernelIdAsUsed(Uri.file('b.ipynb'), remoteLiveKernel3.serverId, remoteLiveKernel3.kernelModel.id!); + + assert.deepEqual(cachedItems[remoteLiveKernel3.serverId][remoteLiveKernel3.kernelModel.id!], [ + Uri.file('a.ipynb').toString(), + Uri.file('b.ipynb').toString() + ]); + + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel1)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel2)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel3)); + + // Remove a kernel connection from some other document. + tracker.trackKernelIdAsNotUsed( + Uri.file('xyz.ipynb'), + remoteLiveKernel1.serverId, + remoteLiveKernel1.kernelModel.id! + ); + + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel1)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel2)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel3)); + + // Remove a kernel connection from a tracked document. + tracker.trackKernelIdAsNotUsed( + Uri.file('a.ipynb'), + remoteLiveKernel1.serverId, + remoteLiveKernel1.kernelModel.id! + ); + + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel1)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel2)); + assert.isTrue(tracker.wasKernelUsed(remoteLiveKernel3)); + + // Forget the Uir connection all together. + onDidRemoveUri.fire(server2Uri); + + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel1)); + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel2)); + assert.isFalse(tracker.wasKernelUsed(remoteLiveKernel3)); + + assert.isEmpty(cachedItems); + }); +}); diff --git a/src/test/kernels/jupyter/remoteKernelConnectionHandler.unit.test.ts b/src/test/kernels/jupyter/remoteKernelConnectionHandler.unit.test.ts new file mode 100644 index 00000000000..893e4172373 --- /dev/null +++ b/src/test/kernels/jupyter/remoteKernelConnectionHandler.unit.test.ts @@ -0,0 +1,231 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { use } from 'chai'; + +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode'; +import { ILiveRemoteKernelConnectionUsageTracker } from '../../../kernels/jupyter/types'; +import { disposeAllDisposables } from '../../../platform/common/helpers'; +import { IDisposable } from '../../../platform/common/types'; +import * as chaiAsPromised from 'chai-as-promised'; +import { + IKernel, + IKernelProvider, + isLocalConnection, + KernelActionSource, + KernelConnectionMetadata, + KernelSocketInformation, + LiveRemoteKernelConnectionMetadata, + LocalKernelConnectionMetadata, + RemoteKernelSpecConnectionMetadata +} from '../../../kernels/types'; +import { PreferredRemoteKernelIdProvider } from '../../../kernels/jupyter/preferredRemoteKernelIdProvider'; +import { RemoteKernelConnectionHandler } from '../../../kernels/jupyter/remoteKernelConnectionHandler'; +import { INotebookControllerManager } from '../../../notebooks/types'; +import { Subject } from 'rxjs'; +import { IVSCodeNotebookController } from '../../../notebooks/controllers/types'; + +use(chaiAsPromised); +suite('Remote kernel connection handler', async () => { + let tracker: ILiveRemoteKernelConnectionUsageTracker; + let preferredRemoteKernelProvider: PreferredRemoteKernelIdProvider; + let onDidStartKernel: EventEmitter; + let onNotebookControllerSelectionChanged: EventEmitter<{ + selected: boolean; + notebook: NotebookDocument; + controller: IVSCodeNotebookController; + }>; + let remoteConnectionHandler: RemoteKernelConnectionHandler; + let controllers: INotebookControllerManager; + let kernelProvider: IKernelProvider; + const disposables: IDisposable[] = []; + // const server2Uri = 'http://one:1234/hello?token=1234'; + const remoteKernelSpec: RemoteKernelSpecConnectionMetadata = { + baseUrl: 'baseUrl', + id: 'remoteKernelSpec1', + kind: 'startUsingRemoteKernelSpec', + serverId: 'server1', + kernelSpec: { + argv: [], + display_name: '', + name: '', + uri: Uri.file('') + } + }; + const localKernelSpec: LocalKernelConnectionMetadata = { + id: 'localKernelSpec1', + kernelSpec: { + argv: [], + display_name: '', + name: '', + uri: Uri.file('') + }, + kind: 'startUsingLocalKernelSpec' + }; + const remoteLiveKernel1: LiveRemoteKernelConnectionMetadata = { + baseUrl: 'baseUrl', + id: 'connectionId', + kind: 'connectToLiveRemoteKernel', + serverId: 'server1', + kernelModel: { + lastActivityTime: new Date(), + id: 'model1', + model: { + id: 'modelId', + kernel: { + id: 'kernelId', + name: 'kernelName' + }, + name: 'modelName', + path: '', + type: '' + }, + name: '', + numberOfConnections: 0 + } + }; + setup(() => { + onDidStartKernel = new EventEmitter(); + kernelProvider = mock(); + controllers = mock(); + tracker = mock(); + preferredRemoteKernelProvider = mock(); + onNotebookControllerSelectionChanged = new EventEmitter<{ + selected: boolean; + notebook: NotebookDocument; + controller: IVSCodeNotebookController; + }>(); + + disposables.push(onDidStartKernel); + disposables.push(onNotebookControllerSelectionChanged); + + when(kernelProvider.onDidStartKernel).thenReturn(onDidStartKernel.event); + when(controllers.onNotebookControllerSelectionChanged).thenReturn(onNotebookControllerSelectionChanged.event); + when(preferredRemoteKernelProvider.storePreferredRemoteKernelId(anything(), anything())).thenResolve(); + when(preferredRemoteKernelProvider.clearPreferredRemoteKernelId(anything())).thenResolve(); + + remoteConnectionHandler = new RemoteKernelConnectionHandler( + disposables, + instance(kernelProvider), + instance(controllers), + instance(tracker), + instance(preferredRemoteKernelProvider) + ); + }); + teardown(() => { + disposeAllDisposables(disposables); + }); + + test('Ensure event handler is added', () => { + remoteConnectionHandler.activate(); + verify(kernelProvider.onDidStartKernel).once(); + }); + function verifyRemoteKernelTracking(connection: KernelConnectionMetadata, source: KernelActionSource) { + const kernel1 = mock(); + when(kernel1.kernelConnectionMetadata).thenReturn(connection); + when(kernel1.creator).thenReturn(source); + const subject = new Subject(); + disposables.push(new Disposable(() => subject.unsubscribe())); + when(kernel1.kernelSocket).thenReturn(subject); + const nbUri = Uri.file('a.ipynb'); + when(kernel1.resourceUri).thenReturn(nbUri); + when(kernel1.disposed).thenReturn(false); + when(kernel1.disposing).thenReturn(false); + + remoteConnectionHandler.activate(); + onDidStartKernel.fire(instance(kernel1)); + + verify(tracker.trackKernelIdAsUsed(anything(), anything(), anything())).never(); + verify(preferredRemoteKernelProvider.storePreferredRemoteKernelId(anything(), anything())).never(); + + const kernelInfo: KernelSocketInformation = { + options: { + clientId: '', + id: 'modelId1', + model: { + id: 'modelId1', + name: '' + }, + userName: '' + } + }; + subject.next(kernelInfo); + + if (connection.kind === 'startUsingRemoteKernelSpec' && source === 'jupyterExtension') { + verify(tracker.trackKernelIdAsUsed(nbUri, remoteKernelSpec.serverId, kernelInfo.options.id)).once(); + verify(preferredRemoteKernelProvider.storePreferredRemoteKernelId(nbUri, kernelInfo.options.id)).once(); + } else { + verify(tracker.trackKernelIdAsUsed(anything(), anything(), anything())).never(); + verify(preferredRemoteKernelProvider.storePreferredRemoteKernelId(anything(), anything())).never(); + } + } + function verifyRemoteKernelTrackingUponKernelSelection(connection: KernelConnectionMetadata, selected: boolean) { + const controller = mock(); + const notebook = mock(); + const nbUri = Uri.file('a.ipynb'); + when(notebook.uri).thenReturn(nbUri); + when(controller.connection).thenReturn(connection); + + remoteConnectionHandler.activate(); + + verify(tracker.trackKernelIdAsUsed(anything(), anything(), anything())).never(); + verify(preferredRemoteKernelProvider.storePreferredRemoteKernelId(anything(), anything())).never(); + + onNotebookControllerSelectionChanged.fire({ + controller: instance(controller), + notebook: instance(notebook), + selected + }); + + if (connection.kind === 'connectToLiveRemoteKernel') { + if (selected) { + verify( + tracker.trackKernelIdAsUsed(nbUri, remoteKernelSpec.serverId, connection.kernelModel.id!) + ).once(); + } else { + verify( + tracker.trackKernelIdAsNotUsed(nbUri, remoteKernelSpec.serverId, connection.kernelModel.id!) + ).once(); + } + } else { + verify(tracker.trackKernelIdAsUsed(anything(), anything(), anything())).never(); + } + + if (selected && isLocalConnection(connection)) { + verify(preferredRemoteKernelProvider.clearPreferredRemoteKernelId(nbUri)).once(); + } + } + test('When starting a remote kernel spec ensure we track this', async () => { + verifyRemoteKernelTracking(remoteKernelSpec, 'jupyterExtension'); + }); + test('When starting a remote kernel spec from a 3rd party extension ensure we do not track this', async () => { + verifyRemoteKernelTracking(remoteKernelSpec, '3rdPartyExtension'); + }); + test('When starting a local kernel spec ensure we do not track this', async () => { + verifyRemoteKernelTracking(localKernelSpec, 'jupyterExtension'); + }); + test('When starting a local kernel spec from a 3rd party extension ensure we do not track this', async () => { + verifyRemoteKernelTracking(localKernelSpec, '3rdPartyExtension'); + }); + test('When starting a kernel related to a live kernel ensure we do not track this, as it will be tracked when the kernel is selected', async () => { + verifyRemoteKernelTracking(remoteLiveKernel1, 'jupyterExtension'); + }); + test('When starting a kernel related to a live kernel from a 3rd party extension ensure we do not track this', async () => { + verifyRemoteKernelTracking(remoteLiveKernel1, '3rdPartyExtension'); + }); + + test('Upon selecting a local kernelspec, ensure we clear the preferred remote kernel & not track this kernel', () => { + verifyRemoteKernelTrackingUponKernelSelection(localKernelSpec, true); + }); + test('Upon selecting a remote kernelspec, ensure we do not not track this kernel', () => { + verifyRemoteKernelTrackingUponKernelSelection(remoteKernelSpec, true); + }); + test('Upon selecting a remote live kernel, ensure we track this kernel', () => { + verifyRemoteKernelTrackingUponKernelSelection(remoteLiveKernel1, true); + }); + test('Upon un-selecting a remote live kernel, ensure we mark this kernel as no longer used', () => { + verifyRemoteKernelTrackingUponKernelSelection(remoteLiveKernel1, false); + }); +}); diff --git a/src/test/mocks/vsc/extHostedTypes.ts b/src/test/mocks/vsc/extHostedTypes.ts index 326e3378d1a..ab94f9ba93b 100644 --- a/src/test/mocks/vsc/extHostedTypes.ts +++ b/src/test/mocks/vsc/extHostedTypes.ts @@ -2249,4 +2249,7 @@ export namespace vscMockExtHostedTypes { export class QuickInputButtons { static readonly Back: vscode.QuickInputButton = {} as any; } + export class NotebookRendererScript { + constructor(public uri: vscode.Uri, public provides: string | string[] = []) {} + } }