Merge pull request #33323 from storybookjs/valentin/fix-telemetry-last-event

Telemetry: Fix race condition in telemetry cache causing malformed JSON

Author: valentinpalkovic

code/core/src/core-server/server-channel/preview-initialized-channel.ts

@@ -1,11 +1,10 @@
 import type { Channel } from 'storybook/internal/channels';
 import { PREVIEW_INITIALIZED } from 'storybook/internal/core-events';
 import { type InitPayload, telemetry } from 'storybook/internal/telemetry';
+import { type CacheEntry, getLastEvents } from 'storybook/internal/telemetry';
+import { getSessionId } from 'storybook/internal/telemetry';
 import type { CoreConfig, Options } from 'storybook/internal/types';
 
-import { type CacheEntry, getLastEvents } from '../../telemetry/event-cache';
-import { getSessionId } from '../../telemetry/session-id';
-
 export const makePayload = (
   userAgent: string,
   lastInit: CacheEntry | undefined,

code/core/src/telemetry/event-cache.test.ts

@@ -3,7 +3,9 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { cache } from 'storybook/internal/common';
 
-import { get, getLastEvents, getPrecedingUpgrade, set } from './event-cache';
+import type { CacheEntry } from './event-cache';
+import { getLastEvents, getPrecedingUpgrade, set } from './event-cache';
+import type { TelemetryEvent } from './types';
 
 vi.mock('storybook/internal/common', { spy: true });
 
@@ -12,54 +14,94 @@ expect.addSnapshotSerializer({
   test: (val) => typeof val !== 'string',
 });
 
+// Helper to create valid TelemetryEvent objects
+const createTelemetryEvent = (
+  eventType: TelemetryEvent['eventType'],
+  eventId: string,
+  overrides?: Partial<TelemetryEvent>
+): TelemetryEvent => ({
+  eventType,
+  eventId,
+  sessionId: 'test-session',
+  context: {},
+  payload: {},
+  ...overrides,
+});
+
 describe('event-cache', () => {
-  const init = { body: { eventType: 'init', eventId: 'init' }, timestamp: 1 };
-  const upgrade = { body: { eventType: 'upgrade', eventId: 'upgrade' }, timestamp: 2 };
-  const dev = { body: { eventType: 'dev', eventId: 'dev' }, timestamp: 3 };
-  const build = { body: { eventType: 'build', eventId: 'build' }, timestamp: 3 };
-  const error = { body: { eventType: 'build', eventId: 'error' }, timestamp: 4 };
-  const versionUpdate = {
-    body: { eventType: 'version-update', eventId: 'version-update' },
+  const init: CacheEntry = {
+    body: createTelemetryEvent('init', 'init'),
+    timestamp: 1,
+  };
+  const upgrade: CacheEntry = {
+    body: createTelemetryEvent('upgrade', 'upgrade'),
+    timestamp: 2,
+  };
+  const dev: CacheEntry = {
+    body: createTelemetryEvent('dev', 'dev'),
+    timestamp: 3,
+  };
+  const build: CacheEntry = {
+    body: createTelemetryEvent('build', 'build'),
+    timestamp: 3,
+  };
+  const error: CacheEntry = {
+    body: createTelemetryEvent('build', 'error'),
+    timestamp: 4,
+  };
+  const versionUpdate: CacheEntry = {
+    body: createTelemetryEvent('version-update', 'version-update'),
     timestamp: 5,
   };
 
   describe('data handling', () => {
     it('errors', async () => {
       const preceding = await getPrecedingUpgrade({
-        init: { timestamp: 1, body: { ...init.body, error: {} } },
+        init: {
+          timestamp: 1,
+          body: { ...init.body, error: {} } as TelemetryEvent & { error: unknown },
+        },
       });
       expect(preceding).toMatchInlineSnapshot(`
         {
           "timestamp": 1,
           "eventType": "init",
-          "eventId": "init"
+          "eventId": "init",
+          "sessionId": "test-session"
         }
       `);
     });
 
     it('session IDs', async () => {
       const preceding = await getPrecedingUpgrade({
-        init: { timestamp: 1, body: { ...init.body, sessionId: 100 } },
+        init: {
+          timestamp: 1,
+          body: { ...init.body, sessionId: '100' },
+        },
       });
       expect(preceding).toMatchInlineSnapshot(`
         {
           "timestamp": 1,
           "eventType": "init",
           "eventId": "init",
-          "sessionId": 100
+          "sessionId": "100"
         }
       `);
     });
 
     it('extra fields', async () => {
       const preceding = await getPrecedingUpgrade({
-        init: { timestamp: 1, body: { ...init.body, foobar: 'baz' } },
+        init: {
+          timestamp: 1,
+          body: { ...init.body, foobar: 'baz' } as TelemetryEvent & { foobar: string },
+        },
       });
       expect(preceding).toMatchInlineSnapshot(`
         {
           "timestamp": 1,
           "eventType": "init",
-          "eventId": "init"
+          "eventId": "init",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -77,7 +119,8 @@ describe('event-cache', () => {
         {
           "timestamp": 1,
           "eventType": "init",
-          "eventId": "init"
+          "eventId": "init",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -88,7 +131,8 @@ describe('event-cache', () => {
         {
           "timestamp": 2,
           "eventType": "upgrade",
-          "eventId": "upgrade"
+          "eventId": "upgrade",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -99,7 +143,8 @@ describe('event-cache', () => {
         {
           "timestamp": 2,
           "eventType": "upgrade",
-          "eventId": "upgrade"
+          "eventId": "upgrade",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -127,31 +172,33 @@ describe('event-cache', () => {
     });
 
     it('both init and upgrade with intervening dev', async () => {
-      const secondUpgrade = {
-        body: { eventType: 'upgrade', eventId: 'secondUpgrade' },
+      const secondUpgrade: CacheEntry = {
+        body: createTelemetryEvent('upgrade', 'secondUpgrade'),
         timestamp: 4,
       };
       const preceding = await getPrecedingUpgrade({ init, dev, upgrade: secondUpgrade });
       expect(preceding).toMatchInlineSnapshot(`
         {
           "timestamp": 4,
           "eventType": "upgrade",
-          "eventId": "secondUpgrade"
+          "eventId": "secondUpgrade",
+          "sessionId": "test-session"
         }
       `);
     });
 
     it('both init and upgrade with non-intervening dev', async () => {
-      const earlyDev = {
-        body: { eventType: 'dev', eventId: 'earlyDev' },
+      const earlyDev: CacheEntry = {
+        body: createTelemetryEvent('dev', 'earlyDev'),
         timestamp: -1,
       };
       const preceding = await getPrecedingUpgrade({ dev: earlyDev, init, upgrade });
       expect(preceding).toMatchInlineSnapshot(`
         {
           "timestamp": 2,
           "eventType": "upgrade",
-          "eventId": "upgrade"
+          "eventId": "upgrade",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -174,7 +221,8 @@ describe('event-cache', () => {
         {
           "timestamp": 2,
           "eventType": "upgrade",
-          "eventId": "upgrade"
+          "eventId": "upgrade",
+          "sessionId": "test-session"
         }
       `);
     });
@@ -192,46 +240,28 @@ describe('event-cache', () => {
 
     it('getLastEvents waits for pending set operations to complete', async () => {
       const initialData = {
-        init: { timestamp: 1, body: { eventType: 'init', eventId: 'init-1' } },
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
       };
       const updatedData = {
-        init: { timestamp: 1, body: { eventType: 'init', eventId: 'init-1' } },
-        upgrade: { timestamp: 2, body: { eventType: 'upgrade', eventId: 'upgrade-1' } },
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+        upgrade: { timestamp: 2, body: createTelemetryEvent('upgrade', 'upgrade-1') },
       };
 
-      // Use a simple delay to simulate async operations
-      let setGetResolved = false;
-      let setSetResolved = false;
+      // Mock cache.get to return initial data first, then updated data
+      cacheGetMock
+        .mockResolvedValueOnce(initialData) // First call in setHelper
+        .mockResolvedValueOnce(updatedData); // Second call in getLastEvents
 
-      cacheGetMock.mockImplementationOnce(async () => {
-        while (!setGetResolved) {
-          await new Promise((resolve) => setTimeout(resolve, 10));
-        }
-        return initialData;
-      });
-
-      cacheSetMock.mockImplementationOnce(async () => {
-        while (!setSetResolved) {
-          await new Promise((resolve) => setTimeout(resolve, 10));
-        }
-      });
+      // Mock cache.set to resolve immediately
+      cacheSetMock.mockResolvedValue(undefined);
 
-      // Mock cache.get to return updated data after set completes
-      cacheGetMock.mockResolvedValueOnce(updatedData);
-
-      // Start a set operation (this will be pending)
-      const setPromiseResult = set('upgrade', { eventType: 'upgrade', eventId: 'upgrade-1' });
+      // Start a set operation (this will be queued and processed)
+      const setPromiseResult = set('upgrade', createTelemetryEvent('upgrade', 'upgrade-1'));
 
       // Immediately call getLastEvents() - it should wait for set() to complete
       const getPromise = getLastEvents();
 
-      // Verify that getLastEvents hasn't resolved yet (it's waiting)
-      await new Promise((resolve) => setTimeout(resolve, 50));
-
-      // Resolve the set operations
-      setGetResolved = true;
-      await new Promise((resolve) => setTimeout(resolve, 50));
-      setSetResolved = true;
+      // Wait for set operation to complete
       await setPromiseResult;
 
       // Now getLastEvents should complete and return the updated data
@@ -242,5 +272,79 @@ describe('event-cache', () => {
       expect(cacheGetMock).toHaveBeenCalledTimes(2); // Once in setHelper, once in getLastEvents
       expect(cacheSetMock).toHaveBeenCalledTimes(1);
     });
+
+    it('queues multiple set operations sequentially', async () => {
+      const initialData = {};
+      const afterFirst = {
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+      };
+      const afterSecond = {
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+        upgrade: { timestamp: 2, body: createTelemetryEvent('upgrade', 'upgrade-1') },
+      };
+      const afterThird = {
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+        upgrade: { timestamp: 2, body: createTelemetryEvent('upgrade', 'upgrade-1') },
+        dev: { timestamp: 3, body: createTelemetryEvent('dev', 'dev-1') },
+      };
+
+      // Mock cache.get to return data in sequence
+      cacheGetMock
+        .mockResolvedValueOnce(initialData) // First set: get initial
+        .mockResolvedValueOnce(afterFirst) // Second set: get after first
+        .mockResolvedValueOnce(afterSecond) // Third set: get after second
+        .mockResolvedValueOnce(afterThird); // getLastEvents: get after third
+
+      // Mock cache.set to resolve immediately
+      cacheSetMock.mockResolvedValue(undefined);
+
+      // Queue multiple set operations
+      const set1 = set('init', createTelemetryEvent('init', 'init-1'));
+      const set2 = set('upgrade', createTelemetryEvent('upgrade', 'upgrade-1'));
+      const set3 = set('dev', createTelemetryEvent('dev', 'dev-1'));
+
+      // Wait for all operations to complete
+      await Promise.all([set1, set2, set3]);
+
+      // Now getLastEvents should return the final state
+      const result = await getLastEvents();
+
+      // Verify all operations were processed sequentially
+      expect(result).toEqual(afterThird);
+      expect(cacheGetMock).toHaveBeenCalledTimes(4); // 3 sets + 1 getLastEvents
+      expect(cacheSetMock).toHaveBeenCalledTimes(3); // One for each set
+    });
+
+    it('handles errors in queued operations', async () => {
+      const initialData = {
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+      };
+      const afterDev = {
+        init: { timestamp: 1, body: createTelemetryEvent('init', 'init-1') },
+        dev: { timestamp: 3, body: createTelemetryEvent('dev', 'dev-1') },
+      };
+
+      // First operation will fail
+      cacheGetMock.mockResolvedValueOnce(initialData);
+      cacheSetMock.mockRejectedValueOnce(new Error('Cache write failed'));
+
+      // Queue an operation that will fail
+      const failedOperation = set('upgrade', createTelemetryEvent('upgrade', 'upgrade-1'));
+      await expect(failedOperation).rejects.toThrow('Cache write failed');
+
+      // Wait a bit to ensure queue processing completes
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // Verify subsequent operations can still be queued and succeed
+      cacheGetMock.mockResolvedValueOnce(initialData);
+      cacheSetMock.mockResolvedValueOnce(undefined);
+      cacheGetMock.mockResolvedValueOnce(afterDev);
+
+      await expect(set('dev', createTelemetryEvent('dev', 'dev-1'))).resolves.toBeUndefined();
+
+      // Verify the successful operation was processed
+      const result = await getLastEvents();
+      expect(result).toEqual(afterDev);
+    });
   });
 });