More disposal / flush logic

Author: lramos15

src/browser/telemetryReporter.ts

@@ -40,7 +40,18 @@ const webAppInsightsClientFactory = async (key: string, replacementOptions?: Rep
 			);
 		},
 		flush: async () => {
-			appInsightsClient?.flush();
+			appInsightsClient?.flush(false);
+		},
+		dispose: async () => {
+			appInsightsClient?.flush(true);
+			const unloadPromise = new Promise<void>((resolve) => {
+				appInsightsClient?.unload(true, () => {
+					resolve();
+					appInsightsClient = undefined;
+				});
+			}
+			);
+			return unloadPromise;
 		}
 	};
 	return telemetryClient;

src/common/1dsClientFactory.ts

@@ -69,7 +69,26 @@ const getAICore = async (key: string, vscodeAPI: typeof vscode, xhrOverride?: IX
  * @param xhrOverride An optional override to use for requests instead of the XHTMLRequest object. Useful for node environments
  */
 export const oneDataSystemClientFactory = async (key: string, vscodeAPI: typeof vscode, xhrOverride?: IXHROverride): Promise<BaseTelemetryClient> => {
-	const appInsightsCore = await getAICore(key, vscodeAPI, xhrOverride);
+	let appInsightsCore: AppInsightsCore | undefined = await getAICore(key, vscodeAPI, xhrOverride);
+	const flushOneDS = async () => {
+		try {
+			const flushPromise = new Promise<void>((resolve, reject) => {
+				if (!appInsightsCore) {
+					resolve();
+					return;
+				}
+				appInsightsCore.flush(true, (completedFlush) => {
+					if (!completedFlush) {
+						reject("Failed to flush app 1DS!");
+						return;
+					}
+				});
+			});
+			return flushPromise;
+		} catch (e: any) {
+			throw new Error("Failed to flush 1DS!\n" + e.message);
+		}
+	};
 	// Shape the app insights core from 1DS into a standard format
 	const telemetryClient: BaseTelemetryClient = {
 		logEvent: (eventName: string, data?: SenderData) => {
@@ -82,28 +101,21 @@ export const oneDataSystemClientFactory = async (key: string, vscodeAPI: typeof
 				throw new Error("Failed to log event to app insights!\n" + e.message);
 			}
 		},
-		flush: async () => {
-			try {
-				const flushPromise = new Promise<void>((resolve, reject) => {
-					if (!appInsightsCore) {
-						resolve();
-						return;
-					}
-					appInsightsCore.flush(true, (completedFlush) => {
-						if (!completedFlush) {
-							reject("Failed to flush app 1DS!");
-							return;
-						}
-						appInsightsCore.unload(true, () => {
-							resolve();
-							return;
-						});
-					});
+		flush: flushOneDS,
+		dispose: async () => {
+			await flushOneDS();
+			const disposePromise = new Promise<void>((resolve) => {
+				if (!appInsightsCore) {
+					resolve();
+					return;
+				}
+				appInsightsCore.unload(true, () => {
+					resolve();
+					appInsightsCore = undefined;
+					return;
 				});
-				return flushPromise;
-			} catch (e: any) {
-				throw new Error("Failed to flush 1DS!\n" + e.message);
-			}
+			});
+			return disposePromise;
 		}
 	};
 	return telemetryClient;

src/common/baseTelemetryReporter.ts

@@ -167,7 +167,8 @@ export class BaseTelemetryReporter {
 	/**
 	 * Disposes of the telemetry reporter
 	 */
-	public dispose(): Promise<any> {
+	public async dispose(): Promise<any> {
+		await this.telemetrySender.dispose();
 		this.telemetryLogger.dispose();
 		return Promise.all(this.disposables.map(d => d.dispose()));
 	}

src/common/baseTelemetrySender.ts

@@ -7,11 +7,13 @@ import { SenderData } from "./baseTelemetryReporter";
 
 export interface BaseTelemetryClient {
 	logEvent(eventName: string, data?: SenderData): void;
-	flush(): void | Promise<void>;
+	flush(): Promise<void>;
+	dispose(): Promise<void>;
 }
 
 export interface ILazyTelemetrySender extends TelemetrySender {
 	instantiateSender(): void
+	dispose(): Promise<void>;
 }
 
 enum InstantiationStatus {
@@ -83,8 +85,13 @@ export class BaseTelemetrySender implements ILazyTelemetrySender {
 	 * Flushes the buffered telemetry data
 	 */
 	async flush(): Promise<void> {
+		return this._telemetryClient?.flush();
+	}
+
+	async dispose(): Promise<void> {
+		await this.flush();
 		if (this._telemetryClient) {
-			await this._telemetryClient.flush();
+			await this._telemetryClient.dispose();
 			this._telemetryClient = undefined;
 		}
 		return;

src/node/telemetryReporter.ts

@@ -5,7 +5,7 @@
 import * as os from "os";
 import * as vscode from "vscode";
 import * as https from "https";
-import type { TelemetryClient } from "applicationinsights";
+import { type TelemetryClient } from "applicationinsights";
 import { SenderData, BaseTelemetryReporter, ReplacementOption } from "../common/baseTelemetryReporter";
 import { BaseTelemetrySender, BaseTelemetryClient } from "../common/baseTelemetrySender";
 import { TelemetryUtil } from "../common/util";
@@ -39,6 +39,7 @@ const appInsightsClientFactory = async (key: string, replacementOptions?: Replac
 				.setAutoDependencyCorrelation(false)
 				.setAutoCollectConsole(false)
 				.setAutoCollectHeartbeat(false)
+				.setAutoCollectIncomingRequestAzureFunctions(false)
 				.setUseDiskRetryCaching(true)
 				.start();
 			appInsightsClient = appInsights.defaultClient;
@@ -76,6 +77,10 @@ const appInsightsClientFactory = async (key: string, replacementOptions?: Replac
 			} catch (e: any) {
 				throw new Error("Failed to flush app insights!\n" + e.message);
 			}
+		},
+		dispose: async () => {
+			appInsightsClient?.flush();
+			appInsightsClient = undefined;
 		}
 	};
 	return telemetryClient;

test/baseTelemetrySender.test.ts

@@ -11,6 +11,7 @@ describe("Base telemetry sender test suite", () => {
 	const telemetryClient: BaseTelemetryClient = {
 		logEvent: sinon.spy(),
 		flush: sinon.spy(),
+		dispose: sinon.spy(),
 	};
 	const telemetryClientFactory: (key: string) => Promise<BaseTelemetryClient> = async () => {
 		return telemetryClient;