fix: Fix a race condition when handling multiple flag updates. (#365)

The setState function is effectively asynchronous, multiple setState
methods may be called in a single cycle before this.state is updated.

So if there were multiple dispatches of changes within a single batch,
then some of the updates could be missed.

The way to handle this in react is to use the previous state provided
using the method version of setState. The operations will be done
sequentially with each having the correct view of the previous state.

Author: kinyoklion

src/provider.test.tsx

@@ -456,10 +456,18 @@ describe('LDProvider', () => {
     const mockSetState = jest.spyOn(instance, 'setState');
 
     await instance.componentDidMount();
-    const setStateFunction = mockSetState.mock?.lastCall?.[0] as (p: ProviderState) => ProviderState;
+
+    // Each set of the state depends on the previous state, so to re-create the final state, we need to call the
+    // setState function for each call.
+    let finalState = previousState;
+
+    for(const call of mockSetState.mock.calls) {
+      const setStateFunction = call[0] as (p: ProviderState) => ProviderState;
+      finalState = setStateFunction(finalState);
+    }
 
     expect(mockLDClient.on).toHaveBeenCalledWith('change', expect.any(Function));
-    expect(setStateFunction(previousState)).toEqual({
+    expect(finalState).toMatchObject({
       flags: { anotherTestFlag: true, testFlag: false },
       unproxiedFlags: { 'another-test-flag': true, 'test-flag': false },
       flagKeyMap: { anotherTestFlag: 'another-test-flag', testFlag: 'test-flag' },
@@ -480,16 +488,56 @@ describe('LDProvider', () => {
     const mockSetState = jest.spyOn(instance, 'setState');
 
     await instance.componentDidMount();
-    const setStateFunction = mockSetState.mock?.lastCall?.[0] as (p: ProviderState) => ProviderState;
+
+    // Each set of the state depends on the previous state, so to re-create the final state, we need to call the
+    // setState function for each call.
+    let finalState = previousState;
+
+    for (const call of mockSetState.mock.calls) {
+      const setStateFunction = call[0] as (p: ProviderState) => ProviderState;
+      finalState = setStateFunction(finalState);
+    }
 
     expect(mockLDClient.on).toHaveBeenCalledWith('change', expect.any(Function));
-    expect(setStateFunction(previousState)).toEqual({
+    expect(finalState).toMatchObject({
       flagKeyMap: {},
       unproxiedFlags: { 'another-test-flag': false, 'test-flag': false },
       flags: { 'another-test-flag': false, 'test-flag': false },
     });
   });
 
+  test('handles deletion of flags', async () => {
+    mockLDClient.on.mockImplementation((_e: string, cb: (c: LDFlagChangeset) => void) => {
+      cb({ 'another-test-flag': { current: undefined, previous: true }, 'test-flag': { current: false, previous: true } });
+    });
+    const props: ProviderConfig = { clientSideID, reactOptions: { useCamelCaseFlagKeys: false } };
+    const LaunchDarklyApp = (
+      <LDProvider {...props}>
+        <App />
+      </LDProvider>
+    );
+    const instance = create(LaunchDarklyApp).root.findByType(LDProvider).instance as EnhancedComponent;
+    const mockSetState = jest.spyOn(instance, 'setState');
+
+    await instance.componentDidMount();
+
+    // Each set of the state depends on the previous state, so to re-create the final state, we need to call the
+    // setState function for each call.
+    let finalState = previousState;
+
+    for (const call of mockSetState.mock.calls) {
+      const setStateFunction = call[0] as (p: ProviderState) => ProviderState;
+      finalState = setStateFunction(finalState);
+    }
+
+    expect(mockLDClient.on).toHaveBeenCalledWith('change', expect.any(Function));
+    expect(finalState).toMatchObject({
+      flagKeyMap: {},
+      unproxiedFlags: { 'test-flag': false },
+      flags: { 'test-flag': false },
+    });
+  });
+
   test(`if props.deferInitialization is true, ld client will only initialize once props.user is defined`, async () => {
     options = { ...options, bootstrap: {} };
     const props: ProviderConfig = { clientSideID, deferInitialization: true, options };

src/provider.tsx

@@ -57,16 +57,19 @@ class LDProvider extends Component<PropsWithChildren<ProviderConfig>, ProviderSt
     ldClient.on('change', (changes: LDFlagChangeset) => {
       const reactOptions = this.getReactOptions();
       const updates = getFlattenedFlagsFromChangeset(changes, targetFlags);
-      const unproxiedFlags = {
-        ...this.state.unproxiedFlags,
-        ...updates,
-      };
+
       if (Object.keys(updates).length > 0) {
-        this.setState((prevState) => ({
-          ...prevState,
-          unproxiedFlags,
-          ...getFlagsProxy(ldClient, unproxiedFlags, reactOptions, targetFlags),
-        }));
+        this.setState((prevState: ProviderState) => {
+          const unproxiedFlags = {
+            ...prevState.unproxiedFlags,
+            ...updates,
+          };
+          return {
+            ...prevState,
+            unproxiedFlags,
+            ...getFlagsProxy(ldClient, unproxiedFlags, reactOptions, targetFlags),
+          }
+        });
       }
     });
   };

src/withLDProvider.test.tsx

@@ -130,10 +130,17 @@ describe('withLDProvider', () => {
     const mockSetState = jest.spyOn(instance, 'setState');
 
     await instance.componentDidMount();
-    const setStateFunction = mockSetState.mock?.lastCall?.[0] as (p: ProviderState) => ProviderState;
+    // Each set of the state depends on the previous state, so to re-create the final state, we need to call the
+    // setState function for each call.
+    let finalState = previousState;
+
+    for (const call of mockSetState.mock.calls) {
+      const setStateFunction = call[0] as (p: ProviderState) => ProviderState;
+      finalState = setStateFunction(finalState);
+    }
 
     expect(mockLDClient.on).toHaveBeenCalledWith('change', expect.any(Function));
-    expect(setStateFunction(previousState)).toEqual({
+    expect(finalState).toMatchObject({
       flags: { anotherTestFlag: true, testFlag: false },
       unproxiedFlags: { 'test-flag': false, 'another-test-flag': true },
       flagKeyMap: { testFlag: 'test-flag', anotherTestFlag: 'another-test-flag' },
@@ -149,10 +156,17 @@ describe('withLDProvider', () => {
     const mockSetState = jest.spyOn(instance, 'setState');
 
     await instance.componentDidMount();
-    const setStateFunction = mockSetState.mock?.lastCall?.[0] as (p: ProviderState) => ProviderState;
+    // Each set of the state depends on the previous state, so to re-create the final state, we need to call the
+    // setState function for each call.
+    let finalState = previousState;
+
+    for (const call of mockSetState.mock.calls) {
+      const setStateFunction = call[0] as (p: ProviderState) => ProviderState;
+      finalState = setStateFunction(finalState);
+    }
 
     expect(mockLDClient.on).toHaveBeenCalledWith('change', expect.any(Function));
-    expect(setStateFunction(previousState)).toEqual({
+    expect(finalState).toMatchObject({
       flags: { 'test-flag': false, 'another-test-flag': false },
       unproxiedFlags: { 'test-flag': false, 'another-test-flag': false },
       flagKeyMap: {},