Add perf marks for query execution (#22082)

src/sql/workbench/api/browser/extensionHost.contribution.ts

@@ -19,6 +19,7 @@ import './mainThreadNotebook';
 import './mainThreadNotebookDocumentsAndEditors';
 import './mainThreadObjectExplorer';
 import './mainThreadQueryEditor';
+import './mainThreadPerf';
 import './mainThreadResourceProvider';
 import './mainThreadErrorDiagnostics';
 import './mainThreadTasks';

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

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { IDisposable, Disposable } from 'vs/base/common/lifecycle';
+import * as perf from 'vs/base/common/performance';
 import {
 	ExtHostDataProtocolShape,
 	MainThreadDataProtocolShape
@@ -116,15 +117,19 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData
 				return Promise.resolve(self._proxy.$cancelQuery(handle, ownerUri));
 			},
 			runQuery(ownerUri: string, selection: azdata.ISelectionData, runOptions?: azdata.ExecutionPlanOptions): Promise<void> {
+				perf.mark(`sql/query/${ownerUri}/main_$runQuery`);
 				return Promise.resolve(self._proxy.$runQuery(handle, ownerUri, selection, runOptions));
 			},
 			runQueryStatement(ownerUri: string, line: number, column: number): Promise<void> {
+				perf.mark(`sql/query/${ownerUri}/main_$runQueryStatement`);
 				return Promise.resolve(self._proxy.$runQueryStatement(handle, ownerUri, line, column));
 			},
 			runQueryString(ownerUri: string, queryString: string): Promise<void> {
+				perf.mark(`sql/query/${ownerUri}/main_$runQueryString`);
 				return Promise.resolve(self._proxy.$runQueryString(handle, ownerUri, queryString));
 			},
 			runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult> {
+				perf.mark(`sql/query/${ownerUri}/main_$runQueryAndReturn`);
 				return Promise.resolve(self._proxy.$runQueryAndReturn(handle, ownerUri, queryString));
 			},
 			parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
@@ -578,6 +583,7 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData
 
 	// Query Management handlers
 	public $onQueryComplete(handle: number, result: azdata.QueryExecuteCompleteNotificationResult): void {
+		perf.mark(`sql/query/${result.ownerUri}/main_$onQueryComplete`);
 		this._queryManagementService.onQueryComplete(result);
 	}
 	public $onBatchStart(handle: number, batchInfo: azdata.QueryExecuteBatchNotificationParams): void {

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

@@ -0,0 +1,24 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { MainThreadPerfShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
+import { SqlMainContext } from 'vs/workbench/api/common/extHost.protocol';
+import * as perf from 'vs/base/common/performance';
+
+@extHostNamedCustomer(SqlMainContext.MainThreadPerf)
+export class MainThreadPerf implements MainThreadPerfShape {
+	constructor(
+		context: IExtHostContext,
+	) { }
+
+	public $mark(name: string) {
+		perf.mark(name);
+	}
+
+	public dispose(): void {
+
+	}
+}

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

@@ -8,7 +8,7 @@ import * as azdata from 'azdata';
 import { Event, Emitter } from 'vs/base/common/event';
 import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { Disposable } from 'vs/workbench/api/common/extHostTypes';
-import { MainThreadDataProtocolShape, ExtHostDataProtocolShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
+import { MainThreadDataProtocolShape, ExtHostDataProtocolShape, MainThreadPerfShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
 import { DataProviderType } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { IURITransformer } from 'vs/base/common/uriIpc';
 import { URI, UriComponents } from 'vs/base/common/uri';
@@ -24,6 +24,7 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	readonly onDidChangeLanguageFlavor: Event<azdata.DidChangeLanguageFlavorParams> = this._onDidChangeLanguageFlavor.event;
 
 	private _proxy: MainThreadDataProtocolShape;
+	private _perfProxy: MainThreadPerfShape;
 
 	private static _handlePool: number = 0;
 	private _adapter = new Map<number, azdata.DataProvider>();
@@ -38,6 +39,7 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	) {
 		super();
 		this._proxy = mainContext.getProxy(SqlMainContext.MainThreadDataProtocol);
+		this._perfProxy = mainContext.getProxy(SqlMainContext.MainThreadPerf);
 	}
 
 	private _getTransformedUri(uri: string, transformMethod: (uri: UriComponents) => UriComponents): string {
@@ -291,19 +293,22 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		if (this.uriTransformer) {
 			ownerUri = this._getTransformedUri(ownerUri, this.uriTransformer.transformIncoming);
 		}
-
+		this._perfProxy.$mark(`sql/query/${ownerUri}/ext_$runQuery`);
 		return this._resolveProvider<azdata.QueryProvider>(handle).runQuery(ownerUri, selection, runOptions);
 	}
 
 	override $runQueryStatement(handle: number, ownerUri: string, line: number, column: number): Thenable<void> {
+		this._perfProxy.$mark(`sql/query/${ownerUri}/ext_$runQueryStatement`);
 		return this._resolveProvider<azdata.QueryProvider>(handle).runQueryStatement(ownerUri, line, column);
 	}
 
 	override $runQueryString(handle: number, ownerUri: string, queryString: string): Thenable<void> {
+		this._perfProxy.$mark(`sql/query/${ownerUri}/ext_$runQueryString`);
 		return this._resolveProvider<azdata.QueryProvider>(handle).runQueryString(ownerUri, queryString);
 	}
 
 	override $runQueryAndReturn(handle: number, ownerUri: string, queryString: string): Thenable<azdata.SimpleExecuteResult> {
+		this._perfProxy.$mark(`sql/query/${ownerUri}/ext_$runQueryAndReturn`);
 		return this._resolveProvider<azdata.QueryProvider>(handle).runQueryAndReturn(ownerUri, queryString);
 	}
 
@@ -349,6 +354,7 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		if (this.uriTransformer) {
 			result.ownerUri = this._getTransformedUri(result.ownerUri, this.uriTransformer.transformOutgoing);
 		}
+		this._perfProxy.$mark(`sql/query/${result.ownerUri}/ext_$onQueryComplete`);
 		// clear messages to maintain the order of things
 		if (this.messageRunner.isScheduled()) {
 			this.messageRunner.cancel();

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

@@ -0,0 +1,19 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IMainContext, SqlMainContext } from 'vs/workbench/api/common/extHost.protocol';
+
+import { MainThreadPerfShape, ExtHostPerfShape } from './sqlExtHost.protocol';
+
+export class ExtHostPerf implements ExtHostPerfShape {
+	private _proxy: MainThreadPerfShape;
+	constructor(mainContext: IMainContext) {
+		this._proxy = mainContext.getProxy(SqlMainContext.MainThreadPerf);
+	}
+
+	$mark(name: string): void {
+		this._proxy.$mark(name);
+	}
+}

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

@@ -1030,3 +1030,11 @@ export interface MainThreadExtensionManagementShape extends IDisposable {
 	$install(vsixPath: string): Thenable<string>;
 	$showObsoleteExtensionApiUsageNotification(message: string): void;
 }
+
+export interface ExtHostPerfShape {
+	$mark(name: string): void;
+}
+
+export interface MainThreadPerfShape {
+	$mark(name: string): void;
+}