Delegate management of notebooks (creation and disposing) to I… (microsoft#10568)

* No need to dispose * Refactor * Remove nb * Oops * Add missing imports * Skip tests * Fix linter

Author: DonJayamanne

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