Re-applying Cleaning up reviver handling for webviews

Reverting the revert of 09d4f006

src/vs/workbench/api/electron-browser/mainThreadWebview.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 import * as map from 'vs/base/common/map';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { ICodeEditorService } from 'vs/editor/browser/services/codeEditorService';
@@ -17,8 +17,8 @@ import { ExtHostContext, ExtHostWebviewsShape, IExtHostContext, MainContext, Mai
 import { editorGroupToViewColumn, EditorViewColumn, viewColumnToEditorGroup } from 'vs/workbench/api/shared/editor';
 import { CodeInsetController } from 'vs/workbench/contrib/codeinset/electron-browser/codeInset.contribution';
 import { WebviewEditor } from 'vs/workbench/contrib/webview/electron-browser/webviewEditor';
-import { RevivedWebviewEditorInput, WebviewEditorInput } from 'vs/workbench/contrib/webview/electron-browser/webviewEditorInput';
-import { ICreateWebViewShowOptions, IWebviewEditorService, WebviewInputOptions, WebviewReviver } from 'vs/workbench/contrib/webview/electron-browser/webviewEditorService';
+import { WebviewEditorInput } from 'vs/workbench/contrib/webview/electron-browser/webviewEditorInput';
+import { ICreateWebViewShowOptions, IWebviewEditorService, WebviewInputOptions } from 'vs/workbench/contrib/webview/electron-browser/webviewEditorService';
 import { WebviewElement } from 'vs/workbench/contrib/webview/electron-browser/webviewElement';
 import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { ACTIVE_GROUP, IEditorService } from 'vs/workbench/services/editor/common/editorService';
@@ -28,9 +28,7 @@ import * as vscode from 'vscode';
 import { extHostNamedCustomer } from './extHostCustomers';
 
 @extHostNamedCustomer(MainContext.MainThreadWebviews)
-export class MainThreadWebviews extends Disposable implements MainThreadWebviewsShape, WebviewReviver {
-
-	private static readonly viewType = 'mainThreadWebview';
+export class MainThreadWebviews extends Disposable implements MainThreadWebviewsShape {
 
 	private static readonly standardSupportedLinkSchemes = ['http', 'https', 'mailto'];
 
@@ -40,18 +38,18 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 	private readonly _proxy: ExtHostWebviewsShape;
 	private readonly _webviews = new Map<WebviewPanelHandle, WebviewEditorInput>();
 	private readonly _webviewsElements = new Map<WebviewInsetHandle, WebviewElement>();
-	private readonly _revivers = new Set<string>();
+	private readonly _revivers = new Map<string, IDisposable>();
 
 	private _activeWebview: WebviewPanelHandle | undefined = undefined;
 
 	constructor(
 		context: IExtHostContext,
 		@ILifecycleService lifecycleService: ILifecycleService,
+		@IExtensionService extensionService: IExtensionService,
 		@IEditorGroupsService private readonly _editorGroupService: IEditorGroupsService,
 		@IEditorService private readonly _editorService: IEditorService,
 		@IWebviewEditorService private readonly _webviewService: IWebviewEditorService,
 		@IOpenerService private readonly _openerService: IOpenerService,
-		@IExtensionService private readonly _extensionService: IExtensionService,
 		@ITelemetryService private readonly _telemetryService: ITelemetryService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@ICodeEditorService private readonly _codeEditorService: ICodeEditorService,
@@ -63,7 +61,18 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 		_editorService.onDidActiveEditorChange(this.onActiveEditorChanged, this, this._toDispose);
 		_editorService.onDidVisibleEditorsChange(this.onVisibleEditorsChanged, this, this._toDispose);
 
-		this._toDispose.push(_webviewService.registerReviver(this));
+		// This reviver's only job is to activate webview extensions
+		// This should trigger the real reviver to be registered from the extension host side.
+		this._toDispose.push(_webviewService.registerReviver({
+			canRevive: (webview) => {
+				const viewType = webview.state.viewType;
+				if (viewType) {
+					extensionService.activateByEvent(`onWebviewPanel:${viewType}`);
+				}
+				return false;
+			},
+			reviveWebview: () => { throw new Error('not implemented'); }
+		}));
 
 		lifecycleService.onBeforeShutdown(e => {
 			e.veto(this._onBeforeShutdown());
@@ -85,7 +94,7 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 			mainThreadShowOptions.group = viewColumnToEditorGroup(this._editorGroupService, showOptions.viewColumn);
 		}
 
-		const webview = this._webviewService.createWebview(MainThreadWebviews.viewType, title, mainThreadShowOptions, reviveWebviewOptions(options), URI.revive(extensionLocation), this.createWebviewEventDelegate(handle));
+		const webview = this._webviewService.createWebview(this.getInternalWebviewId(viewType), title, mainThreadShowOptions, reviveWebviewOptions(options), URI.revive(extensionLocation), this.createWebviewEventDelegate(handle));
 		webview.state = {
 			viewType: viewType,
 			state: undefined
@@ -206,50 +215,55 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 	}
 
 	public $registerSerializer(viewType: string): void {
-		this._revivers.add(viewType);
-	}
-
-	public $unregisterSerializer(viewType: string): void {
-		this._revivers.delete(viewType);
-	}
-
-	public reviveWebview(webview: WebviewEditorInput): Promise<void> {
-		const viewType = webview.state.viewType;
-		return Promise.resolve(this._extensionService.activateByEvent(`onWebviewPanel:${viewType}`).then(() => {
-			const handle = 'revival-' + MainThreadWebviews.revivalPool++;
-			this._webviews.set(handle, webview);
-			webview._events = this.createWebviewEventDelegate(handle);
-
-			let state = undefined;
-			if (webview.state.state) {
-				try {
-					state = JSON.parse(webview.state.state);
-				} catch {
-					// noop
+		if (this._revivers.has(viewType)) {
+			throw new Error(`Reviver for ${viewType} already registered`);
+		}
+		this._revivers.set(viewType, this._webviewService.registerReviver({
+			canRevive: (webview) => {
+				return !webview.isDisposed() && webview.state && webview.state.viewType === viewType;
+			},
+			reviveWebview: async (webview): Promise<void> => {
+				const viewType = webview.state.viewType;
+				const handle = 'revival-' + MainThreadWebviews.revivalPool++;
+				this._webviews.set(handle, webview);
+				webview._events = this.createWebviewEventDelegate(handle);
+				let state = undefined;
+				if (webview.state.state) {
+					try {
+						state = JSON.parse(webview.state.state);
+					} catch {
+						// noop
+					}
 				}
-			}
 
-			return this._proxy.$deserializeWebviewPanel(handle, webview.state.viewType, webview.getTitle(), state, editorGroupToViewColumn(this._editorGroupService, webview.group || ACTIVE_GROUP), webview.options)
-				.then(undefined, error => {
+				try {
+					await this._proxy.$deserializeWebviewPanel(handle, webview.state.viewType, webview.getTitle(), state, editorGroupToViewColumn(this._editorGroupService, webview.group || ACTIVE_GROUP), webview.options);
+				} catch (error) {
 					onUnexpectedError(error);
-
 					webview.html = MainThreadWebviews.getDeserializationFailedContents(viewType);
-				});
+				}
+			}
 		}));
 	}
 
-	public canRevive(webview: WebviewEditorInput): boolean {
-		if (webview.isDisposed() || !webview.state || webview.viewType !== MainThreadWebviews.viewType) {
-			return false;
+	public $unregisterSerializer(viewType: string): void {
+		const reviver = this._revivers.get(viewType);
+		if (!reviver) {
+			throw new Error(`No reviver for ${viewType} registered`);
 		}
 
-		return this._revivers.has(webview.state.viewType) || !!(webview as RevivedWebviewEditorInput).reviver;
+		reviver.dispose();
+		this._revivers.delete(viewType);
+	}
+
+	private getInternalWebviewId(viewType: string): string {
+		return `mainThreadWebview-${viewType}`;
 	}
 
 	private _onBeforeShutdown(): boolean {
-		this._webviews.forEach((view) => {
-			if (this.canRevive(view)) {
-				view.state.state = view.webviewState;
+		this._webviews.forEach((webview) => {
+			if (!webview.isDisposed() && webview.state && this._revivers.has(webview.state.viewType)) {
+				webview.state.state = webview.webviewState;
 			}
 		});
 		return false; // Don't veto shutdown

src/vs/workbench/contrib/webview/electron-browser/webviewEditorService.ts

@@ -92,7 +92,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 	_serviceBrand: any;
 
 	private readonly _revivers = new Set<WebviewReviver>();
-	private _awaitingRevival: { input: WebviewEditorInput, resolve: () => void }[] = [];
+	private _awaitingRevival: Array<{ input: WebviewEditorInput, resolve: () => void }> = [];
 
 	constructor(
 		@IEditorService private readonly _editorService: IEditorService,