Add UnownedDisposable to make it clearer that webviewEditorInput takes ownership of the webview passed to it

mjbvz · mjbvz · commit b893fd1e26eb · 2019-07-19T17:36:17.000-07:00
diff --git a/src/vs/base/common/lifecycle.ts b/src/vs/base/common/lifecycle.ts
@@ -207,6 +207,43 @@ export class MutableDisposable<T extends IDisposable> implements IDisposable {
 	}
 }
 
+/**
+ * Wrapper class that stores a disposable that is not currently "owned" by anyone.
+ *
+ * Example use cases:
+ *
+ * - Express that a function/method will take ownership of a disposable parameter.
+ * - Express that a function returns a disposable that the caller must explicitly take ownership of.
+ */
+export class UnownedDisposable<T extends IDisposable> extends Disposable {
+	private _hasBeenAcquired = false;
+	private _value?: T;
+
+	public constructor(value: T) {
+		super();
+		this._value = value;
+	}
+
+	public acquire(): T {
+		if (this._hasBeenAcquired) {
+			throw new Error('This disposable has already been acquired');
+		}
+		this._hasBeenAcquired = true;
+		const value = this._value!;
+		this._value = undefined;
+		return value;
+	}
+
+	public dispose() {
+		super.dispose();
+		if (!this._hasBeenAcquired) {
+			this._hasBeenAcquired = true;
+			this._value!.dispose();
+			this._value = undefined;
+		}
+	}
+}
+
 export interface IReference<T> extends IDisposable {
 	readonly object: T;
 }
diff --git a/src/vs/workbench/contrib/webview/browser/webviewEditorInput.ts b/src/vs/workbench/contrib/webview/browser/webviewEditorInput.ts
@@ -9,6 +9,7 @@ import { IEditorModel } from 'vs/platform/editor/common/editor';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { EditorInput, EditorModel, GroupIdentifier, IEditorInput } from 'vs/workbench/common/editor';
 import { WebviewEditorOverlay } from 'vs/workbench/contrib/webview/common/webview';
+import { UnownedDisposable as Unowned } from 'vs/base/common/lifecycle';
 
 class WebviewIconsManager {
 	private readonly _icons = new Map<string, { light: URI, dark: URI }>();
@@ -58,6 +59,7 @@ export class WebviewEditorInput extends EditorInput {
 	private _name: string;
 	private _iconPath?: { light: URI, dark: URI };
 	private _group?: GroupIdentifier;
+	private readonly _webview: WebviewEditorOverlay;
 
 	constructor(
 		public readonly id: string,
@@ -67,14 +69,14 @@ export class WebviewEditorInput extends EditorInput {
 			readonly location: URI;
 			readonly id: ExtensionIdentifier;
 		},
-		public readonly webview: WebviewEditorOverlay,
+		webview: Unowned<WebviewEditorOverlay>,
 	) {
 		super();
 
 		this._name = name;
 		this.extension = extension;
 
-		this._register(webview); // The input owns this webview
+		this._webview = this._register(webview.acquire()); // The input owns this webview
 	}
 
 	public getTypeId(): string {
@@ -105,6 +107,10 @@ export class WebviewEditorInput extends EditorInput {
 		this._onDidChangeLabel.fire();
 	}
 
+	public get webview() {
+		return this._webview;
+	}
+
 	public get iconPath() {
 		return this._iconPath;
 	}
@@ -147,7 +153,7 @@ export class RevivedWebviewEditorInput extends WebviewEditorInput {
 			readonly id: ExtensionIdentifier
 		},
 		private readonly reviver: (input: WebviewEditorInput) => Promise<void>,
-		webview: WebviewEditorOverlay,
+		webview: Unowned<WebviewEditorOverlay>,
 	) {
 		super(id, viewType, name, extension, webview);
 	}
diff --git a/src/vs/workbench/contrib/webview/browser/webviewEditorService.ts b/src/vs/workbench/contrib/webview/browser/webviewEditorService.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { equals } from 'vs/base/common/arrays';
-import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
+import { IDisposable, toDisposable, UnownedDisposable } from 'vs/base/common/lifecycle';
 import { values } from 'vs/base/common/map';
 import { URI } from 'vs/base/common/uri';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
@@ -144,7 +144,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 	): WebviewEditorInput {
 		const webview = this.createWebiew(id, extension, options);
 
-		const webviewInput = this._instantiationService.createInstance(WebviewEditorInput, id, viewType, title, extension, webview);
+		const webviewInput = this._instantiationService.createInstance(WebviewEditorInput, id, viewType, title, extension, new UnownedDisposable(webview));
 		this._editorService.openEditor(webviewInput, { pinned: true, preserveFocus: showOptions.preserveFocus }, showOptions.group);
 		return webviewInput;
 	}
@@ -191,7 +191,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 			const promise = new Promise<void>(r => { resolve = r; });
 			this._revivalPool.add(webview, resolve!);
 			return promise;
-		}, webview);
+		}, new UnownedDisposable(webview));
 
 		webviewInput.iconPath = iconPath;
 
