Split NotebookControllerManager (#10537)

* Partial refactor * Update telemetry * Fixup rest of changes * Update telemetry * Update telemetry * Fixup build errors * Update telemetry * Fix inversify issue * Fix some unnecessary code * Handle potential unhandled error case * Review feedback * Update src/notebooks/controllers/controllerRegistration.ts Co-authored-by: Don Jayamanne <don.jayamanne@outlook.com> * Going back to old code * Fix interactive window * Update telemetry * Fix build issue Co-authored-by: Don Jayamanne <don.jayamanne@outlook.com>
microsoft · Jun 22, 2022 · 4af7b2a · 4af7b2a
1 parent 5b7e66d
commit 4af7b2a
Show file tree

Hide file tree

Showing 58 changed files with 1,580 additions and 1,421 deletions.
diff --git a/TELEMETRY.md b/TELEMETRY.md
diff --git a/src/intellisense/intellisenseProvider.node.ts b/src/intellisense/intellisenseProvider.node.ts
@@ -18,9 +18,9 @@ import { IDisposableRegistry, IConfigurationService, IsPreRelease } from '../pla
 import { IInterpreterService } from '../platform/interpreter/contracts';
 import { PythonEnvironment } from '../platform/pythonEnvironments/info';
 import { getInterpreterId } from '../platform/pythonEnvironments/info/interpreter';
-import { INotebookControllerManager, INotebookCompletionProvider, INotebookEditorProvider } from '../notebooks/types';
+import { INotebookCompletionProvider, INotebookEditorProvider } from '../notebooks/types';
 import { LanguageServer } from './languageServer.node';
-import { IVSCodeNotebookController } from '../notebooks/controllers/types';
+import { IControllerSelection, IVSCodeNotebookController } from '../notebooks/controllers/types';
 import { getComparisonKey } from '../platform/vscode-path/resources';
 import { CompletionRequest } from 'vscode-languageclient';
 import { NotebookPythonPathService } from './notebookPythonPathService.node';
@@ -43,7 +43,7 @@ export class IntellisenseProvider implements INotebookCompletionProvider, IExten
 
     constructor(
         @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
-        @inject(INotebookControllerManager) private readonly notebookControllerManager: INotebookControllerManager,
+        @inject(IControllerSelection) private readonly notebookControllerSelection: IControllerSelection,
         @inject(INotebookEditorProvider) private readonly notebookEditorProvider: INotebookEditorProvider,
         @inject(IVSCodeNotebook) private readonly notebooks: IVSCodeNotebook,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
@@ -56,7 +56,7 @@ export class IntellisenseProvider implements INotebookCompletionProvider, IExten
 
     public activate() {
         // Sign up for kernel change events on notebooks
-        this.notebookControllerManager.onNotebookControllerSelected(this.controllerChanged, this, this.disposables);
+        this.notebookControllerSelection.onControllerSelected(this.controllerChanged, this, this.disposables);
         // Sign up for notebook open and close events.
         this.notebooks.onDidOpenNotebookDocument(this.openedNotebook, this, this.disposables);
         this.notebooks.onDidCloseNotebookDocument(this.closedNotebook, this, this.disposables);
@@ -74,7 +74,7 @@ export class IntellisenseProvider implements INotebookCompletionProvider, IExten
     }
 
     public async getLanguageClient(notebook: NotebookDocument) {
-        const controller = this.notebookControllerManager.getSelectedNotebookController(notebook);
+        const controller = this.notebookControllerSelection.getSelected(notebook);
         const interpreter = controller
             ? controller.connection.interpreter
             : await this.interpreterService.getActiveInterpreter(notebook.uri);
@@ -165,7 +165,7 @@ export class IntellisenseProvider implements INotebookCompletionProvider, IExten
             !this.notebookPythonPathService.isPylanceUsingLspNotebooks()
         ) {
             // Create a language server as soon as we open. Otherwise intellisense will wait until we run.
-            const controller = this.notebookControllerManager.getSelectedNotebookController(n);
+            const controller = this.notebookControllerSelection.getSelected(n);
 
             // Save mapping from notebook to controller
             if (controller) {
@@ -207,9 +207,7 @@ export class IntellisenseProvider implements INotebookCompletionProvider, IExten
         // We should allow intellisense for a URI when the interpreter matches
         // the controller for the uri
         const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(uri);
-        const controller = notebook
-            ? this.notebookControllerManager.getSelectedNotebookController(notebook)
-            : undefined;
+        const controller = notebook ? this.notebookControllerSelection.getSelected(notebook) : undefined;
         const notebookInterpreter = controller ? controller.connection.interpreter : this.getActiveInterpreterSync(uri);
         let notebookId = notebookInterpreter ? this.getInterpreterIdFromCache(notebookInterpreter) : undefined;
 

diff --git a/src/intellisense/languageServer.node.ts b/src/intellisense/languageServer.node.ts
@@ -24,7 +24,7 @@ import { FileBasedCancellationStrategy } from './fileBasedCancellationStrategy.n
 import { createNotebookMiddleware, createPylanceMiddleware, NotebookMiddleware } from '@vscode/jupyter-lsp-middleware';
 import * as uuid from 'uuid/v4';
 import { NOTEBOOK_SELECTOR, PYTHON_LANGUAGE } from '../platform/common/constants';
-import { traceInfo } from '../platform/logging';
+import { traceInfo, traceInfoIfCI } from '../platform/logging';
 import { getInterpreterId } from '../platform/pythonEnvironments/info/interpreter';
 import { noop } from '../platform/common/utils/misc';
 import { PythonEnvironment } from '../platform/pythonEnvironments/info';
@@ -142,7 +142,11 @@ export class LanguageServer implements Disposable {
 
     public stopWatching(notebook: NotebookDocument) {
         // Tell the middleware to stop watching this document
-        this.middleware.stopWatching(notebook);
+        try {
+            this.middleware.stopWatching(notebook);
+        } catch (ex) {
+            traceInfoIfCI(`Error shutting down the LS.`, ex);
+        }
     }
 
     public startWatching(notebook: NotebookDocument) {

diff --git a/src/intellisense/notebookPythonPathService.node.ts b/src/intellisense/notebookPythonPathService.node.ts
@@ -4,7 +4,7 @@
 
 import { inject, injectable } from 'inversify';
 import { Disposable, extensions, Uri, workspace } from 'vscode';
-import { INotebookControllerManager, INotebookEditorProvider } from '../notebooks/types';
+import { INotebookEditorProvider } from '../notebooks/types';
 import { IExtensionSingleActivationService } from '../platform/activation/types';
 import { IPythonApiProvider } from '../platform/api/types';
 import { PylanceExtension, PythonExtension } from '../platform/common/constants';
@@ -13,6 +13,7 @@ import { IConfigurationService } from '../platform/common/types';
 import { IInterpreterService } from '../platform/interpreter/contracts';
 import * as semver from 'semver';
 import { traceInfo, traceVerbose } from '../platform/logging';
+import { IControllerSelection } from '../notebooks/controllers/types';
 
 /**
  * Manages use of the Python extension's registerJupyterPythonPathFunction API which
@@ -28,7 +29,7 @@ export class NotebookPythonPathService implements IExtensionSingleActivationServ
     constructor(
         @inject(IPythonApiProvider) private readonly apiProvider: IPythonApiProvider,
         @inject(INotebookEditorProvider) private readonly notebookEditorProvider: INotebookEditorProvider,
-        @inject(INotebookControllerManager) private readonly notebookControllerManager: INotebookControllerManager,
+        @inject(IControllerSelection) private readonly notebookControllerSelection: IControllerSelection,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
         @inject(IConfigurationService) private readonly configService: IConfigurationService
     ) {
@@ -121,7 +122,7 @@ export class NotebookPythonPathService implements IExtensionSingleActivationServ
             return undefined;
         }
 
-        const controller = this.notebookControllerManager.getSelectedNotebookController(notebook);
+        const controller = this.notebookControllerSelection.getSelected(notebook);
         const interpreter = controller
             ? controller.connection.interpreter
             : await this.interpreterService.getActiveInterpreter(uri);

diff --git a/src/interactive-window/debugger/jupyter/debuggingManager.ts b/src/interactive-window/debugger/jupyter/debuggingManager.ts
@@ -22,7 +22,6 @@ import {
 import { IKernelProvider } from '../../../kernels/types';
 import { IpykernelCheckResult, assertIsDebugConfig } from '../../../kernels/debugger/helper';
 import { KernelDebugAdapter } from './kernelDebugAdapter';
-import { INotebookControllerManager } from '../../../notebooks/types';
 import { IExtensionSingleActivationService } from '../../../platform/activation/types';
 import { ICommandManager, IApplicationShell, IVSCodeNotebook } from '../../../platform/common/application/types';
 import { IPlatformService } from '../../../platform/common/platform/types';
@@ -36,6 +35,7 @@ import { IFileGeneratedCodes } from '../../editor-integration/types';
 import { buildSourceMap } from '../helper';
 import { noop } from '../../../platform/common/utils/misc';
 import { IInteractiveWindowDebuggingManager } from '../../types';
+import { IControllerLoader, IControllerSelection } from '../../../notebooks/controllers/types';
 
 /**
  * The DebuggingManager maintains the mapping between notebook documents and debug sessions.
@@ -47,7 +47,8 @@ export class InteractiveWindowDebuggingManager
 {
     public constructor(
         @inject(IKernelProvider) kernelProvider: IKernelProvider,
-        @inject(INotebookControllerManager) notebookControllerManager: INotebookControllerManager,
+        @inject(IControllerSelection) controllerSelection: IControllerSelection,
+        @inject(IControllerLoader) controllerLoader: IControllerLoader,
         @inject(ICommandManager) commandManager: ICommandManager,
         @inject(IApplicationShell) appShell: IApplicationShell,
         @inject(IVSCodeNotebook) vscNotebook: IVSCodeNotebook,
@@ -56,7 +57,7 @@ export class InteractiveWindowDebuggingManager
         private readonly debugLocationTrackerFactory: IDebugLocationTrackerFactory,
         @inject(IConfigurationService) private readonly configService: IConfigurationService
     ) {
-        super(kernelProvider, notebookControllerManager, commandManager, appShell, vscNotebook);
+        super(kernelProvider, controllerLoader, controllerSelection, commandManager, appShell, vscNotebook);
     }
 
     public override async activate(): Promise<void> {

diff --git a/src/interactive-window/generatedCodeStoreManager.ts b/src/interactive-window/generatedCodeStoreManager.ts
@@ -5,7 +5,7 @@ import { inject, injectable } from 'inversify';
 import { NotebookDocument } from 'vscode';
 import { getAssociatedNotebookDocument } from '../kernels/helpers';
 import { IKernel, IKernelProvider } from '../kernels/types';
-import { INotebookControllerManager } from '../notebooks/types';
+import { IControllerSelection } from '../notebooks/controllers/types';
 import { IExtensionSyncActivationService } from '../platform/activation/types';
 import { InteractiveWindowView } from '../platform/common/constants';
 import { disposeAllDisposables } from '../platform/common/helpers';
@@ -20,7 +20,7 @@ export class GeneratedCodeStorageManager implements IExtensionSyncActivationServ
         @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(ICodeGeneratorFactory) private readonly codeGeneratorFactory: ICodeGeneratorFactory,
         @inject(IGeneratedCodeStorageFactory) private readonly storageFactory: IGeneratedCodeStorageFactory,
-        @inject(INotebookControllerManager) private readonly controllers: INotebookControllerManager
+        @inject(IControllerSelection) private readonly controllers: IControllerSelection
     ) {
         disposables.push(this);
     }
@@ -29,7 +29,7 @@ export class GeneratedCodeStorageManager implements IExtensionSyncActivationServ
     }
     activate(): void {
         this.kernelProvider.onDidCreateKernel(this.onDidCreateKernel, this, this.disposables);
-        this.controllers.onNotebookControllerSelected(this.onNotebookControllerSelected, this, this.disposables);
+        this.controllers.onControllerSelected(this.onNotebookControllerSelected, this, this.disposables);
     }
     private onNotebookControllerSelected({ notebook }: { notebook: NotebookDocument }) {
         this.storageFactory.get({ notebook })?.clear();

diff --git a/src/interactive-window/interactiveWindow.ts b/src/interactive-window/interactiveWindow.ts
@@ -42,7 +42,6 @@ import {
     KernelConnectionMetadata,
     NotebookCellRunState
 } from '../kernels/types';
-import { INotebookControllerManager } from '../notebooks/types';
 import { chainable } from '../platform/common/utils/decorators';
 import { InteractiveCellResultError } from '../platform/errors/interactiveCellResultError';
 import { DataScience } from '../platform/common/utils/localize';
@@ -56,7 +55,7 @@ import { generateCellsFromNotebookDocument } from './editor-integration/cellFact
 import { CellMatcher } from './editor-integration/cellMatcher';
 import { IInteractiveWindowLoadable, IInteractiveWindowDebugger, IInteractiveWindowDebuggingManager } from './types';
 import { generateInteractiveCode } from './helpers';
-import { IVSCodeNotebookController } from '../notebooks/controllers/types';
+import { IControllerSelection, IVSCodeNotebookController } from '../notebooks/controllers/types';
 import { DisplayOptions } from '../kernels/displayOptions';
 import { getInteractiveCellMetadata } from './helpers';
 import { KernelConnector } from '../notebooks/controllers/kernelConnector';
@@ -121,7 +120,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
         private _owner: Resource,
         private mode: InteractiveWindowMode,
         private readonly exportDialog: IExportDialog,
-        private readonly notebookControllerManager: INotebookControllerManager,
+        private readonly notebookControllerSelection: IControllerSelection,
         private readonly serviceContainer: IServiceContainer,
         private readonly interactiveWindowDebugger: IInteractiveWindowDebugger | undefined,
         private readonly errorHandler: IDataScienceErrorHandler,
@@ -380,13 +379,13 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
     }
 
     private listenForControllerSelection() {
-        const controller = this.notebookControllerManager.getSelectedNotebookController(this.notebookEditor.notebook);
+        const controller = this.notebookControllerSelection.getSelected(this.notebookEditor.notebook);
         if (controller !== undefined) {
             this.registerControllerChangeListener(controller);
         }
 
         // Ensure we hear about any controller changes so we can update our cached promises
-        this.notebookControllerManager.onNotebookControllerSelected(
+        this.notebookControllerSelection.onControllerSelected(
             (e: { notebook: NotebookDocument; controller: IVSCodeNotebookController }) => {
                 if (e.notebook !== this.notebookEditor.notebook) {
                     return;
@@ -426,7 +425,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
         const insertionIndex =
             notebookCell && notebookCell.index >= 0 ? notebookCell.index : this.notebookEditor.notebook.cellCount;
         // If possible display the error message in the cell.
-        const controller = this.notebookControllerManager.getSelectedNotebookController(this.notebookEditor.notebook);
+        const controller = this.notebookControllerSelection.getSelected(this.notebookEditor.notebook);
         const output = createOutputWithErrorMessageForDisplay(message);
         if (this.notebookEditor.notebook.cellCount === 0 || !controller || !output || !notebookCell) {
             const edit = new WorkspaceEdit();

diff --git a/src/interactive-window/interactiveWindowCommandListener.ts b/src/interactive-window/interactiveWindowCommandListener.ts
@@ -23,21 +23,15 @@ import {
     IDocumentManager,
     IVSCodeNotebook
 } from '../platform/common/application/types';
-import {
-    Commands,
-    JupyterNotebookView,
-    JVSC_EXTENSION_ID,
-    PYTHON_LANGUAGE,
-    Telemetry
-} from '../platform/common/constants';
+import { Commands, JVSC_EXTENSION_ID, PYTHON_LANGUAGE, Telemetry } from '../platform/common/constants';
 import { traceError, traceInfo } from '../platform/logging';
 import { IFileSystem } from '../platform/common/platform/types';
 import { IConfigurationService, IDataScienceCommandListener, IDisposableRegistry } from '../platform/common/types';
 import * as localize from '../platform/common/utils/localize';
 import { captureTelemetry } from '../telemetry';
 import { CommandSource } from '../platform/testing/common/constants';
 import { JupyterInstallError } from '../platform/errors/jupyterInstallError';
-import { INotebookControllerManager, INotebookEditorProvider } from '../notebooks/types';
+import { INotebookEditorProvider } from '../notebooks/types';
 import { KernelConnectionMetadata } from '../kernels/types';
 import { INotebookExporter, IJupyterExecution } from '../kernels/jupyter/types';
 import { IDataScienceErrorHandler } from '../kernels/errors/types';
@@ -49,6 +43,7 @@ import { getDisplayPath, getFilePath } from '../platform/common/platform/fs-path
 import { chainWithPendingUpdates } from '../kernels/execution/notebookUpdater';
 import { openAndShowNotebook } from '../platform/common/utils/notebooks';
 import { noop } from '../platform/common/utils/misc';
+import { IControllerPreferredService } from '../notebooks/controllers/types';
 
 @injectable()
 export class InteractiveWindowCommandListener implements IDataScienceCommandListener {
@@ -69,7 +64,7 @@ export class InteractiveWindowCommandListener implements IDataScienceCommandList
         @inject(IClipboard) private clipboard: IClipboard,
         @inject(IVSCodeNotebook) private notebook: IVSCodeNotebook,
         @inject(ICommandManager) private commandManager: ICommandManager,
-        @inject(INotebookControllerManager) private controllerManager: INotebookControllerManager
+        @inject(IControllerPreferredService) private controllerPreferredService: IControllerPreferredService
     ) {}
 
     public register(commandManager: ICommandManager): void {
@@ -279,17 +274,14 @@ export class InteractiveWindowCommandListener implements IDataScienceCommandList
                         getDisplayPath(file)
                     );
                     // Next open this notebook & execute it.
-                    await this.notebook.showNotebookDocument(uri, {
+                    const editor = await this.notebook.showNotebookDocument(uri, {
                         preserveFocus: false,
                         viewColumn: ViewColumn.Beside
                     });
-                    const preferredController = await this.controllerManager.getActiveInterpreterOrDefaultController(
-                        JupyterNotebookView,
-                        file
-                    );
-                    if (preferredController) {
+                    const { controller } = await this.controllerPreferredService.computePreferred(editor.notebook);
+                    if (controller) {
                         await this.commandManager.executeCommand('notebook.selectKernel', {
-                            id: preferredController.id,
+                            id: controller.id,
                             extension: JVSC_EXTENSION_ID
                         });
                     }