Delegate management of notebooks (creation and disposing) to INotebookProvider

src/client/datascience/interactive-common/interactiveBase.ts

@@ -89,6 +89,7 @@ import {
     IMessageCell,
     INotebook,
     INotebookExporter,
+    INotebookProvider,
     INotebookServer,
     INotebookServerOptions,
     InterruptResult,
@@ -152,7 +153,8 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
         @unmanaged() title: string,
         @unmanaged() viewColumn: ViewColumn,
         @unmanaged() experimentsManager: IExperimentsManager,
-        @unmanaged() private switcher: KernelSwitcher
+        @unmanaged() private switcher: KernelSwitcher,
+        @unmanaged() private readonly notebookProvider: INotebookProvider
     ) {
         super(
             configuration,
@@ -326,13 +328,6 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
         super.dispose();
         this.listeners.forEach(l => l.dispose());
         this.updateContexts(undefined);
-
-        // When closing an editor, dispose of the notebook associated with it.
-        // This won't work when we have multiple views of the notebook though. Notebook ownership
-        // should probably move to whatever owns the backing model.
-        return this.notebook?.dispose().then(() => {
-            this._notebook = undefined;
-        });
     }
 
     public startProgress() {
@@ -1093,11 +1088,7 @@ export abstract class InteractiveBase extends WebViewHost<IInteractiveWindowMapp
                 this.getNotebookOptions()
             ]);
             try {
-                // We could have multiple native editors opened for the same file/model.
-                notebook = uri ? await server.getNotebook(uri) : undefined;
-                if (!notebook) {
-                    notebook = uri ? await server.createNotebook(resource, uri, options?.metadata) : undefined;
-                }
+                notebook = uri ? await this.notebookProvider.getNotebook(server, uri, options?.metadata) : undefined;
             } catch (e) {
                 // If we get an invalid kernel error, make sure to ask the user to switch
                 if (e instanceof JupyterInvalidKernelError && server && server.getConnectionInfo()?.localLaunch) {

src/client/datascience/interactive-common/notebookProvider.ts

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { nbformat } from '@jupyterlab/coreutils';
+import { injectable } from 'inversify';
+import { Uri } from 'vscode';
+import { noop } from '../../common/utils/misc';
+import { INotebook, INotebookProvider, INotebookServer } from '../types';
+
+@injectable()
+export class BaseNotebookProvider implements INotebookProvider {
+    protected readonly notebooks = new Map<string, Promise<INotebook>>();
+    public async getNotebook(
+        server: INotebookServer,
+        resource: Uri,
+        metadata?: nbformat.INotebookMetadata | undefined
+    ): Promise<INotebook> {
+        // We could have multiple native editors opened for the same file/model.
+        const notebook = await server.getNotebook(resource);
+        if (notebook) {
+            return notebook;
+        }
+
+        if (this.notebooks.get(resource.fsPath)) {
+            return this.notebooks.get(resource.fsPath)!!;
+        }
+
+        const promise = server.createNotebook(resource, resource, metadata);
+        this.notebooks.set(resource.fsPath, promise);
+
+        // Remove promise from cache if the same promise still exists.
+        const removeFromCache = () => {
+            const cachedPromise = this.notebooks.get(resource.fsPath);
+            if (cachedPromise === promise) {
+                this.notebooks.delete(resource.fsPath);
+            }
+        };
+
+        // If the notebook is disposed, remove from cache.
+        promise.then(nb => nb.onDisposed(removeFromCache)).catch(noop);
+
+        // If promise fails, then remove the promise from cache.
+        promise.catch(removeFromCache);
+
+        return promise;
+    }
+    protected async disposeNotebook(resource: Uri) {
+        // First find all notebooks associated with this editor (ipynb file).
+        const notebookPromise = this.notebooks.get(resource.fsPath);
+        if (!notebookPromise) {
+            // Possible it was closed before a notebook could be created.
+            return;
+        }
+        this.notebooks.delete(resource.fsPath);
+        const notebook = await notebookPromise.catch(noop);
+        if (!notebook) {
+            return;
+        }
+
+        await notebook.dispose().catch(noop);
+    }
+}

src/client/datascience/interactive-ipynb/nativeEditor.ts

@@ -78,6 +78,7 @@ import {
     INotebookExporter,
     INotebookImporter,
     INotebookModel,
+    INotebookProvider,
     INotebookServerOptions,
     IStatusProvider,
     IThemeFinder,
@@ -90,6 +91,7 @@ import { nbformat } from '@jupyterlab/coreutils';
 import cloneDeep = require('lodash/cloneDeep');
 import { concatMultilineStringInput } from '../../../datascience-ui/common';
 import { KernelSwitcher } from '../jupyter/kernels/kernelSwitcher';
+import { NativeNotebookProvider } from './notebookProvider';
 
 const nativeEditorDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 @injectable()
@@ -176,7 +178,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         @inject(ProgressReporter) progressReporter: ProgressReporter,
         @inject(IExperimentsManager) experimentsManager: IExperimentsManager,
         @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry,
-        @inject(KernelSwitcher) switcher: KernelSwitcher
+        @inject(KernelSwitcher) switcher: KernelSwitcher,
+        @inject(NativeNotebookProvider) notebookProvider: INotebookProvider
     ) {
         super(
             progressReporter,
@@ -210,7 +213,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
             localize.DataScience.nativeEditorTitle(),
             ViewColumn.Active,
             experimentsManager,
-            switcher
+            switcher,
+            notebookProvider
         );
         asyncRegistry.push(this);
 

src/client/datascience/interactive-ipynb/nativeEditorOldWebView.ts

@@ -45,12 +45,14 @@ import {
     INotebookExporter,
     INotebookImporter,
     INotebookModel,
+    INotebookProvider,
     IStatusProvider,
     IThemeFinder
 } from '../types';
 import { NativeEditor } from './nativeEditor';
 import { NativeEditorStorage } from './nativeEditorStorage';
 import { NativeEditorSynchronizer } from './nativeEditorSynchronizer';
+import { NativeNotebookProvider } from './notebookProvider';
 
 enum AskForSaveResult {
     Yes,
@@ -97,7 +99,8 @@ export class NativeEditorOldWebView extends NativeEditor {
         @inject(ProgressReporter) progressReporter: ProgressReporter,
         @inject(IExperimentsManager) experimentsManager: IExperimentsManager,
         @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry,
-        @inject(KernelSwitcher) switcher: KernelSwitcher
+        @inject(KernelSwitcher) switcher: KernelSwitcher,
+        @inject(NativeNotebookProvider) notebookProvider: INotebookProvider
     ) {
         super(
             listeners,
@@ -127,7 +130,8 @@ export class NativeEditorOldWebView extends NativeEditor {
             progressReporter,
             experimentsManager,
             asyncRegistry,
-            switcher
+            switcher,
+            notebookProvider
         );
         asyncRegistry.push(this);
         // No ui syncing in old notebooks.

src/client/datascience/interactive-ipynb/notebookProvider.ts

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { IFileSystem } from '../../common/platform/types';
+import { IDisposableRegistry } from '../../common/types';
+import { BaseNotebookProvider } from '../interactive-common/notebookProvider';
+import { INotebookEditor, INotebookEditorProvider } from '../types';
+
+@injectable()
+export class NativeNotebookProvider extends BaseNotebookProvider {
+    constructor(
+        @inject(IFileSystem) private readonly fs: IFileSystem,
+        @inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry
+    ) {
+        super();
+        disposables.push(editorProvider.onDidCloseNotebookEditor(this.onDidCloseNotebookEditor, this));
+    }
+
+    protected async onDidCloseNotebookEditor(editor: INotebookEditor) {
+        // First find all notebooks associated with this editor (ipynb file).
+        const editors = this.editorProvider.editors.filter(
+            e => this.fs.arePathsSame(e.file.fsPath, editor.file.fsPath) && e !== editor
+        );
+
+        // If we have no editors for this file, then dispose the notebook.
+        if (editors.length === 0) {
+            await super.disposeNotebook(editor.file);
+        }
+    }
+}

src/client/datascience/interactive-window/interactiveWindow.ts

@@ -53,11 +53,13 @@ import {
     IJupyterKernelSpec,
     IJupyterVariables,
     INotebookExporter,
+    INotebookProvider,
     INotebookServerOptions,
     IStatusProvider,
     IThemeFinder,
     WebViewViewChangeEventArgs
 } from '../types';
+import { InteractiveWindowNotebookProvider } from './notebookProvider';
 
 const historyReactDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 
@@ -108,7 +110,8 @@ export class InteractiveWindow extends InteractiveBase implements IInteractiveWi
         @inject(IMemento) @named(GLOBAL_MEMENTO) globalStorage: Memento,
         @inject(ProgressReporter) progressReporter: ProgressReporter,
         @inject(IExperimentsManager) experimentsManager: IExperimentsManager,
-        @inject(KernelSwitcher) switcher: KernelSwitcher
+        @inject(KernelSwitcher) switcher: KernelSwitcher,
+        @inject(InteractiveWindowNotebookProvider) notebookProvider: INotebookProvider
     ) {
         super(
             progressReporter,
@@ -142,7 +145,8 @@ export class InteractiveWindow extends InteractiveBase implements IInteractiveWi
             localize.DataScience.historyTitle(),
             ViewColumn.Two,
             experimentsManager,
-            switcher
+            switcher,
+            notebookProvider
         );
 
         // Send a telemetry event to indicate window is opening

src/client/datascience/interactive-window/notebookProvider.ts

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { IDisposableRegistry } from '../../common/types';
+import { noop } from '../../common/utils/misc';
+import { BaseNotebookProvider } from '../interactive-common/notebookProvider';
+import { IInteractiveWindowProvider } from '../types';
+
+@injectable()
+export class InteractiveWindowNotebookProvider extends BaseNotebookProvider {
+    constructor(
+        @inject(IInteractiveWindowProvider) private readonly interactiveWindowProvider: IInteractiveWindowProvider,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry
+    ) {
+        super();
+        disposables.push(
+            interactiveWindowProvider.onDidChangeActiveInteractiveWindow(this.checkAndDisposeNotebook, this)
+        );
+    }
+
+    /**
+     * Interactive windows have just one window.
+     * When that it closed, just close all of the notebooks associated with interactive windows.
+     */
+    protected checkAndDisposeNotebook() {
+        if (this.interactiveWindowProvider.getActive()) {
+            return;
+        }
+
+        Array.from(this.notebooks.values()).forEach(promise => {
+            promise.then(notebook => notebook.dispose()).catch(noop);
+        });
+
+        this.notebooks.clear();
+    }
+}

src/client/datascience/jupyter/jupyterNotebook.ts

@@ -156,10 +156,14 @@ export class JupyterNotebookBase implements INotebook {
     private _workingDirectory: string | undefined;
     private _loggers: INotebookExecutionLogger[] = [];
     private onStatusChangedEvent: EventEmitter<ServerStatus> | undefined;
+    public get onDisposed(): Event<void> {
+        return this.disposed.event;
+    }
     public get onKernelChanged(): Event<IJupyterKernelSpec | LiveKernelModel> {
         return this.kernelChanged.event;
     }
     private kernelChanged = new EventEmitter<IJupyterKernelSpec | LiveKernelModel>();
+    private disposed = new EventEmitter<void>();
     private sessionStatusChanged: Disposable | undefined;
     private initializedMatplotlib = false;
 
@@ -211,6 +215,7 @@ export class JupyterNotebookBase implements INotebook {
             const dispose = this.session ? this.session.dispose() : undefined;
             return dispose ? dispose : Promise.resolve();
         }
+        this.disposed.fire();
         return Promise.resolve();
     }
 

src/client/datascience/jupyter/liveshare/guestJupyterNotebook.ts

@@ -36,6 +36,7 @@ export class GuestJupyterNotebook
     public onKernelChanged: Event<IJupyterKernelSpec | LiveKernelModel> = new EventEmitter<
         IJupyterKernelSpec | LiveKernelModel
     >().event;
+    public onDisposed = new EventEmitter<void>().event;
     private responseQueue: ResponseQueue = new ResponseQueue();
     private onStatusChangedEvent: EventEmitter<ServerStatus> | undefined;
 

src/client/datascience/jupyter/liveshare/hostJupyterNotebook.ts

@@ -40,7 +40,7 @@ export class HostJupyterNotebook
     private localResponses: ResponseQueue = new ResponseQueue();
     private requestLog: Map<string, number> = new Map<string, number>();
     private catchupPendingCount: number = 0;
-    private disposed = false;
+    private isDisposed = false;
     constructor(
         liveShare: ILiveShareApi,
         session: IJupyterSession,
@@ -74,8 +74,8 @@ export class HostJupyterNotebook
     }
 
     public dispose = async (): Promise<void> => {
-        if (!this.disposed) {
-            this.disposed = true;
+        if (!this.isDisposed) {
+            this.isDisposed = true;
             await super.dispose();
             const api = await this.api;
             return this.onDetach(api);
@@ -85,7 +85,7 @@ export class HostJupyterNotebook
     public async onAttach(api: vsls.LiveShare | null): Promise<void> {
         await super.onAttach(api);
 
-        if (api && !this.disposed) {
+        if (api && !this.isDisposed) {
             const service = await this.waitForService();
 
             // Attach event handlers to different requests

src/client/datascience/serviceRegistry.ts

@@ -39,9 +39,11 @@ import { NativeEditorProvider } from './interactive-ipynb/nativeEditorProvider';
 import { NativeEditorProviderOld } from './interactive-ipynb/nativeEditorProviderOld';
 import { NativeEditorStorage } from './interactive-ipynb/nativeEditorStorage';
 import { NativeEditorSynchronizer } from './interactive-ipynb/nativeEditorSynchronizer';
+import { NativeNotebookProvider } from './interactive-ipynb/notebookProvider';
 import { InteractiveWindow } from './interactive-window/interactiveWindow';
 import { InteractiveWindowCommandListener } from './interactive-window/interactiveWindowCommandListener';
 import { InteractiveWindowProvider } from './interactive-window/interactiveWindowProvider';
+import { InteractiveWindowNotebookProvider } from './interactive-window/notebookProvider';
 import { JupyterCommandLineSelector } from './jupyter/commandLineSelector';
 import { JupyterCommandFactory } from './jupyter/interpreter/jupyterCommand';
 import { JupyterCommandFinder } from './jupyter/interpreter/jupyterCommandFinder';
@@ -108,6 +110,7 @@ import {
     INotebookExecutionLogger,
     INotebookExporter,
     INotebookImporter,
+    INotebookProvider,
     INotebookServer,
     INotebookStorage,
     IPlotViewer,
@@ -188,6 +191,8 @@ export function registerTypes(serviceManager: IServiceManager) {
     serviceManager.addSingleton<NotebookStarter>(NotebookStarter, NotebookStarter);
     serviceManager.addSingleton<ProgressReporter>(ProgressReporter, ProgressReporter);
     serviceManager.addSingleton<NativeEditorSynchronizer>(NativeEditorSynchronizer, NativeEditorSynchronizer);
+    serviceManager.addSingleton<INotebookProvider>(InteractiveWindowNotebookProvider, InteractiveWindowNotebookProvider);
+    serviceManager.addSingleton<INotebookProvider>(NativeNotebookProvider, NativeNotebookProvider);
 
     // Temporary code, to allow users to revert to the old behavior.
     const cfg = serviceManager.get<IWorkspaceService>(IWorkspaceService).getConfiguration('python.dataScience', undefined);

src/client/datascience/types.ts

@@ -113,6 +113,7 @@ export interface INotebook extends IAsyncDisposable {
     readonly server: INotebookServer;
     readonly status: ServerStatus;
     onSessionStatusChanged: Event<ServerStatus>;
+    onDisposed: Event<void>;
     onKernelChanged: Event<IJupyterKernelSpec | LiveKernelModel>;
     clear(id: string): void;
     executeObservable(code: string, file: string, line: number, id: string, silent: boolean): Observable<ICell[]>;
@@ -842,3 +843,10 @@ type WebViewViewState = {
     readonly active: boolean;
 };
 export type WebViewViewChangeEventArgs = { current: WebViewViewState; previous: WebViewViewState };
+
+export interface INotebookProvider {
+    /**
+     * Gets or creates a notebook, and manages the lifetime of notebooks.
+     */
+    getNotebook(server: INotebookServer, resource: Uri, options?: nbformat.INotebookMetadata): Promise<INotebook>;
+}