From 4a550ab3c95d7849c4a2e14df694c4333b81ec17 Mon Sep 17 00:00:00 2001
From: rebornix <penn.lv@gmail.com>
Date: Fri, 21 Apr 2023 10:57:19 -0700
Subject: [PATCH] Re #178715. Wait for notebook model to be ready for dispose

---
 .../notebook/common/notebookEditorModel.ts    | 12 ++++++
 .../notebookEditorModelResolverServiceImpl.ts | 43 +++++++++++++++----
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts b/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
index 9745a324fcde5..08dde80552425 100644
--- a/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
+++ b/src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
@@ -76,6 +76,18 @@ export class SimpleNotebookEditorModel extends EditorModel implements INotebookE
 		return Boolean(this._workingCopy?.model?.notebookModel);
 	}
 
+	async canDispose(): Promise<boolean> {
+		if (!this._workingCopy) {
+			return true;
+		}
+
+		if (SimpleNotebookEditorModel._isStoredFileWorkingCopy(this._workingCopy)) {
+			return this._workingCopyManager.stored.canDispose(this._workingCopy);
+		} else {
+			return true;
+		}
+	}
+
 	isDirty(): boolean {
 		return this._workingCopy?.isDirty() ?? false;
 	}
diff --git a/src/vs/workbench/contrib/notebook/common/notebookEditorModelResolverServiceImpl.ts b/src/vs/workbench/contrib/notebook/common/notebookEditorModelResolverServiceImpl.ts
index 2b6428646bf44..bb0966a1f83c5 100644
--- a/src/vs/workbench/contrib/notebook/common/notebookEditorModelResolverServiceImpl.ts
+++ b/src/vs/workbench/contrib/notebook/common/notebookEditorModelResolverServiceImpl.ts
@@ -35,6 +35,7 @@ class NotebookModelReferenceCollection extends ReferenceCollection<Promise<IReso
 
 	private readonly _dirtyStates = new ResourceMap<boolean>();
 
+	private readonly modelsToDispose = new Set<string>();
 	constructor(
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@INotebookService private readonly _notebookService: INotebookService,
@@ -56,6 +57,9 @@ class NotebookModelReferenceCollection extends ReferenceCollection<Promise<IReso
 	}
 
 	protected async createReferencedObject(key: string, viewType: string, hasAssociatedFilePath: boolean): Promise<IResolvedNotebookEditorModel> {
+		// Untrack as being disposed
+		this.modelsToDispose.delete(key);
+
 		const uri = URI.parse(key);
 
 		const workingCopyTypeId = NotebookWorkingCopyTypeIdentifier.create(viewType);
@@ -101,14 +105,37 @@ class NotebookModelReferenceCollection extends ReferenceCollection<Promise<IReso
 		return result;
 	}
 
-	protected destroyReferencedObject(_key: string, object: Promise<IResolvedNotebookEditorModel>): void {
-		object.then(model => {
-			this._modelListener.get(model)?.dispose();
-			this._modelListener.delete(model);
-			model.dispose();
-		}).catch(err => {
-			this._logService.error('FAILED to destory notebook', err);
-		});
+	protected destroyReferencedObject(key: string, object: Promise<IResolvedNotebookEditorModel>): void {
+		this.modelsToDispose.add(key);
+
+		(async () => {
+			try {
+				const model = await object;
+
+				if (!this.modelsToDispose.has(key)) {
+					// return if model has been acquired again meanwhile
+					return;
+				}
+
+				if (model instanceof SimpleNotebookEditorModel) {
+					await model.canDispose();
+				}
+
+				if (!this.modelsToDispose.has(key)) {
+					// return if model has been acquired again meanwhile
+					return;
+				}
+
+				// Finally we can dispose the model
+				this._modelListener.get(model)?.dispose();
+				this._modelListener.delete(model);
+				model.dispose();
+			} catch (err) {
+				this._logService.error('FAILED to destory notebook', err);
+			} finally {
+				this.modelsToDispose.delete(key); // Untrack as being disposed
+			}
+		})();
 	}
 }