Batch messages on the exthost to not freeze ads (#8949)

* batch messages on the exthost to not freeze ads

* clear out messages on query complete

src/sql/base/common/map.ts

@@ -62,3 +62,26 @@ export class ReverseLookUpMap<K, V> {
 		return this.forward.size;
 	}
 }
+
+/**
+ * ASSUMES THAT THE VALUES ARE ALREADY SERIALIZABLE
+ */
+export function mapToSerializable<T>(map: Map<string, T>): [string, T][] {
+	const serializable: [string, T][] = [];
+
+	map.forEach((value, key) => {
+		serializable.push([key, value]);
+	});
+
+	return serializable;
+}
+
+export function serializableToMap<T>(serializable: [string, T][]): Map<string, T> {
+	const items = new Map<string, T>();
+
+	for (const [key, value] of serializable) {
+		items.set(key, value);
+	}
+
+	return items;
+}

src/sql/workbench/api/browser/mainThreadDataProtocol.ts

@@ -26,6 +26,7 @@ import { IFileBrowserService } from 'sql/workbench/services/fileBrowser/common/i
 import { IExtHostContext } from 'vs/workbench/api/common/extHost.protocol';
 import { extHostNamedCustomer } from 'vs/workbench/api/common/extHostCustomers';
 import { assign } from 'vs/base/common/objects';
+import { serializableToMap } from 'sql/base/common/map';
 
 /**
  * Main thread class for handling data protocol management registration.
@@ -500,8 +501,8 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData
 	public $onResultSetUpdated(handle: number, resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void {
 		this._queryManagementService.onResultSetUpdated(resultSetInfo);
 	}
-	public $onQueryMessage(handle: number, message: azdata.QueryExecuteMessageParams): void {
+	public $onQueryMessage(messages: [string, azdata.QueryExecuteMessageParams[]][]): void {
-		this._queryManagementService.onMessage(message);
+		this._queryManagementService.onMessage(serializableToMap(messages));
 	}
 	public $onEditSessionReady(handle: number, ownerUri: string, success: boolean, message: string): void {
 		this._queryManagementService.onEditSessionReady(ownerUri, success, message);

src/sql/workbench/api/common/extHostDataProtocol.ts

@@ -13,6 +13,8 @@ import { DataProviderType } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { IURITransformer } from 'vs/base/common/uriIpc';
 import { URI } from 'vs/base/common/uri';
 import { find } from 'vs/base/common/arrays';
+import { RunOnceScheduler } from 'vs/base/common/async';
+import { mapToSerializable } from 'sql/base/common/map';
 
 export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 
@@ -26,6 +28,9 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	private _adapter = new Map<number, azdata.DataProvider>();
 	private _providersByType = new Map<azdata.DataProviderType, azdata.DataProvider[]>();
 
+	private readonly messageRunner = new RunOnceScheduler(() => this.sendMessages(), 1000);
+	private readonly queuedMessages = new Map<string, azdata.QueryExecuteMessageParams[]>();
+
 	constructor(
 		mainContext: IMainContext,
 		private uriTransformer: IURITransformer | null
@@ -302,6 +307,11 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		if (this.uriTransformer) {
 			result.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(result.ownerUri))).toString(true);
 		}
+		// clear messages to maintain the order of things
+		if (this.messageRunner.isScheduled()) {
+			this.messageRunner.cancel();
+			this.sendMessages();
+		}
 		this._proxy.$onQueryComplete(handle, result);
 	}
 	$onBatchStart(handle: number, batchInfo: azdata.QueryExecuteBatchNotificationParams): void {
@@ -328,11 +338,23 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		}
 		this._proxy.$onResultSetUpdated(handle, resultSetInfo);
 	}
-	$onQueryMessage(handle: number, message: azdata.QueryExecuteMessageParams): void {
+	$onQueryMessage(message: azdata.QueryExecuteMessageParams): void {
 		if (this.uriTransformer) {
 			message.ownerUri = URI.from(this.uriTransformer.transformOutgoing(URI.parse(message.ownerUri))).toString(true);
 		}
-		this._proxy.$onQueryMessage(handle, message);
+		if (!this.queuedMessages.has(message.ownerUri)) {
+			this.queuedMessages.set(message.ownerUri, []);
+		}
+		this.queuedMessages.get(message.ownerUri).push(message);
+		if (!this.messageRunner.isScheduled()) {
+			this.messageRunner.schedule();
+		}
+	}
+
+	private sendMessages() {
+		const messages = mapToSerializable(this.queuedMessages);
+		this.queuedMessages.clear();
+		this._proxy.$onQueryMessage(messages);
 	}
 
 	$saveResults(handle: number, requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult> {

src/sql/workbench/api/common/sqlExtHost.api.impl.ts

@@ -240,7 +240,7 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
 				});
 
 				provider.registerOnMessage((message: azdata.QueryExecuteMessageParams) => {
-					extHostDataProvider.$onQueryMessage(provider.handle, message);
+					extHostDataProvider.$onQueryMessage(message);
 				});
 
 				provider.registerOnEditSessionReady((ownerUri: string, success: boolean, message: string) => {

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -219,7 +219,7 @@ export abstract class ExtHostDataProtocolShape {
 	/**
 	 * Callback when a message generated during query execution is issued
 	 */
-	$onQueryMessage(handle: number, message: azdata.QueryExecuteMessageParams): void { throw ni(); }
+	$onQueryMessage(message: azdata.QueryExecuteMessageParams): void { throw ni(); }
 
 	/**
 	 * Requests saving of the results from a result set into a specific format (CSV, JSON, Excel)
@@ -559,7 +559,7 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 	$onBatchComplete(handle: number, batchInfo: azdata.QueryExecuteBatchNotificationParams): void;
 	$onResultSetAvailable(handle: number, resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
 	$onResultSetUpdated(handle: number, resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
-	$onQueryMessage(handle: number, message: azdata.QueryExecuteMessageParams): void;
+	$onQueryMessage(message: [string, azdata.QueryExecuteMessageParams[]][]): void;
 	$onObjectExplorerSessionCreated(handle: number, message: azdata.ObjectExplorerSession): void;
 	$onObjectExplorerSessionDisconnected(handle: number, message: azdata.ObjectExplorerSession): void;
 	$onObjectExplorerNodeExpanded(providerId: string, message: azdata.ObjectExplorerExpandInfo): void;

src/sql/workbench/contrib/query/browser/messagePanel.ts

@@ -20,7 +20,7 @@ import { IThemeService, ITheme } from 'vs/platform/theme/common/themeService';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { OpenMode, ClickBehavior, ICancelableEvent, IControllerOptions } from 'vs/base/parts/tree/browser/treeDefaults';
 import { WorkbenchTreeController } from 'vs/platform/list/browser/listService';
-import { isArray } from 'vs/base/common/types';
+import { isArray, isString } from 'vs/base/common/types';
 import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { ScrollbarVisibility } from 'vs/base/common/scrollable';
@@ -30,8 +30,13 @@ import { ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { resultsErrorColor } from 'sql/platform/theme/common/colors';
 import { MessagePanelState } from 'sql/workbench/common/editor/query/messagePanelState';
 
-export interface IResultMessageIntern extends IQueryMessage {
+export interface IResultMessageIntern {
 	id?: string;
+	batchId?: number;
+	isError: boolean;
+	time?: string | Date;
+	message: string;
+	selection?: ISelectionData;
 }
 
 export interface IMessagePanelMessage {
@@ -190,7 +195,7 @@ export class MessagePanel extends Disposable {
 		let expandableTree: IExpandableTree = <IExpandableTree>this.tree;
 		if (this.state && this.state.scrollPosition) {
 			const previousScroll = this.state.scrollPosition;
-			this.tree.refresh(this.model).then(() => {
+			this.tree.refresh(this.model, false).then(() => {
 				// Restore the previous scroll position when switching between tabs
 				expandableTree.setScrollPosition(previousScroll);
 			});
@@ -290,13 +295,13 @@ class MessageDataSource implements IDataSource {
 
 class MessageRenderer implements IRenderer {
 
-	getHeight(tree: ITree, element: IQueryMessage): number {
+	getHeight(tree: ITree, element: IResultMessageIntern): number {
 		const lineHeight = 22;
 		let lines = element.message.split('\n').length;
 		return lineHeight * lines;
 	}
 
-	getTemplateId(tree: ITree, element: IQueryMessage): string {
+	getTemplateId(tree: ITree, element: IResultMessageIntern): string {
 		if (element instanceof Model) {
 			return TemplateIds.MODEL;
 		} else if (element.selection) {
@@ -333,13 +338,16 @@ class MessageRenderer implements IRenderer {
 		}
 	}
 
-	renderElement(tree: ITree, element: IQueryMessage, templateId: string, templateData: IMessageTemplate | IBatchTemplate): void {
+	renderElement(tree: ITree, element: IResultMessageIntern, templateId: string, templateData: IMessageTemplate | IBatchTemplate): void {
 		if (templateId === TemplateIds.MESSAGE || templateId === TemplateIds.ERROR) {
 			let data: IMessageTemplate = templateData;
 			data.message.innerText = element.message;
 		} else if (templateId === TemplateIds.BATCH) {
 			let data = templateData as IBatchTemplate;
-			data.timeStamp.innerText = element.time;
+			if (isString(element.time)) {
+				element.time = new Date(element.time!);
+			}
+			data.timeStamp.innerText = (element.time as Date).toLocaleTimeString();
 			data.message.innerText = element.message;
 		}
 	}

src/sql/workbench/services/insights/browser/insightsDialogController.ts

@@ -142,9 +142,10 @@ export class InsightsDialogController {
 				this._errorMessageService.showDialog(Severity.Error, nls.localize("insightsError", "Insights error"), error);
 			});
 		});
-		queryRunner.onMessage(message => {
+		queryRunner.onMessage(messages => {
-			if (message.isError) {
+			const errorMessage = messages.find(m => m.isError);
-				this._errorMessageService.showDialog(Severity.Error, nls.localize("insightsError", "Insights error"), message.message);
+			if (errorMessage) {
+				this._errorMessageService.showDialog(Severity.Error, nls.localize("insightsError", "Insights error"), errorMessage.message);
 			}
 		});
 	}

src/sql/workbench/services/insights/test/browser/insightsDialogController.test.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { InsightsDialogController } from 'sql/workbench/services/insights/browser/insightsDialogController';
-import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
+import QueryRunner, { IQueryMessage } from 'sql/workbench/services/query/common/queryRunner';
 import { ConnectionManagementService } from 'sql/workbench/services/connection/browser/connectionManagementService';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
@@ -112,7 +112,7 @@ function getPrimedQueryRunner(data: string[][], columns: string[]): IPrimedQuery
 	const emitter = new Emitter<string>();
 	const querymock = Mock.ofType(QueryRunner, MockBehavior.Strict);
 	querymock.setup(x => x.onQueryEnd).returns(x => emitter.event);
-	querymock.setup(x => x.onMessage).returns(x => new Emitter<azdata.IResultMessage>().event);
+	querymock.setup(x => x.onMessage).returns(x => new Emitter<[IQueryMessage]>().event);
 	querymock.setup(x => x.batchSets).returns(x => {
 		return <Array<azdata.BatchSummary>>[
 			{

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -340,10 +340,12 @@ class SqlKernel extends Disposable implements nb.IKernel {
 				this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error);
 			});
 		}));
-		this._register(queryRunner.onMessage(message => {
+		this._register(queryRunner.onMessage(messages => {
 			// TODO handle showing a messages output (should be updated with all messages, only changing 1 output in total)
-			if (this._future && isUndefinedOrNull(message.selection)) {
+			for (const message of messages) {
-				this._future.handleMessage(message);
+				if (this._future && isUndefinedOrNull(message.selection)) {
+					this._future.handleMessage(message);
+				}
 			}
 		}));
 		this._register(queryRunner.onBatchEnd(batch => {

src/sql/workbench/services/query/common/queryManagement.ts

@@ -45,7 +45,7 @@ export interface IQueryManagementService {
 	onBatchComplete(batchInfo: azdata.QueryExecuteBatchNotificationParams): void;
 	onResultSetAvailable(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
 	onResultSetUpdated(resultSetInfo: azdata.QueryExecuteResultSetNotificationParams): void;
-	onMessage(message: azdata.QueryExecuteMessageParams): void;
+	onMessage(message: Map<string, azdata.QueryExecuteMessageParams[]>): void;
 
 	// Edit Data Callbacks
 	onEditSessionReady(ownerUri: string, success: boolean, message: string): void;
@@ -283,10 +283,12 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 
-	public onMessage(message: azdata.QueryExecuteMessageParams): void {
+	public onMessage(messagesMap: Map<string, azdata.QueryExecuteMessageParams[]>): void {
-		this._notify(message.ownerUri, (runner: QueryRunner) => {
+		for (const [uri, messages] of messagesMap) {
-			runner.handleMessage(message);
+			this._notify(uri, (runner: QueryRunner) => {
-		});
+				runner.handleMessage(messages);
+			});
+		}
 	}
 
 	// Edit Data Functions

src/sql/workbench/services/query/common/queryRunner.ts

@@ -34,7 +34,11 @@ export interface IEditSessionReadyEvent {
 	message: string;
 }
 
-export interface IQueryMessage extends azdata.IResultMessage {
+export interface IQueryMessage {
+	batchId?: number;
+	isError: boolean;
+	time?: string;
+	message: string;
 	selection?: azdata.ISelectionData;
 }
 
@@ -58,8 +62,8 @@ export default class QueryRunner extends Disposable {
 	private _planXml = new Deferred<string>();
 	public get planXml(): Promise<string> { return this._planXml.promise; }
 
-	private _onMessage = this._register(new Emitter<IQueryMessage>());
+	private _onMessage = this._register(new Emitter<IQueryMessage[]>());
-	public get onMessage(): Event<IQueryMessage> { return this._onMessage.event; } // this is the only way typemoq can moq this... needs investigation @todo anthonydresser 5/2/2019
+	public get onMessage(): Event<IQueryMessage[]> { return this._onMessage.event; } // this is the only way typemoq can moq this... needs investigation @todo anthonydresser 5/2/2019
 
 	private _onResultSet = this._register(new Emitter<azdata.ResultSetSummary>());
 	public readonly onResultSet = this._onResultSet.event;
@@ -225,13 +229,13 @@ export default class QueryRunner extends Disposable {
 			error = error.message;
 		}
 		let message = nls.localize('query.ExecutionFailedError', "Execution failed due to an unexpected error: {0}\t{1}", eol, error);
-		this.handleMessage(<azdata.QueryExecuteMessageParams>{
+		this.handleMessage([<azdata.QueryExecuteMessageParams>{
 			ownerUri: this.uri,
 			message: {
 				isError: true,
 				message: message
 			}
-		});
+		}]);
 		this.handleQueryComplete(<azdata.QueryExecuteCompleteNotificationResult>{ ownerUri: this.uri });
 	}
 
@@ -267,7 +271,7 @@ export default class QueryRunner extends Disposable {
 		this._messages.push(message);
 
 		this._onQueryEnd.fire(timeStamp);
-		this._onMessage.fire(message);
+		this._onMessage.fire([message]);
 	}
 
 	/**
@@ -293,12 +297,12 @@ export default class QueryRunner extends Disposable {
 		let message = {
 			// account for index by 1
 			message: nls.localize('query.message.startQuery', "Started executing query at Line {0}", batch.selection.startLine + 1),
-			time: new Date(batch.executionStart).toLocaleTimeString(),
+			time: batch.executionStart,
 			selection: batch.selection,
 			isError: false
 		};
 		this._messages.push(message);
-		this._onMessage.fire(message);
+		this._onMessage.fire([message]);
 		this._onBatchStart.fire(batch);
 	}
 
@@ -403,13 +407,12 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Handle a Mssage from the service layer
 	 */
-	public handleMessage(obj: azdata.QueryExecuteMessageParams): void {
+	public handleMessage(messagesObj: azdata.QueryExecuteMessageParams[]): void {
-		let message = obj.message;
+		const messages = messagesObj.map(m => m.message);
-		message.time = new Date(message.time!).toLocaleTimeString();
+		this._messages.push(...messages);
-		this._messages.push(message);
 
 		// Send the message to the results pane
-		this._onMessage.fire(message);
+		this._onMessage.fire(messages);
 	}
 
 	/**
@@ -580,7 +583,7 @@ export default class QueryRunner extends Disposable {
 			};
 			this._messages.push(message);
 			// Send the message to the results pane
-			this._onMessage.fire(message);
+			this._onMessage.fire([message]);
 		}
 	}