From 31a7001dc11e9ab9302d8500e8e96779801aad32 Mon Sep 17 00:00:00 2001
From: Konstantin <kshabelko@microsoft.com>
Date: Fri, 26 Apr 2024 14:52:34 -0400
Subject: [PATCH] Make performanceClient.discardMeasurements() flush aux cache
 data in addition to measurements (#7061)

- Make `performanceClient.discardMeasurements()` flush aux cache data in
addition to measurements.
---
 ...-e0ac26de-f149-4faf-adef-154f60c566b8.json |  7 ++
 .../performance/PerformanceClient.ts          | 13 +---
 .../test/telemetry/PerformanceClient.spec.ts  | 76 ++++++++++++++++++-
 3 files changed, 84 insertions(+), 12 deletions(-)
 create mode 100644 change/@azure-msal-common-e0ac26de-f149-4faf-adef-154f60c566b8.json

diff --git a/change/@azure-msal-common-e0ac26de-f149-4faf-adef-154f60c566b8.json b/change/@azure-msal-common-e0ac26de-f149-4faf-adef-154f60c566b8.json
new file mode 100644
index 0000000000..c6e509982a
--- /dev/null
+++ b/change/@azure-msal-common-e0ac26de-f149-4faf-adef-154f60c566b8.json
@@ -0,0 +1,7 @@
+{
+  "type": "minor",
+  "comment": "Make performanceClient.discardMeasurements() flush aux cache data in addition to measurements #7061",
+  "packageName": "@azure/msal-common",
+  "email": "kshabelko@microsoft.com",
+  "dependentChangeType": "patch"
+}
diff --git a/lib/msal-common/src/telemetry/performance/PerformanceClient.ts b/lib/msal-common/src/telemetry/performance/PerformanceClient.ts
index 532cbb86fa..02edc27b5c 100644
--- a/lib/msal-common/src/telemetry/performance/PerformanceClient.ts
+++ b/lib/msal-common/src/telemetry/performance/PerformanceClient.ts
@@ -620,7 +620,7 @@ export abstract class PerformanceClient implements IPerformanceClient {
 
         if (isRoot) {
             queueInfo = this.getQueueInfo(event.correlationId);
-            this.discardCache(rootEvent.correlationId);
+            this.discardMeasurements(rootEvent.correlationId);
         } else {
             rootEvent.incompleteSubMeasurements?.delete(event.eventId);
         }
@@ -781,7 +781,7 @@ export abstract class PerformanceClient implements IPerformanceClient {
     }
 
     /**
-     * Removes measurements for a given correlation id.
+     * Removes measurements and aux data for a given correlation id.
      *
      * @param {string} correlationId
      */
@@ -791,15 +791,6 @@ export abstract class PerformanceClient implements IPerformanceClient {
             correlationId
         );
         this.eventsByCorrelationId.delete(correlationId);
-    }
-
-    /**
-     * Removes cache for a given correlation id.
-     *
-     * @param {string} correlationId correlation identifier
-     */
-    private discardCache(correlationId: string): void {
-        this.discardMeasurements(correlationId);
 
         this.logger.trace(
             "PerformanceClient: QueueMeasurements discarded",
diff --git a/lib/msal-common/test/telemetry/PerformanceClient.spec.ts b/lib/msal-common/test/telemetry/PerformanceClient.spec.ts
index 1d9444aaf6..e640b2622c 100644
--- a/lib/msal-common/test/telemetry/PerformanceClient.spec.ts
+++ b/lib/msal-common/test/telemetry/PerformanceClient.spec.ts
@@ -72,7 +72,10 @@ export class MockPerformanceClient
         eventName: PerformanceEvents,
         correlationId?: string | undefined
     ): void {
-        return;
+        this.preQueueTimeByCorrelationId.set(correlationId || "", {
+            name: eventName,
+            time: 12345,
+        });
     }
 
     getDurationMs(startTimeMs: number): number {
@@ -1122,4 +1125,75 @@ describe("PerformanceClient.spec.ts", () => {
             rootEvent.end({ success: true });
         });
     });
+
+    describe("discard", () => {
+        it("discards cache data", () => {
+            const mockPerfClient = new MockPerformanceClient();
+            const correlationId = "test-correlation-id";
+            const dummyCorrelationId = "dummy-correlation-id";
+
+            const rootEvent = mockPerfClient.startMeasurement(
+                PerformanceEvents.AcquireTokenSilent,
+                correlationId
+            );
+            const firstEvent = mockPerfClient.startMeasurement(
+                PerformanceEvents.AcquireTokenSilentAsync,
+                correlationId
+            );
+            mockPerfClient.setPreQueueTime(
+                PerformanceEvents.AcquireTokenSilentAsync,
+                correlationId
+            );
+            const secondEvent = mockPerfClient.startMeasurement(
+                PerformanceEvents.AcquireTokenFromCache,
+                correlationId
+            );
+            mockPerfClient.setPreQueueTime(
+                PerformanceEvents.AcquireTokenFromCache,
+                correlationId
+            );
+            secondEvent.end({ success: true });
+            firstEvent.end({ success: true });
+            rootEvent.discard();
+
+            mockPerfClient.startMeasurement(
+                PerformanceEvents.AcquireTokenSilent,
+                dummyCorrelationId
+            );
+            mockPerfClient.setPreQueueTime(
+                PerformanceEvents.AcquireTokenSilent,
+                dummyCorrelationId
+            );
+
+            expect(
+                // @ts-ignore
+                mockPerfClient.eventsByCorrelationId.has(correlationId)
+            ).toBeFalsy();
+            expect(
+                // @ts-ignore
+                mockPerfClient.preQueueTimeByCorrelationId.has(correlationId)
+            ).toBeFalsy();
+            expect(
+                // @ts-ignore
+                mockPerfClient.queueMeasurements.has(correlationId)
+            ).toBeFalsy();
+            // @ts-ignore
+            expect(mockPerfClient.eventStack.has(correlationId)).toBeFalsy();
+
+            expect(
+                // @ts-ignore
+                mockPerfClient.eventsByCorrelationId.has(dummyCorrelationId)
+            ).toBeTruthy();
+            expect(
+                // @ts-ignore
+                mockPerfClient.preQueueTimeByCorrelationId.has(
+                    dummyCorrelationId
+                )
+            ).toBeTruthy();
+            expect(
+                // @ts-ignore
+                mockPerfClient.eventStack.has(dummyCorrelationId)
+            ).toBeTruthy();
+        });
+    });
 });