Fix race when loading notebook webviews (microsoft#162472) (microsoft#162694)

* Fix race when loading notebook webviews

This tries to fix microsoft#161520 by doing the following:

- Wait for basic initialization to finish before posting messages
- `initializedMarkup` should wait  for the correct response instead of the first response

This fixes the issue in my testing, but causes the layout to shift around during load

* wait for js preload to be initialized in webview

* Removed pending message queue

* Remove extra async

Co-authored-by: rebornix <[email protected]>

Co-authored-by: rebornix <[email protected]>

Author: mjbvz

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -47,6 +47,7 @@ import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editor
 import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService';
 import { FromWebviewMessage, IAckOutputHeight, IClickedDataUrlMessage, ICodeBlockHighlightRequest, IContentWidgetTopRequest, IControllerPreload, ICreationContent, ICreationRequestMessage, IFindMatch, IMarkupCellInitialization, RendererMetadata, ToWebviewMessage } from './webviewMessages';
 import { compressOutputItemStreams } from 'vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor';
+import { DeferredPromise } from 'vs/base/common/async';
 
 export interface ICachedInset<K extends ICommonCellInfo> {
 	outputId: string;
@@ -110,10 +111,14 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Disposable {
 	private readonly _onMessage = this._register(new Emitter<INotebookWebviewMessage>());
 	private readonly _preloadsCache = new Set<string>();
 	public readonly onMessage: Event<INotebookWebviewMessage> = this._onMessage.event;
-	private _initialized?: Promise<void>;
 	private _disposed = false;
 	private _currentKernel?: INotebookKernel;
 
+	private _initialized?: DeferredPromise<void>;
+	private _webviewPreloadInitialized?: DeferredPromise<void>;
+	private firstInit = true;
+	private initializeMarkupPromise?: { readonly requestId: string; readonly p: DeferredPromise<void>; readonly isFirstInit: boolean };
+
 	private readonly nonce = UUID.generateUuid();
 
 	constructor(
@@ -452,11 +457,9 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Disposable {
 		}
 
 		let coreDependencies = '';
-		let resolveFunc: () => void;
 
-		this._initialized = new Promise<void>((resolve) => {
-			resolveFunc = resolve;
-		});
+		this._initialized = new DeferredPromise();
+		this._webviewPreloadInitialized = new DeferredPromise();
 
 		if (!isWeb) {
 			const loaderUri = FileAccess.asFileUri('vs/loader.js', require);
@@ -469,7 +472,7 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Disposable {
 			</script>`;
 			const htmlContent = this.generateContent(coreDependencies, baseUrl.toString());
 			this._initialize(htmlContent);
-			resolveFunc!();
+			this._initialized.complete();
 		} else {
 			const loaderUri = FileAccess.asBrowserUri('vs/loader.js', require);
 
@@ -493,16 +496,16 @@ var requirejs = (function() {
 
 				const htmlContent = this.generateContent(coreDependencies, baseUrl.toString());
 				this._initialize(htmlContent);
-				resolveFunc!();
+				this._initialized!.complete();
 			}, error => {
 				// the fetch request is rejected
 				const htmlContent = this.generateContent(coreDependencies, baseUrl.toString());
 				this._initialize(htmlContent);
-				resolveFunc!();
+				this._initialized!.complete();
 			});
 		}
 
-		await this._initialized;
+		await this._initialized.p;
 	}
 
 	private getNotebookBaseUri() {
@@ -590,9 +593,17 @@ var requirejs = (function() {
 
 			switch (data.type) {
 				case 'initialized': {
+					this._webviewPreloadInitialized?.complete();
 					this.initializeWebViewState();
 					break;
 				}
+				case 'initializedMarkup': {
+					if (this.initializeMarkupPromise?.requestId === data.requestId) {
+						this.initializeMarkupPromise?.p.complete();
+						this.initializeMarkupPromise = undefined;
+					}
+					break;
+				}
 				case 'dimension': {
 					for (const update of data.updates) {
 						const height = update.height;
@@ -932,8 +943,6 @@ var requirejs = (function() {
 		];
 	}
 
-	private firstInit = true;
-
 	private initializeWebViewState() {
 		this._preloadsCache.clear();
 		if (this._currentKernel) {
@@ -944,9 +953,9 @@ var requirejs = (function() {
 			this._sendMessageToWebview({ ...inset.cachedCreation, initiallyHidden: this.hiddenInsetMapping.has(output) });
 		}
 
-		if (this.firstInit) {
+		if (this.initializeMarkupPromise?.isFirstInit) {
 			// On first run the contents have already been initialized so we don't need to init them again
-			this.firstInit = false;
+			// no op
 		} else {
 			const mdCells = [...this.markupPreviewMapping.values()];
 			this.markupPreviewMapping.clear();
@@ -1160,15 +1169,16 @@ var requirejs = (function() {
 			return;
 		}
 
-		// TODO: use proper handler
-		const p = new Promise<void>(resolve => {
-			const sub = this.webview?.onMessage(e => {
-				if (e.message.type === 'initializedMarkup') {
-					resolve();
-					sub?.dispose();
-				}
-			});
-		});
+		this.initializeMarkupPromise?.p.complete();
+		const requestId = UUID.generateUuid();
+		this.initializeMarkupPromise = { p: new DeferredPromise(), requestId, isFirstInit: this.firstInit };
+
+		if (this._webviewPreloadInitialized) {
+			// wait for webview preload script module to be loaded
+			await this._webviewPreloadInitialized.p;
+		}
+
+		this.firstInit = false;
 
 		for (const cell of cells) {
 			this.markupPreviewMapping.set(cell.cellId, cell);
@@ -1177,9 +1187,10 @@ var requirejs = (function() {
 		this._sendMessageToWebview({
 			type: 'initializeMarkup',
 			cells,
+			requestId,
 		});
 
-		await p;
+		return this.initializeMarkupPromise.p.p;
 	}
 
 	/**

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewMessages.ts

@@ -138,6 +138,7 @@ export interface ICellDragEndMessage extends BaseToWebviewMessage {
 
 export interface IInitializedMarkupMessage extends BaseToWebviewMessage {
 	readonly type: 'initializedMarkup';
+	readonly requestId: string;
 }
 
 export interface ICodeBlockHighlightRequest {
@@ -351,6 +352,7 @@ export interface IMarkupCellInitialization {
 export interface IInitializeMarkupCells {
 	readonly type: 'initializeMarkup';
 	readonly cells: readonly IMarkupCellInitialization[];
+	readonly requestId: string;
 }
 
 export interface INotebookStylesMessage {

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads.ts

@@ -1048,7 +1048,7 @@ async function webviewPreloads(ctx: PreloadContext) {
 					await Promise.all(event.data.cells.map(info => viewModel.ensureMarkupCell(info)));
 				} finally {
 					dimensionUpdater.updateImmediately();
-					postNotebookMessage('initializedMarkup', {});
+					postNotebookMessage('initializedMarkup', { requestId: event.data.requestId });
 				}
 				break;
 			}