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;
+}

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

@@ -7,6 +7,7 @@ import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { IDisposable } from 'vs/base/common/lifecycle';
+import * as perf from 'vs/base/common/performance';
 import * as azdata from 'azdata';
 import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
 import { Event, Emitter } from 'vs/base/common/event';
@@ -234,6 +235,7 @@ export class QueryManagementService implements IQueryManagementService {
 
 	public runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void> {
 		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQuery, ownerUri, runOptions);
+		perf.mark(`sql/query/${ownerUri}/runQuery`);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQuery(ownerUri, rangeToSelectionData(range), runOptions);
 		});
@@ -241,6 +243,7 @@ export class QueryManagementService implements IQueryManagementService {
 
 	public runQueryStatement(ownerUri: string, line: number, column: number): Promise<void> {
 		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQueryStatement, ownerUri);
+		perf.mark(`sql/query/${ownerUri}/runQueryStatement`);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryStatement(ownerUri, line - 1, column - 1); // we are taking in a vscode IRange which is 1 indexed, but our api expected a 0 index
 		});
@@ -248,6 +251,7 @@ export class QueryManagementService implements IQueryManagementService {
 
 	public runQueryString(ownerUri: string, queryString: string): Promise<void> {
 		this.addTelemetry(TelemetryKeys.TelemetryAction.RunQueryString, ownerUri);
+		perf.mark(`sql/query/${ownerUri}/runQueryString`);
 		return this._runAction(ownerUri, (runner) => {
 			return runner.runQueryString(ownerUri, queryString);
 		});

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

@@ -19,6 +19,7 @@ import { Emitter, Event } from 'vs/base/common/event';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { URI } from 'vs/base/common/uri';
+import * as perf from 'vs/base/common/performance';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { IGridDataProvider, getResultsString, getTableHeaderString } from 'sql/workbench/services/query/common/gridDataProvider';
 import { getErrorMessage } from 'vs/base/common/errors';
@@ -243,6 +244,9 @@ export default class QueryRunner extends Disposable {
 		// this also isn't exact but its the best we can do
 		this._queryEndTime = new Date();
 
+		perf.mark(`sql/query/${this.uri}/handleQueryComplete`);
+		this.logPerfMarks();
+
 		// Store the batch sets we got back as a source of "truth"
 		this._isExecuting = false;
 		this._hasCompleted = true;
@@ -513,6 +517,25 @@ export default class QueryRunner extends Disposable {
 		};
 		this._onVisualize.fire(result);
 	}
+
+	private logPerfMarks(): void {
+		let marks = perf.getMarks().filter(mark => mark.name.startsWith(`sql/query/${this.uri}`));
+		// One query editor can have one set of marks for each query gets executed,
+		// we only log the latest set of marks for the latest query.
+		let reverseIdx = marks.slice().reverse().findIndex(m => m.name.startsWith(`sql/query/${this.uri}/runQuery`));
+		marks = marks.slice(-(reverseIdx + 1));
+		let head = `Name\tTimestamp\tDelta\tTotal\n`;
+		let traceMsg = '';
+		let lastStartTime = -1;
+		let total = 0;
+		for (const { name, startTime } of marks) {
+			let delta = lastStartTime !== -1 ? startTime - lastStartTime : 0;
+			total += delta;
+			traceMsg += `${name}\t${startTime}\t${delta}\t${total}\n`;
+			lastStartTime = startTime;
+		}
+		this.logService.debug(`Traces for sql/query/${this.uri}:\n${head}${traceMsg}`);
+	}
 }
 
 export class QueryGridDataProvider implements IGridDataProvider {

src/vs/workbench/api/common/extHost.protocol.ts

@@ -85,7 +85,9 @@ import {
 	ExtHostBackgroundTaskManagementShape, ExtHostDashboardWebviewsShape, ExtHostModelViewShape, ExtHostModelViewTreeViewsShape,
 	ExtHostDashboardShape, ExtHostModelViewDialogShape, ExtHostQueryEditorShape, ExtHostExtensionManagementShape, ExtHostAzureBlobShape,
 	ExtHostNotebookShape as SqlExtHostNotebookShape, ExtHostWorkspaceShape as SqlExtHostWorkspaceShape,
-	ExtHostNotebookDocumentsAndEditorsShape as SqlExtHostNotebookDocumentsAndEditorsShape
+	ExtHostNotebookDocumentsAndEditorsShape as SqlExtHostNotebookDocumentsAndEditorsShape,
+	ExtHostPerfShape,
+	MainThreadPerfShape
 } from 'sql/workbench/api/common/sqlExtHost.protocol';
 
 export interface IWorkspaceData extends IStaticWorkspaceData {
@@ -2406,6 +2408,7 @@ export const SqlMainContext = {
 	MainThreadExtensionManagement: createProxyIdentifier<MainThreadExtensionManagementShape>('MainThreadExtensionManagement'),
 	MainThreadWorkspace: createProxyIdentifier<SqlMainThreadWorkspaceShape>('MainThreadWorkspace'),
 	MainThreadAzureBlob: createProxyIdentifier<MainThreadAzureBlobShape>('MainThreadAzureBlob'),
+	MainThreadPerf: createProxyIdentifier<MainThreadPerfShape>('MainThreadPerf')
 };
 
 export const SqlExtHostContext = {
@@ -2430,5 +2433,6 @@ export const SqlExtHostContext = {
 	ExtHostNotebookDocumentsAndEditors: createProxyIdentifier<SqlExtHostNotebookDocumentsAndEditorsShape>('ExtHostNotebookDocumentsAndEditors'),
 	ExtHostExtensionManagement: createProxyIdentifier<ExtHostExtensionManagementShape>('ExtHostExtensionManagement'),
 	ExtHostWorkspace: createProxyIdentifier<SqlExtHostWorkspaceShape>('ExtHostWorkspace'),
-	ExtHostAzureBlob: createProxyIdentifier<ExtHostAzureBlobShape>('ExtHostAzureBlob')
+	ExtHostAzureBlob: createProxyIdentifier<ExtHostAzureBlobShape>('ExtHostAzureBlob'),
+	ExtHostPerf: createProxyIdentifier<ExtHostPerfShape>('ExtHostPerf')
 };