Merge branch 'joh/nbRefCount2'

src/vs/base/common/lifecycle.ts

@@ -219,11 +219,11 @@ export abstract class ReferenceCollection<T> {
 
 	private readonly references: Map<string, { readonly object: T; counter: number; }> = new Map();
 
-	acquire(key: string): IReference<T> {
+	acquire(key: string, ...args: any[]): IReference<T> {
 		let reference = this.references.get(key);
 
 		if (!reference) {
-			reference = { counter: 0, object: this.createReferencedObject(key) };
+			reference = { counter: 0, object: this.createReferencedObject(key, ...args) };
 			this.references.set(key, reference);
 		}
 
@@ -240,7 +240,7 @@ export abstract class ReferenceCollection<T> {
 		return { object, dispose };
 	}
 
-	protected abstract createReferencedObject(key: string): T;
+	protected abstract createReferencedObject(key: string, ...args: any[]): T;
 	protected abstract destroyReferencedObject(key: string, object: T): void;
 }
 

src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -29,7 +29,7 @@ import { EditorInput, Extensions as EditorInputExtensions, IEditorInput, IEditor
 import { IBackupFileService } from 'vs/workbench/services/backup/common/backup';
 import { NotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookEditor';
 import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
-import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
+import { INotebookService, INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { NotebookService } from 'vs/workbench/contrib/notebook/browser/notebookServiceImpl';
 import { CellKind, CellUri, NotebookDocumentBackupData, NotebookEditorPriority } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
@@ -42,6 +42,7 @@ import { NotebookEditorOptions } from 'vs/workbench/contrib/notebook/browser/not
 import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
 import { NotebookRegistry } from 'vs/workbench/contrib/notebook/browser/notebookRegistry';
+import { NotebookModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
 
 // Editor Contribution
 
@@ -431,6 +432,7 @@ workbenchContributionsRegistry.registerWorkbenchContribution(NotebookContributio
 workbenchContributionsRegistry.registerWorkbenchContribution(CellContentProvider, LifecyclePhase.Starting);
 
 registerSingleton(INotebookService, NotebookService);
+registerSingleton(INotebookEditorModelResolverService, NotebookModelResolverService, true);
 
 const configurationRegistry = Registry.as<IConfigurationRegistry>(Extensions.Configuration);
 configurationRegistry.registerConfiguration({

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -559,6 +559,7 @@ export const NOTEBOOK_EDITOR_CURSOR_BOUNDARY = new RawContextKey<'none' | 'top'
 
 
 export interface INotebookEditorModel extends IEditorModel {
+	viewType: string;
 	notebook: NotebookTextModel;
 	isDirty(): boolean;
 	save(): Promise<boolean>;

src/vs/workbench/contrib/notebook/common/notebookEditorModelResolverService.ts

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { createDecorator, IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { URI } from 'vs/base/common/uri';
+import { INotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { NotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookEditorModel';
+import { IReference, ReferenceCollection } from 'vs/base/common/lifecycle';
+import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
+import { ILogService } from 'vs/platform/log/common/log';
+
+export const INotebookEditorModelResolverService = createDecorator<INotebookEditorModelResolverService>('INotebookModelResolverService');
+
+export interface INotebookEditorModelResolverService {
+	readonly _serviceBrand: undefined;
+	resolve(resource: URI, viewType: string, editorId?: string): Promise<IReference<INotebookEditorModel>>;
+}
+
+
+export class NotebookModelReferenceCollection extends ReferenceCollection<Promise<INotebookEditorModel>> {
+
+	constructor(
+		@IInstantiationService readonly _instantiationService: IInstantiationService,
+		@INotebookService private readonly _notebookService: INotebookService,
+		@ILogService private readonly _logService: ILogService,
+	) {
+		super();
+	}
+
+	protected createReferencedObject(key: string, ...args: any[]): Promise<INotebookEditorModel> {
+		const [viewType, editorId] = args as [string, string | undefined];
+
+		const resource = URI.parse(key);
+		const model = this._instantiationService.createInstance(NotebookEditorModel, resource, viewType);
+		const promise = model.load({ editorId });
+		return promise;
+	}
+
+	protected destroyReferencedObject(_key: string, object: Promise<INotebookEditorModel>): void {
+		object.then(model => {
+			this._notebookService.destoryNotebookDocument(model.viewType, model.notebook);
+			model.dispose();
+		}).catch(err => {
+			this._logService.critical('FAILED to destory notebook', err);
+		});
+	}
+}
+
+export class NotebookModelResolverService implements INotebookEditorModelResolverService {
+
+	readonly _serviceBrand: undefined;
+
+	private readonly _data: NotebookModelReferenceCollection;
+
+	constructor(
+		@IInstantiationService instantiationService: IInstantiationService
+	) {
+		this._data = instantiationService.createInstance(NotebookModelReferenceCollection);
+	}
+
+	async resolve(resource: URI, viewType: string, editorId?: string | undefined): Promise<IReference<INotebookEditorModel>> {
+		const reference = this._data.acquire(resource.toString(), viewType, editorId);
+		const model = await reference.object;
+		return {
+			object: model,
+			dispose() { reference.dispose(); }
+		};
+	}
+}

src/vs/workbench/contrib/notebook/common/notebookService.ts

@@ -8,11 +8,21 @@ import { URI } from 'vs/base/common/uri';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
 import { NotebookExtensionDescription } from 'vs/workbench/api/common/extHost.protocol';
 import { Event } from 'vs/base/common/event';
-import { INotebookTextModel, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto, INotebookTextModelBackup, IEditor, ICellEditOperation, NotebookCellOutputsSplice, IOrderedMimeType, IProcessedOutput } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto, INotebookTextModelBackup, IEditor, ICellEditOperation, NotebookCellOutputsSplice, IOrderedMimeType, IProcessedOutput, INotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { INotebookEditorModelManager } from 'vs/workbench/contrib/notebook/common/notebookEditorModel';
+import { IReference } from 'vs/base/common/lifecycle';
+
+
+export const INotebookEditorModelResolverService = createDecorator<INotebookEditorModelResolverService>('INotebookEditorModelResolverService');
+
+export interface INotebookEditorModelResolverService {
+	readonly _serviceBrand: undefined;
+	resolve(resource: URI, viewType: string, editorId?: string): Promise<IReference<INotebookEditorModel>>;
+}
+
 
 export const INotebookService = createDecorator<INotebookService>('notebookService');