fix: prevent locking before setting new user data (#751)

Author: oscb

packages/core/src/__tests__/methods/alias.test.ts

@@ -1,6 +1,5 @@
-import { SegmentClient } from '../../analytics';
-import { getMockLogger } from '../__helpers__/mockLogger';
-import { MockSegmentStore } from '../__helpers__/mockSegmentStore';
+import { EventType, SegmentEvent } from '../../types';
+import { createTestClient } from '../__helpers__/setupSegmentClient';
 
 jest.mock('../../uuid');
 
@@ -9,41 +8,30 @@ jest
   .mockReturnValue('2010-01-01T00:00:00.000Z');
 
 describe('methods #alias', () => {
-  const store = new MockSegmentStore({
-    userInfo: {
-      anonymousId: 'anonymousId',
-      userId: 'current-user-id',
-    },
-  });
-
-  const clientArgs = {
-    config: {
-      writeKey: '123-456',
-    },
-    logger: getMockLogger(),
-    store: store,
+  const initialUserInfo = {
+    anonymousId: 'anonymousId',
+    userId: 'current-user-id',
   };
 
+  const { store, client, expectEvent } = createTestClient({
+    userInfo: initialUserInfo,
+  });
+
   beforeEach(() => {
     store.reset();
     jest.clearAllMocks();
   });
 
   it('adds the alias event correctly', async () => {
-    const client = new SegmentClient(clientArgs);
-
-    jest.spyOn(client, 'process');
-
     await client.alias('new-user-id');
 
     const expectedEvent = {
       previousId: 'current-user-id',
-      type: 'alias',
+      type: EventType.AliasEvent,
       userId: 'new-user-id',
     };
 
-    expect(client.process).toHaveBeenCalledTimes(1);
-    expect(client.process).toHaveBeenCalledWith(expectedEvent);
+    expectEvent(expectedEvent);
 
     expect(client.userInfo.get()).toEqual({
       anonymousId: 'anonymousId',
@@ -53,32 +41,45 @@ describe('methods #alias', () => {
   });
 
   it('uses anonymousId in event if no userId in store', async () => {
-    const client = new SegmentClient({
-      ...clientArgs,
-      store: new MockSegmentStore({
-        userInfo: {
-          anonymousId: 'anonymousId',
-          userId: undefined,
-        },
-      }),
+    await store.userInfo.set({
+      anonymousId: 'anonymousId',
+      userId: undefined,
     });
-    jest.spyOn(client, 'process');
 
     await client.alias('new-user-id');
 
     const expectedEvent = {
       previousId: 'anonymousId',
-      type: 'alias',
+      type: EventType.AliasEvent,
       userId: 'new-user-id',
     };
 
-    expect(client.process).toHaveBeenCalledTimes(1);
-    expect(client.process).toHaveBeenCalledWith(expectedEvent);
+    expectEvent(expectedEvent);
 
     expect(client.userInfo.get()).toEqual({
       anonymousId: 'anonymousId',
       userId: 'new-user-id',
       traits: undefined,
     });
   });
+
+  it('is concurrency safe', async () => {
+    // We trigger an alias and do not await it, we do a track immediately and await.
+    // The track call should have the correct values injected into it.
+    client.alias('new-user-id');
+    await client.track('something');
+
+    const expectedTrackEvent: Partial<SegmentEvent> = {
+      event: 'something',
+      userId: 'new-user-id',
+      type: EventType.TrackEvent,
+    };
+
+    expectEvent(expectedTrackEvent);
+
+    expect(client.userInfo.get()).toEqual({
+      ...initialUserInfo,
+      userId: 'new-user-id',
+    });
+  });
 });

packages/core/src/__tests__/methods/identify.test.ts

@@ -58,8 +58,6 @@ describe('methods #identify', () => {
   });
 
   it('does not update userId when userId is undefined', async () => {
-    jest.spyOn(client, 'process');
-
     await client.identify(undefined, { name: 'Mary' });
 
     const expectedEvent = {
@@ -110,4 +108,24 @@ describe('methods #identify', () => {
       userId: 'new-user-id',
     });
   });
+
+  it('is concurrency safe', async () => {
+    // We trigger an identify and do not await it, we do a track immediately and await.
+    // The track call should have the correct values injected into it.
+    client.identify('new-user-id');
+    await client.track('something');
+
+    const expectedTrackEvent: Partial<SegmentEvent> = {
+      event: 'something',
+      userId: 'new-user-id',
+      type: EventType.TrackEvent,
+    };
+
+    expectEvent(expectedTrackEvent);
+
+    expect(client.userInfo.get()).toEqual({
+      ...initialUserInfo,
+      userId: 'new-user-id',
+    });
+  });
 });

packages/core/src/analytics.ts

@@ -507,8 +507,9 @@ export class SegmentClient {
   }
 
   async alias(newUserId: string) {
-    const { anonymousId, userId: previousUserId } =
-      await this.store.userInfo.get(true);
+    // We don't use a concurrency safe version of get here as we don't want to lock the values yet,
+    // we will update the values correctly when InjectUserInfo processes the change
+    const { anonymousId, userId: previousUserId } = this.store.userInfo.get();
 
     const event = createAliasEvent({
       anonymousId,

packages/core/src/plugins/InjectUserInfo.ts

@@ -14,15 +14,12 @@ export class InjectUserInfo extends PlatformPlugin {
   type = PluginType.before;
 
   async execute(event: SegmentEvent): Promise<SegmentEvent> {
-    const userInfo = await this.analytics!.userInfo.get(true);
-    const injectedEvent: SegmentEvent = {
-      ...event,
-      anonymousId: userInfo.anonymousId,
-      userId: userInfo.userId,
-    };
-
+    // Order here is IMPORTANT!
+    // Identify and Alias userInfo set operations have to come as soon as possible
+    // Do not block the set by doing a safe get first as it might cause a race condition
+    // within events procesing in the timeline asyncronously
     if (event.type === EventType.IdentifyEvent) {
-      await this.analytics!.userInfo.set((state) => ({
+      const userInfo = await this.analytics!.userInfo.set((state) => ({
         ...state,
         userId: event.userId ?? state.userId,
         traits: {
@@ -32,27 +29,41 @@ export class InjectUserInfo extends PlatformPlugin {
       }));
 
       return {
-        ...injectedEvent,
+        ...event,
+        anonymousId: userInfo.anonymousId,
         userId: event.userId ?? userInfo.userId,
         traits: {
           ...userInfo.traits,
           ...event.traits,
         },
       } as IdentifyEventType;
     } else if (event.type === EventType.AliasEvent) {
-      const { anonymousId, userId: previousUserId } = userInfo;
+      let previousUserId: string;
 
-      await this.analytics!.userInfo.set((state) => ({
-        ...state,
-        userId: event.userId,
-      }));
+      const userInfo = await this.analytics!.userInfo.set((state) => {
+        previousUserId = state.userId ?? state.anonymousId;
+
+        return {
+          ...state,
+          userId: event.userId,
+        };
+      });
 
       return {
-        ...injectedEvent,
+        ...event,
+        anonymousId: userInfo.anonymousId,
         userId: event.userId,
-        previousId: previousUserId || anonymousId,
+        previousId: previousUserId!,
       } as AliasEventType;
     }
+
+    const userInfo = await this.analytics!.userInfo.get(true);
+    const injectedEvent: SegmentEvent = {
+      ...event,
+      anonymousId: userInfo.anonymousId,
+      userId: userInfo.userId,
+    };
+
     return injectedEvent;
   }
 }

packages/core/src/plugins/__tests__/InjectUserInfo.test.ts

@@ -0,0 +1,109 @@
+import { createIdentifyEvent, createTrackEvent } from '../../events';
+import { type SegmentEvent, EventType, UserTraits } from '../../types';
+import { createTestClient } from '../../__tests__/__helpers__/setupSegmentClient';
+import { InjectUserInfo } from '../InjectUserInfo';
+
+describe('InjectContext', () => {
+  const currentUserId = 'current-user-id';
+  const anonymousId = 'very-anonymous';
+
+  const initialUserInfo = {
+    traits: {
+      name: 'Stacy',
+      age: 30,
+    },
+    userId: currentUserId,
+    anonymousId: anonymousId,
+  };
+  const { store, client } = createTestClient({
+    userInfo: initialUserInfo,
+  });
+
+  beforeEach(() => {
+    store.reset();
+    jest.clearAllMocks();
+  });
+
+  it('adds userId and anonymousId to events', async () => {
+    const expectedEvent: Partial<SegmentEvent> = {
+      event: 'something',
+      userId: currentUserId,
+      anonymousId: anonymousId,
+      type: EventType.TrackEvent,
+    };
+
+    const plugin = new InjectUserInfo();
+    plugin.configure(client);
+    const event = await plugin.execute(
+      createTrackEvent({
+        event: 'something',
+      })
+    );
+
+    expect(event).toMatchObject(expectedEvent);
+  });
+
+  it('updates userId and traits on identify', async () => {
+    const newUserId = 'new-user-id';
+    const newTraits: UserTraits = {
+      age: 30,
+    };
+
+    const expectedEvent: Partial<SegmentEvent> = {
+      userId: newUserId,
+      anonymousId: anonymousId,
+      traits: newTraits,
+      type: EventType.IdentifyEvent,
+    };
+
+    const plugin = new InjectUserInfo();
+    plugin.configure(client);
+    const event = await plugin.execute(
+      createIdentifyEvent({
+        userId: newUserId,
+        userTraits: newTraits,
+      })
+    );
+
+    expect(event).toMatchObject(expectedEvent);
+    expect(client.userInfo.get()).toEqual({
+      ...initialUserInfo,
+      userId: 'new-user-id',
+    });
+  });
+
+  it('is concurrency safe', async () => {
+    const newUserId = 'new-user-id';
+    const newTraits: UserTraits = {
+      age: 30,
+    };
+
+    const expectedEvent: Partial<SegmentEvent> = {
+      userId: newUserId,
+      anonymousId: anonymousId,
+      type: EventType.TrackEvent,
+    };
+
+    const plugin = new InjectUserInfo();
+    plugin.configure(client);
+
+    plugin.execute(
+      createIdentifyEvent({
+        userId: newUserId,
+        userTraits: newTraits,
+      })
+    );
+
+    const event = await plugin.execute(
+      createTrackEvent({
+        event: 'something',
+      })
+    );
+
+    expect(event).toMatchObject(expectedEvent);
+    expect(client.userInfo.get()).toEqual({
+      ...initialUserInfo,
+      userId: 'new-user-id',
+    });
+  });
+});