untitled - use ref counted lifecycle (#194371)

* wip

* cleanup

* also use ref in main

* fix one test

* adop `canDispose` for untitled as well

* untitled - ignore dirty on change on revert

src/vs/workbench/api/browser/mainThreadDocuments.ts

@@ -17,7 +17,7 @@ import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/
 import { toLocalResource, extUri, IExtUri } from 'vs/base/common/resources';
 import { IWorkingCopyFileService } from 'vs/workbench/services/workingCopy/common/workingCopyFileService';
 import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity';
-import { Emitter } from 'vs/base/common/event';
+import { Emitter, Event } from 'vs/base/common/event';
 import { IPathService } from 'vs/workbench/services/path/common/pathService';
 import { ResourceMap } from 'vs/base/common/map';
 import { IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
@@ -267,15 +267,19 @@ export class MainThreadDocuments extends Disposable implements MainThreadDocumen
 	}
 
 	private async _doCreateUntitled(associatedResource?: URI, languageId?: string, initialValue?: string): Promise<URI> {
-		const model = await this._textFileService.untitled.resolve({
+		const model = this._textFileService.untitled.create({
 			associatedResource,
 			languageId,
 			initialValue
 		});
 		const resource = model.resource;
+		const ref = await this._textModelResolverService.createModelReference(resource);
 		if (!this._modelTrackers.has(resource)) {
+			ref.dispose();
 			throw new Error(`expected URI ${resource.toString()} to have come to LIFE`);
 		}
+		this._modelReferenceCollection.add(resource, ref, ref.object.textEditorModel.getValueLength());
+		Event.once(model.onDidRevert)(() => this._modelReferenceCollection.remove(resource));
 		this._proxy.$acceptDirtyStateChanged(resource, true); // mark as dirty
 		return resource;
 	}

src/vs/workbench/common/editor/textResourceEditorInput.ts

@@ -99,7 +99,7 @@ export class TextResourceEditorInput extends AbstractTextResourceEditorInput imp
 		private description: string | undefined,
 		private preferredLanguageId: string | undefined,
 		private preferredContents: string | undefined,
-		@ITextModelService private readonly textModelResolverService: ITextModelService,
+		@ITextModelService private readonly textModelService: ITextModelService,
 		@ITextFileService textFileService: ITextFileService,
 		@IEditorService editorService: IEditorService,
 		@IFileService fileService: IFileService,
@@ -160,7 +160,7 @@ export class TextResourceEditorInput extends AbstractTextResourceEditorInput imp
 		this.preferredLanguageId = undefined;
 
 		if (!this.modelReference) {
-			this.modelReference = this.textModelResolverService.createModelReference(this.resource);
+			this.modelReference = this.textModelService.createModelReference(this.resource);
 		}
 
 		const ref = await this.modelReference;

src/vs/workbench/contrib/files/browser/editors/fileEditorInput.ts

@@ -93,7 +93,7 @@ export class FileEditorInput extends AbstractTextResourceEditorInput implements
 		preferredContents: string | undefined,
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@ITextFileService textFileService: ITextFileService,
-		@ITextModelService private readonly textModelResolverService: ITextModelService,
+		@ITextModelService private readonly textModelService: ITextModelService,
 		@ILabelService labelService: ILabelService,
 		@IFileService fileService: IFileService,
 		@IFilesConfigurationService filesConfigurationService: IFilesConfigurationService,
@@ -348,7 +348,7 @@ export class FileEditorInput extends AbstractTextResourceEditorInput implements
 			// resolve() ensures we are not creating model references for these kind of resources.
 			// In addition we have a bit of payload to take into account (encoding, reload) that the text resolver does not handle yet.
 			if (!this.cachedTextFileModelReference) {
-				this.cachedTextFileModelReference = await this.textModelResolverService.createModelReference(this.resource) as IReference<ITextFileEditorModel>;
+				this.cachedTextFileModelReference = await this.textModelService.createModelReference(this.resource) as IReference<ITextFileEditorModel>;
 			}
 
 			const model = this.cachedTextFileModelReference.object;

src/vs/workbench/services/textfile/common/textEditorService.ts

@@ -146,20 +146,7 @@ export class TextEditorService extends Disposable implements ITextEditorService
 				untitledModel = this.untitledTextEditorService.create({ associatedResource: untitledInput.resource, ...untitledOptions });
 			}
 
-			return this.createOrGetCached(untitledModel.resource, () => {
-
-				// Factory function for new untitled editor
-				const input = this.instantiationService.createInstance(UntitledTextEditorInput, untitledModel);
-
-				// We dispose the untitled model once the editor
-				// is being disposed. Even though we may have not
-				// created the model initially, the lifecycle for
-				// untitled is tightly coupled with the editor
-				// lifecycle for now.
-				Event.once(input.onWillDispose)(() => untitledModel.dispose());
-
-				return input;
-			});
+			return this.createOrGetCached(untitledModel.resource, () => this.instantiationService.createInstance(UntitledTextEditorInput, untitledModel));
 		}
 
 		// Text File/Resource Editor Support

src/vs/workbench/services/textmodelResolver/common/textModelResolverService.ts

@@ -18,6 +18,7 @@ import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
 import { ModelUndoRedoParticipant } from 'vs/editor/common/services/modelUndoRedoParticipant';
 import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity';
+import { UntitledTextEditorModel } from 'vs/workbench/services/untitled/common/untitledTextEditorModel';
 
 class ResourceModelCollection extends ReferenceCollection<Promise<IResolvedTextEditorModel>> {
 
@@ -102,9 +103,9 @@ class ResourceModelCollection extends ReferenceCollection<Promise<IResolvedTextE
 
 	protected destroyReferencedObject(key: string, modelPromise: Promise<ITextEditorModel>): void {
 
-		// untitled and inMemory are bound to a different lifecycle
+		// inMemory is bound to a different lifecycle
 		const resource = URI.parse(key);
-		if (resource.scheme === Schemas.untitled || resource.scheme === Schemas.inMemory) {
+		if (resource.scheme === Schemas.inMemory) {
 			return;
 		}
 
@@ -125,6 +126,10 @@ class ResourceModelCollection extends ReferenceCollection<Promise<IResolvedTextE
 					// text file models have conditions that prevent them
 					// from dispose, so we have to wait until we can dispose
 					await this.textFileService.files.canDispose(model);
+				} else if (model instanceof UntitledTextEditorModel) {
+					// untitled file models have conditions that prevent them
+					// from dispose, so we have to wait until we can dispose
+					await this.textFileService.untitled.canDispose(model);
 				}
 
 				if (!this.modelsToDispose.has(key)) {

src/vs/workbench/services/untitled/common/untitledTextEditorInput.ts

@@ -17,6 +17,8 @@ import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/
 import { IPathService } from 'vs/workbench/services/path/common/pathService';
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { IFilesConfigurationService } from 'vs/workbench/services/filesConfiguration/common/filesConfigurationService';
+import { ITextModelService } from 'vs/editor/common/services/resolverService';
+import { dispose, IReference } from 'vs/base/common/lifecycle';
 
 /**
  * An editor input to be used for untitled text buffers.
@@ -34,6 +36,7 @@ export class UntitledTextEditorInput extends AbstractTextResourceEditorInput imp
 	}
 
 	private modelResolve: Promise<void> | undefined = undefined;
+	private cachedUntitledTextEditorModelReference: IReference<IUntitledTextEditorModel> | undefined = undefined;
 
 	constructor(
 		readonly model: IUntitledTextEditorModel,
@@ -43,7 +46,8 @@ export class UntitledTextEditorInput extends AbstractTextResourceEditorInput imp
 		@IFileService fileService: IFileService,
 		@IWorkbenchEnvironmentService private readonly environmentService: IWorkbenchEnvironmentService,
 		@IPathService private readonly pathService: IPathService,
-		@IFilesConfigurationService filesConfigurationService: IFilesConfigurationService
+		@IFilesConfigurationService filesConfigurationService: IFilesConfigurationService,
+		@ITextModelService private readonly textModelService: ITextModelService
 	) {
 		super(model.resource, undefined, editorService, textFileService, labelService, fileService, filesConfigurationService);
 
@@ -120,11 +124,22 @@ export class UntitledTextEditorInput extends AbstractTextResourceEditorInput imp
 
 	override async resolve(): Promise<IUntitledTextEditorModel> {
 		if (!this.modelResolve) {
-			this.modelResolve = this.model.resolve();
+			this.modelResolve = (async () => {
+
+				// Acquire a model reference
+				this.cachedUntitledTextEditorModelReference = await this.textModelService.createModelReference(this.resource) as IReference<IUntitledTextEditorModel>;
+			})();
 		}
 
 		await this.modelResolve;
 
+		// It is possible that this input was disposed before the model
+		// finished resolving. As such, we need to make sure to dispose
+		// the model reference to not leak it.
+		if (this.isDisposed()) {
+			this.disposeModelReference();
+		}
+
 		return this.model;
 	}
 
@@ -176,8 +191,18 @@ export class UntitledTextEditorInput extends AbstractTextResourceEditorInput imp
 	}
 
 	override dispose(): void {
+
+		// Model
 		this.modelResolve = undefined;
 
+		// Model reference
+		this.disposeModelReference();
+
 		super.dispose();
 	}
+
+	private disposeModelReference(): void {
+		dispose(this.cachedUntitledTextEditorModelReference);
+		this.cachedUntitledTextEditorModelReference = undefined;
+	}
 }

src/vs/workbench/services/untitled/common/untitledTextEditorModel.ts

@@ -12,7 +12,7 @@ import { Event, Emitter } from 'vs/base/common/event';
 import { IWorkingCopyBackupService } from 'vs/workbench/services/workingCopy/common/workingCopyBackup';
 import { ITextResourceConfigurationChangeEvent, ITextResourceConfigurationService } from 'vs/editor/common/services/textResourceConfiguration';
 import { ITextModel } from 'vs/editor/common/model';
-import { createTextBufferFactoryFromStream } from 'vs/editor/common/model/textModel';
+import { createTextBufferFactory, createTextBufferFactoryFromStream } from 'vs/editor/common/model/textModel';
 import { ITextEditorModel } from 'vs/editor/common/services/resolverService';
 import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
 import { IWorkingCopy, WorkingCopyCapabilities, IWorkingCopyBackup, NO_TYPE_ID, IWorkingCopySaveEvent } from 'vs/workbench/services/workingCopy/common/workingCopy';
@@ -273,16 +273,19 @@ export class UntitledTextEditorModel extends BaseTextEditorModel implements IUnt
 
 	async revert(): Promise<void> {
 
+		// Reset contents to be empty
+		this.ignoreDirtyOnModelContentChange = true;
+		try {
+			this.updateTextEditorModel(createTextBufferFactory(''));
+		} finally {
+			this.ignoreDirtyOnModelContentChange = false;
+		}
+
 		// No longer dirty
 		this.setDirty(false);
 
 		// Emit as event
 		this._onDidRevert.fire();
-
-		// A reverted untitled model is invalid because it has
-		// no actual source on disk to revert to. As such we
-		// dispose the model.
-		this.dispose();
 	}
 
 	async backup(token: CancellationToken): Promise<IWorkingCopyBackup> {
@@ -307,6 +310,8 @@ export class UntitledTextEditorModel extends BaseTextEditorModel implements IUnt
 
 	//#region Resolve
 
+	private ignoreDirtyOnModelContentChange = false;
+
 	override async resolve(): Promise<void> {
 
 		// Create text editor model if not yet done
@@ -375,16 +380,18 @@ export class UntitledTextEditorModel extends BaseTextEditorModel implements IUnt
 	}
 
 	private onModelContentChanged(textEditorModel: ITextModel, e: IModelContentChangedEvent): void {
+		if (!this.ignoreDirtyOnModelContentChange) {
 
-		// mark the untitled text editor as non-dirty once its content becomes empty and we do
-		// not have an associated path set. we never want dirty indicator in that case.
-		if (!this.hasAssociatedFilePath && textEditorModel.getLineCount() === 1 && textEditorModel.getLineLength(1) === 0) {
-			this.setDirty(false);
-		}
+			// mark the untitled text editor as non-dirty once its content becomes empty and we do
+			// not have an associated path set. we never want dirty indicator in that case.
+			if (!this.hasAssociatedFilePath && textEditorModel.getLineCount() === 1 && textEditorModel.getLineLength(1) === 0) {
+				this.setDirty(false);
+			}
 
-		// turn dirty otherwise
-		else {
-			this.setDirty(true);
+			// turn dirty otherwise
+			else {
+				this.setDirty(true);
+			}
 		}
 
 		// Check for name change if first line changed in the range of 0-FIRST_LINE_NAME_CANDIDATE_MAX_LENGTH columns

src/vs/workbench/services/untitled/common/untitledTextEditorService.ts

@@ -117,6 +117,13 @@ export interface IUntitledTextEditorModelManager {
 	 * Figures out if the given resource has an associated resource or not.
 	 */
 	isUntitledWithAssociatedResource(resource: URI): boolean;
+
+	/**
+	 * Waits for the model to be ready to be disposed. There may be conditions
+	 * under which the model cannot be disposed, e.g. when it is dirty. Once the
+	 * promise is settled, it is safe to dispose the model.
+	 */
+	canDispose(model: IUntitledTextEditorModel): true | Promise<true>;
 }
 
 export interface IUntitledTextEditorService extends IUntitledTextEditorModelManager {
@@ -270,6 +277,29 @@ export class UntitledTextEditorService extends Disposable implements IUntitledTe
 	isUntitledWithAssociatedResource(resource: URI): boolean {
 		return resource.scheme === Schemas.untitled && resource.path.length > 1 && !UntitledTextEditorService.UNTITLED_WITHOUT_ASSOCIATED_RESOURCE_REGEX.test(resource.path);
 	}
+
+	canDispose(model: UntitledTextEditorModel): true | Promise<true> {
+		if (model.isDisposed()) {
+			return true; // quick return if model already disposed
+		}
+
+		// promise based return in all other cases
+		return this.doCanDispose(model);
+	}
+
+	private async doCanDispose(model: UntitledTextEditorModel): Promise<true> {
+
+		// dirty model: we do not allow to dispose dirty models to prevent
+		// data loss cases. dirty models can only be disposed when they are
+		// either saved or reverted
+		if (model.isDirty()) {
+			await Event.toPromise(model.onDidChangeDirty);
+
+			return this.canDispose(model);
+		}
+
+		return true;
+	}
 }
 
 registerSingleton(IUntitledTextEditorService, UntitledTextEditorService, InstantiationType.Delayed);

src/vs/workbench/services/untitled/test/browser/untitledTextEditor.test.ts

@@ -14,14 +14,15 @@ import { PLAINTEXT_LANGUAGE_ID } from 'vs/editor/common/languages/modesRegistry'
 import { ISingleEditOperation } from 'vs/editor/common/core/editOperation';
 import { Range } from 'vs/editor/common/core/range';
 import { UntitledTextEditorInput } from 'vs/workbench/services/untitled/common/untitledTextEditorInput';
-import { IUntitledTextEditorModel } from 'vs/workbench/services/untitled/common/untitledTextEditorModel';
+import { IUntitledTextEditorModel, UntitledTextEditorModel } from 'vs/workbench/services/untitled/common/untitledTextEditorModel';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { EditorInputCapabilities } from 'vs/workbench/common/editor';
 import { DisposableStore } from 'vs/base/common/lifecycle';
 import { isReadable, isReadableStream } from 'vs/base/common/stream';
 import { readableToBuffer, streamToBuffer, VSBufferReadable, VSBufferReadableStream } from 'vs/base/common/buffer';
 import { LanguageDetectionLanguageEventSource } from 'vs/workbench/services/languageDetection/common/languageDetectionWorkerService';
 import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
+import { timeout } from 'vs/base/common/async';
 
 suite('Untitled text editors', () => {
 
@@ -573,5 +574,32 @@ suite('Untitled text editors', () => {
 		assert.strictEqual(counter, 1, 'Another change to same encoding does not fire event');
 	});
 
+	test('canDispose with dirty model', async function () {
+		const service = accessor.untitledTextEditorService;
+		const input = disposables.add(instantiationService.createInstance(UntitledTextEditorInput, service.create()));
+
+		const model = disposables.add(await input.resolve());
+
+		model.textEditorModel?.setValue('foo');
+
+		const canDisposePromise = service.canDispose(model as UntitledTextEditorModel);
+		assert.ok(canDisposePromise instanceof Promise);
+
+		let canDispose = false;
+		(async () => {
+			canDispose = await canDisposePromise;
+		})();
+
+		assert.strictEqual(canDispose, false);
+		model.revert({ soft: true });
+
+		await timeout(0);
+
+		assert.strictEqual(canDispose, true);
+
+		const canDispose2 = service.canDispose(model as UntitledTextEditorModel);
+		assert.strictEqual(canDispose2, true);
+	});
+
 	ensureNoDisposablesAreLeakedInTestSuite();
 });