Use RefCollection to track NotebookEditorInputs (#207529)

src/vs/workbench/contrib/interactive/browser/interactiveEditorInput.ts

@@ -18,7 +18,7 @@ import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { IInteractiveDocumentService } from 'vs/workbench/contrib/interactive/browser/interactiveDocumentService';
 import { IInteractiveHistoryService } from 'vs/workbench/contrib/interactive/browser/interactiveHistoryService';
 import { IResolvedNotebookEditorModel, NotebookSetting } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { ICompositeNotebookEditorInput, NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
+import { ICompositeNotebookEditorInput, INotebookEditorInputFactory, NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 
 export class InteractiveEditorInput extends EditorInput implements ICompositeNotebookEditorInput {
@@ -96,11 +96,11 @@ export class InteractiveEditorInput extends EditorInput implements ICompositeNot
 		@IInteractiveHistoryService historyService: IInteractiveHistoryService,
 		@INotebookService private readonly _notebookService: INotebookService,
 		@IFileDialogService private readonly _fileDialogService: IFileDialogService,
-		@IConfigurationService configurationService: IConfigurationService
+		@IConfigurationService configurationService: IConfigurationService,
 	) {
-		const input = NotebookEditorInput.getOrCreate(instantiationService, resource, undefined, 'interactive', {});
 		super();
 		this.isScratchpad = configurationService.getValue<boolean>(NotebookSetting.InteractiveWindowPromptToSave) !== true;
+		const input = instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory).getOrCreate(resource, undefined, 'interactive', {}));
 		this._notebookEditorInput = input;
 		this._register(this._notebookEditorInput);
 		this.name = title ?? InteractiveEditorInput.windowNames[resource.path] ?? paths.basename(resource.path, paths.extname(resource.path));

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

@@ -26,7 +26,7 @@ import { Extensions as WorkbenchExtensions, IWorkbenchContribution, IWorkbenchCo
 import { IEditorSerializer, IEditorFactoryRegistry, EditorExtensions } from 'vs/workbench/common/editor';
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { NotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookEditor';
-import { NotebookEditorInput, NotebookEditorInputOptions } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
+import { INotebookEditorInputFactory, NotebookEditorInput, NotebookEditorInputFactory, NotebookEditorInputOptions } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { NotebookService } from 'vs/workbench/contrib/notebook/browser/services/notebookServiceImpl';
 import { CellKind, CellUri, IResolvedNotebookEditorModel, NotebookWorkingCopyTypeIdentifier, NotebookSetting, ICellOutput, ICell } from 'vs/workbench/contrib/notebook/common/notebookCommon';
@@ -208,7 +208,7 @@ class NotebookEditorSerializer implements IEditorSerializer {
 			return undefined;
 		}
 
-		const input = NotebookEditorInput.getOrCreate(instantiationService, resource, preferredResource, viewType, options);
+		const input = instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory).getOrCreate(resource, preferredResource, viewType, options));
 		return input;
 	}
 }
@@ -658,7 +658,7 @@ class SimpleNotebookWorkingCopyEditorHandler extends Disposable implements IWork
 	}
 
 	createEditor(workingCopy: IWorkingCopyIdentifier): EditorInput {
-		return NotebookEditorInput.getOrCreate(this._instantiationService, workingCopy.resource, undefined, this._getViewType(workingCopy)!);
+		return this._instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory).getOrCreate(workingCopy.resource, undefined, this._getViewType(workingCopy)!));
 	}
 
 	private async _installHandler(): Promise<void> {
@@ -754,6 +754,7 @@ registerSingleton(INotebookExecutionStateService, NotebookExecutionStateService,
 registerSingleton(INotebookRendererMessagingService, NotebookRendererMessagingService, InstantiationType.Delayed);
 registerSingleton(INotebookKeymapService, NotebookKeymapService, InstantiationType.Delayed);
 registerSingleton(INotebookLoggingService, NotebookLoggingService, InstantiationType.Delayed);
+registerSingleton(INotebookEditorInputFactory, NotebookEditorInputFactory, InstantiationType.Delayed);
 
 const schemas: IJSONSchemaMap = {};
 function isConfigurationPropertySchema(x: IConfigurationPropertySchema | { [path: string]: IConfigurationPropertySchema }): x is IConfigurationPropertySchema {

src/vs/workbench/contrib/notebook/browser/services/notebookServiceImpl.ts

@@ -29,7 +29,7 @@ import { NotebookDiffEditorInput } from 'vs/workbench/contrib/notebook/common/no
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellUri, NotebookSetting, INotebookContributionData, INotebookExclusiveDocumentFilter, INotebookRendererInfo, INotebookTextModel, IOrderedMimeType, IOutputDto, MimeTypeDisplayOrder, NotebookData, NotebookEditorPriority, NotebookRendererMatch, NOTEBOOK_DISPLAY_ORDER, RENDERER_EQUIVALENT_EXTENSIONS, RENDERER_NOT_AVAILABLE, TransientOptions, NotebookExtensionDescription, INotebookStaticPreloadInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
+import { INotebookEditorInputFactory } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
 import { INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
 import { NotebookOutputRendererInfo, NotebookStaticPreloadInfo as NotebookStaticPreloadInfo } from 'vs/workbench/contrib/notebook/common/notebookOutputRenderer';
 import { NotebookEditorDescriptor, NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
@@ -196,7 +196,7 @@ export class NotebookProviderInfoStore extends Disposable {
 				}
 
 				const notebookOptions = { ...options, cellOptions } as INotebookEditorOptions;
-				const editor = NotebookEditorInput.getOrCreate(this._instantiationService, notebookUri, preferredResource, notebookProviderInfo.id);
+				const editor = this._instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory).getOrCreate(notebookUri, preferredResource, notebookProviderInfo.id));
 				return { editor, options: notebookOptions };
 			};
 
@@ -209,7 +209,7 @@ export class NotebookProviderInfoStore extends Disposable {
 					ref.dispose();
 				});
 
-				return { editor: NotebookEditorInput.getOrCreate(this._instantiationService, ref.object.resource, undefined, notebookProviderInfo.id), options };
+				return { editor: this._instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory).getOrCreate(ref.object.resource, undefined, notebookProviderInfo.id)), options };
 			};
 			const notebookDiffEditorInputFactory: DiffEditorInputFactoryFunction = ({ modified, original, label, description }) => {
 				return { editor: NotebookDiffEditorInput.create(this._instantiationService, modified.resource!, label, description, original.resource!, notebookProviderInfo.id) };

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

@@ -10,7 +10,7 @@ import { URI } from 'vs/base/common/uri';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { INotebookDiffEditorModel, IResolvedNotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { DiffEditorInput } from 'vs/workbench/common/editor/diffEditorInput';
-import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
+import { INotebookEditorInputFactory, NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 
 class NotebookDiffEditorModel extends EditorModel implements INotebookDiffEditorModel {
@@ -24,8 +24,9 @@ class NotebookDiffEditorModel extends EditorModel implements INotebookDiffEditor
 
 export class NotebookDiffEditorInput extends DiffEditorInput {
 	static create(instantiationService: IInstantiationService, resource: URI, name: string | undefined, description: string | undefined, originalResource: URI, viewType: string) {
-		const original = NotebookEditorInput.getOrCreate(instantiationService, originalResource, undefined, viewType);
-		const modified = NotebookEditorInput.getOrCreate(instantiationService, resource, undefined, viewType);
+		const factor = instantiationService.invokeFunction(accessor => accessor.get(INotebookEditorInputFactory));
+		const original = factor.getOrCreate(originalResource, undefined, viewType);
+		const modified = factor.getOrCreate(resource, undefined, viewType);
 		return instantiationService.createInstance(NotebookDiffEditorInput, name, description, original, modified, viewType);
 	}
 

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

@@ -9,10 +9,10 @@ import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { URI } from 'vs/base/common/uri';
 import { isEqual, joinPath } from 'vs/base/common/resources';
-import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { IInstantiationService, createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { IFileDialogService } from 'vs/platform/dialogs/common/dialogs';
 import { INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
-import { IDisposable, IReference } from 'vs/base/common/lifecycle';
+import { IDisposable, IReference, ReferenceCollection } from 'vs/base/common/lifecycle';
 import { CellEditType, IResolvedNotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { ILabelService } from 'vs/platform/label/common/label';
 import { Schemas } from 'vs/base/common/network';
@@ -40,6 +40,49 @@ export interface NotebookEditorInputOptions {
 	_workingCopy?: IWorkingCopyIdentifier;
 }
 
+class NotebookEditorInputReferenceCollection extends ReferenceCollection<NotebookEditorInput> {
+	private readonly _disposables = new WeakMap<NotebookEditorInput, Function>();
+	constructor(@IInstantiationService private readonly instantiationService: IInstantiationService) {
+		super();
+	}
+	protected override createReferencedObject(key: string, resource: URI, preferredResource: URI | undefined, viewType: string, options: NotebookEditorInputOptions): NotebookEditorInput {
+		const ref = this.instantiationService.createInstance(NotebookEditorInput, resource, preferredResource, viewType, options);
+		// Keep track of the disposable func, as this will be overridden later.
+		this._disposables.set(ref, ref.dispose.bind(ref));
+		return ref;
+	}
+	protected override destroyReferencedObject(key: string, object: NotebookEditorInput): void {
+		this._disposables.get(object)?.();
+		this._disposables.delete(object);
+	}
+
+}
+
+
+export const INotebookEditorInputFactory = createDecorator<INotebookEditorInputFactory>('INotebookEditorInputFactory');
+
+export interface INotebookEditorInputFactory {
+	getOrCreate(resource: URI, preferredResource: URI | undefined, viewType: string, options?: NotebookEditorInputOptions): NotebookEditorInput;
+}
+
+export class NotebookEditorInputFactory implements INotebookEditorInputFactory {
+	private readonly _data: NotebookEditorInputReferenceCollection;
+	constructor(@IInstantiationService instantiationService: IInstantiationService) {
+		this._data = instantiationService.createInstance(NotebookEditorInputReferenceCollection);
+	}
+
+	getOrCreate(resource: URI, preferredResource: URI | undefined, viewType: string, options: NotebookEditorInputOptions = {}): NotebookEditorInput {
+		const cacheId = `${resource.toString()}|${viewType}|${options._workingCopy?.typeId}`;
+		const { object, dispose } = this._data.acquire(cacheId, resource, preferredResource, viewType, options);
+		// Deref the object,
+		// Ref collection will automatically call the `dispose` on the original ref when its no longer referenced.
+		object.dispose = () => {
+			dispose();
+		};
+		return object;
+	}
+}
+
 export class NotebookEditorInput extends AbstractResourceEditorInput {
 
 	private static EditorCache: Record<string, NotebookEditorInput> = {};