Simplify starting how we start local Jupyter, remote Jupyter and raw kernels (part 2)

src/kernels/common/baseJupyterSession.ts

@@ -96,13 +96,17 @@ export abstract class BaseJupyterSession implements IBaseKernelConnectionSession
     private _wrappedKernel?: KernelConnectionWrapper;
     private _isDisposed?: boolean;
     private readonly _disposed = new EventEmitter<void>();
+    private readonly didShutdown = new EventEmitter<void>();
     protected readonly disposables: IDisposable[] = [];
     public get disposed() {
         return this._isDisposed === true;
     }
     public get onDidDispose() {
         return this._disposed.event;
     }
+    public get onDidShutdown() {
+        return this.didShutdown.event;
+    }
     protected get session(): ISessionWithSocket | undefined {
         return this._session;
     }
@@ -525,6 +529,8 @@ export abstract class BaseJupyterSession implements IBaseKernelConnectionSession
             this.restartSessionPromise = undefined;
             this.onStatusChangedEvent.fire('dead');
             this._disposed.fire();
+            this.didShutdown.fire();
+            this.didShutdown.dispose();
             this._disposed.dispose();
             this.onStatusChangedEvent.dispose();
             this.previousAnyMessageHandler?.dispose();

src/kernels/common/kernelConnectionSessionCreator.ts

@@ -0,0 +1,160 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { inject, injectable, optional } from 'inversify';
+import {
+    IJupyterConnection,
+    IKernelConnectionSession,
+    IKernelConnectionSessionCreator,
+    isLocalConnection,
+    isRemoteConnection,
+    KernelConnectionSessionCreationOptions
+} from '../types';
+import { Cancellation } from '../../platform/common/cancellation';
+import { IRawKernelConnectionSessionCreator } from '../raw/types';
+import { IJupyterServerProvider, IJupyterSessionManagerFactory } from '../jupyter/types';
+import { JupyterKernelConnectionSessionCreator } from '../jupyter/session/jupyterKernelConnectionSessionCreator';
+import { JupyterConnection } from '../jupyter/connection/jupyterConnection';
+import { Telemetry, sendTelemetryEvent } from '../../telemetry';
+import { JupyterSelfCertsError } from '../../platform/errors/jupyterSelfCertsError';
+import { JupyterSelfCertsExpiredError } from '../../platform/errors/jupyterSelfCertsExpiredError';
+import { RemoteJupyterServerConnectionError } from '../../platform/errors/remoteJupyterServerConnectionError';
+import { disposeAllDisposables } from '../../platform/common/helpers';
+import { IAsyncDisposableRegistry, IDisposable } from '../../platform/common/types';
+import { KernelProgressReporter } from '../../platform/progress/kernelProgressReporter';
+import { DataScience } from '../../platform/common/utils/localize';
+import { Disposable } from 'vscode';
+import { noop } from '../../platform/common/utils/misc';
+import { LocalJupyterServerConnectionError } from '../../platform/errors/localJupyterServerConnectionError';
+import { BaseError } from '../../platform/errors/types';
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+const LocalHosts = ['localhost', '127.0.0.1', '::1'];
+
+/**
+ * Generic class for connecting to a server. Probably could be renamed as it doesn't provide notebooks, but rather connections.
+ */
+@injectable()
+export class KernelConnectionSessionCreator implements IKernelConnectionSessionCreator {
+    constructor(
+        @inject(IRawKernelConnectionSessionCreator)
+        @optional()
+        private readonly rawKernelSessionCreator: IRawKernelConnectionSessionCreator | undefined,
+        @inject(IJupyterServerProvider)
+        private readonly jupyterNotebookProvider: IJupyterServerProvider,
+        @inject(IJupyterSessionManagerFactory) private readonly sessionManagerFactory: IJupyterSessionManagerFactory,
+        @inject(JupyterKernelConnectionSessionCreator)
+        private readonly jupyterSessionCreator: JupyterKernelConnectionSessionCreator,
+        @inject(JupyterConnection) private readonly jupyterConnection: JupyterConnection,
+        @inject(IAsyncDisposableRegistry) private readonly asyncDisposables: IAsyncDisposableRegistry
+    ) {}
+
+    public async create(options: KernelConnectionSessionCreationOptions): Promise<IKernelConnectionSession> {
+        const kernelConnection = options.kernelConnection;
+        const isLocal = isLocalConnection(kernelConnection);
+
+        if (this.rawKernelSessionCreator?.isSupported && isLocal) {
+            return this.createRawKernelSession(this.rawKernelSessionCreator, options);
+        }
+
+        const disposables: IDisposable[] = [];
+        let progressReporter: IDisposable | undefined;
+        const createProgressReporter = () => {
+            if (options.ui.disableUI || progressReporter) {
+                return;
+            }
+            // Status depends upon if we're about to connect to existing server or not.
+            progressReporter = KernelProgressReporter.createProgressReporter(
+                options.resource,
+                isRemoteConnection(options.kernelConnection)
+                    ? DataScience.connectingToJupyter
+                    : DataScience.startingJupyter
+            );
+            disposables.push(progressReporter);
+        };
+        if (options.ui.disableUI) {
+            options.ui.onDidChangeDisableUI(createProgressReporter, this, disposables);
+        }
+        createProgressReporter();
+
+        return this.createJupyterKernelSession(options).finally(() => disposeAllDisposables(disposables));
+    }
+    private createRawKernelSession(
+        factory: IRawKernelConnectionSessionCreator,
+        options: KernelConnectionSessionCreationOptions
+    ): Promise<IKernelConnectionSession> {
+        return factory.create(options.resource, options.kernelConnection, options.ui, options.token);
+    }
+    private async createJupyterKernelSession(
+        options: KernelConnectionSessionCreationOptions
+    ): Promise<IKernelConnectionSession> {
+        let connection: undefined | IJupyterConnection;
+
+        // Check to see if we support ipykernel or not
+        const disposablesWhenThereAreFailures: IDisposable[] = [];
+        try {
+            connection = isRemoteConnection(options.kernelConnection)
+                ? await this.jupyterConnection.createConnectionInfo({
+                      serverId: options.kernelConnection.serverId
+                  })
+                : await this.jupyterNotebookProvider.getOrCreateServer({
+                      resource: options.resource,
+                      token: options.token,
+                      ui: options.ui
+                  });
+
+            if (!connection.localLaunch && LocalHosts.includes(connection.hostName.toLowerCase())) {
+                sendTelemetryEvent(Telemetry.ConnectRemoteJupyterViaLocalHost);
+            }
+
+            Cancellation.throwIfCanceled(options.token);
+
+            const sessionManager = await this.sessionManagerFactory.create(connection);
+            this.asyncDisposables.push(sessionManager);
+            disposablesWhenThereAreFailures.push(new Disposable(() => sessionManager.dispose().catch(noop)));
+
+            Cancellation.throwIfCanceled(options.token);
+            // Disposing session manager will dispose all sessions that were started by that session manager.
+            // Hence Session managers should be disposed only if the corresponding session is shutdown.
+            const session = await this.jupyterSessionCreator.create({
+                creator: options.creator,
+                kernelConnection: options.kernelConnection,
+                resource: options.resource,
+                sessionManager,
+                token: options.token,
+                ui: options.ui
+            });
+            session.onDidShutdown(() => sessionManager.dispose());
+            return session;
+        } catch (ex) {
+            disposeAllDisposables(disposablesWhenThereAreFailures);
+
+            if (isRemoteConnection(options.kernelConnection)) {
+                sendTelemetryEvent(Telemetry.ConnectRemoteFailedJupyter, undefined, undefined, ex);
+                // Check for the self signed certs error specifically
+                if (!connection) {
+                    throw ex;
+                } else if (JupyterSelfCertsError.isSelfCertsError(ex)) {
+                    sendTelemetryEvent(Telemetry.ConnectRemoteSelfCertFailedJupyter);
+                    throw new JupyterSelfCertsError(connection.baseUrl);
+                } else if (JupyterSelfCertsExpiredError.isSelfCertsExpiredError(ex)) {
+                    sendTelemetryEvent(Telemetry.ConnectRemoteExpiredCertFailedJupyter);
+                    throw new JupyterSelfCertsExpiredError(connection.baseUrl);
+                } else {
+                    throw new RemoteJupyterServerConnectionError(
+                        connection.baseUrl,
+                        options.kernelConnection.serverId,
+                        ex
+                    );
+                }
+            } else {
+                sendTelemetryEvent(Telemetry.ConnectFailedJupyter, undefined, undefined, ex);
+                if (ex instanceof BaseError) {
+                    throw ex;
+                } else {
+                    throw new LocalJupyterServerConnectionError(ex);
+                }
+            }
+        }
+    }
+}

src/kernels/jupyter/launcher/kernelConnectionSessionCreator.unit.test.ts → src/kernels/common/kernelConnectionSessionCreator.unit.test.ts

@@ -1,16 +1,21 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import { IDisposable } from '@fluentui/react';
 import { expect } from 'chai';
 import { anything, instance, mock, when } from 'ts-mockito';
 import * as vscode from 'vscode';
-import { PythonExtensionChecker } from '../../../platform/api/pythonApi';
-import { IJupyterKernelConnectionSession, KernelConnectionMetadata } from '../../types';
-import { DisplayOptions } from '../../displayOptions';
-import { IJupyterNotebookProvider } from '../types';
-import { IRawKernelConnectionSessionCreator } from '../../raw/types';
-import { IDisposable } from '../../../platform/common/types';
-import { disposeAllDisposables } from '../../../platform/common/helpers';
+import { EventEmitter } from 'vscode';
+import { PythonExtensionChecker } from '../../platform/api/pythonApi';
+import { AsyncDisposableRegistry } from '../../platform/common/asyncDisposableRegistry';
+import { disposeAllDisposables } from '../../platform/common/helpers';
+import { IAsyncDisposableRegistry } from '../../platform/common/types';
+import { DisplayOptions } from '../displayOptions';
+import { JupyterConnection } from '../jupyter/connection/jupyterConnection';
+import { JupyterKernelConnectionSessionCreator } from '../jupyter/session/jupyterKernelConnectionSessionCreator';
+import { IJupyterServerProvider, IJupyterSessionManagerFactory } from '../jupyter/types';
+import { IRawKernelConnectionSessionCreator } from '../raw/types';
+import { IJupyterKernelConnectionSession, KernelConnectionMetadata } from '../types';
 import { KernelConnectionSessionCreator } from './kernelConnectionSessionCreator';
 
 function Uri(filename: string): vscode.Uri {
@@ -20,12 +25,14 @@ function Uri(filename: string): vscode.Uri {
 /* eslint-disable  */
 suite('NotebookProvider', () => {
     let kernelConnectionSessionCreator: KernelConnectionSessionCreator;
-    let jupyterNotebookProvider: IJupyterNotebookProvider;
+    let jupyterNotebookProvider: IJupyterServerProvider;
     let rawKernelSessionCreator: IRawKernelConnectionSessionCreator;
     let cancelToken: vscode.CancellationTokenSource;
     const disposables: IDisposable[] = [];
+    let asyncDisposables: IAsyncDisposableRegistry;
+    let onDidShutdown: EventEmitter<void>;
     setup(() => {
-        jupyterNotebookProvider = mock<IJupyterNotebookProvider>();
+        jupyterNotebookProvider = mock<IJupyterServerProvider>();
         rawKernelSessionCreator = mock<IRawKernelConnectionSessionCreator>();
         cancelToken = new vscode.CancellationTokenSource();
         disposables.push(cancelToken);
@@ -34,21 +41,37 @@ suite('NotebookProvider', () => {
         when(extensionChecker.isPythonExtensionInstalled).thenReturn(true);
         const onDidChangeEvent = new vscode.EventEmitter<void>();
         disposables.push(onDidChangeEvent);
-
+        onDidShutdown = new vscode.EventEmitter<void>();
+        disposables.push(onDidShutdown);
+        const sessionManagerFactory = mock<IJupyterSessionManagerFactory>();
+        const jupyterSessionCreator = mock<JupyterKernelConnectionSessionCreator>();
+        const jupyterConnection = mock<JupyterConnection>();
+        when(jupyterConnection.createConnectionInfo(anything())).thenResolve({
+            localLaunch: true,
+            baseUrl: 'http://localhost:8888'
+        } as any);
+        const mockSession = mock<IJupyterKernelConnectionSession>();
+        when(mockSession.onDidShutdown).thenReturn(onDidShutdown.event);
+        instance(mockSession as any).then = undefined;
+        when(jupyterSessionCreator.create(anything())).thenResolve(instance(mockSession));
+        asyncDisposables = new AsyncDisposableRegistry();
         kernelConnectionSessionCreator = new KernelConnectionSessionCreator(
             instance(rawKernelSessionCreator),
-            instance(jupyterNotebookProvider)
+            instance(jupyterNotebookProvider),
+            instance(sessionManagerFactory),
+            instance(jupyterSessionCreator),
+            instance(jupyterConnection),
+            asyncDisposables
         );
     });
-    teardown(() => disposeAllDisposables(disposables));
+    teardown(async () => {
+        disposeAllDisposables(disposables);
+        await asyncDisposables.dispose();
+    });
     test('NotebookProvider getOrCreateNotebook jupyter provider does not have notebook already', async () => {
-        const mockSession = mock<IJupyterKernelConnectionSession>();
-        instance(mockSession as any).then = undefined;
-        when(jupyterNotebookProvider.createNotebook(anything())).thenResolve(instance(mockSession));
-        when(jupyterNotebookProvider.connect(anything())).thenResolve({} as any);
+        when(jupyterNotebookProvider.getOrCreateServer(anything())).thenResolve({} as any);
         const doc = mock<vscode.NotebookDocument>();
         when(doc.uri).thenReturn(Uri('C:\\\\foo.py'));
-
         const session = await kernelConnectionSessionCreator.create({
             resource: Uri('C:\\\\foo.py'),
             kernelConnection: instance(mock<KernelConnectionMetadata>()),
@@ -60,10 +83,7 @@ suite('NotebookProvider', () => {
     });
 
     test('NotebookProvider getOrCreateNotebook second request should return the notebook already cached', async () => {
-        const mockSession = mock<IJupyterKernelConnectionSession>();
-        instance(mockSession as any).then = undefined;
-        when(jupyterNotebookProvider.createNotebook(anything())).thenResolve(instance(mockSession));
-        when(jupyterNotebookProvider.connect(anything())).thenResolve({} as any);
+        when(jupyterNotebookProvider.getOrCreateServer(anything())).thenResolve({} as any);
         const doc = mock<vscode.NotebookDocument>();
         when(doc.uri).thenReturn(Uri('C:\\\\foo.py'));