From 2dd385764bd9c01b3f88f34b2da80b023b52cad1 Mon Sep 17 00:00:00 2001
From: Bill Lynch <llynch.bill@gmail.com>
Date: Fri, 13 Oct 2023 13:57:24 +0100
Subject: [PATCH] fix(react): issue with multiple context updates

---
 packages/integration-react/package.json       |   6 +-
 .../src/ReactAdapter.test.tsx                 |  98 ++++----
 .../integration-react/src/ReactAdaptor.tsx    | 238 ++++++------------
 3 files changed, 133 insertions(+), 209 deletions(-)

diff --git a/packages/integration-react/package.json b/packages/integration-react/package.json
index 6d567c8c..95ed7720 100644
--- a/packages/integration-react/package.json
+++ b/packages/integration-react/package.json
@@ -19,9 +19,13 @@
     "@spotify-confidence/openfeature-web-provider": "^0.0.3",
     "@testing-library/jest-dom": "^5.16.5",
     "@testing-library/react": "^14.0.0",
-    "@types/node-fetch": "^2.6.4"
+    "@types/node-fetch": "^2.6.4",
+    "@types/use-sync-external-store": "^0.0.4"
   },
   "publishConfig": {
     "access": "public"
+  },
+  "dependencies": {
+    "use-sync-external-store": "^1.2.0"
   }
 }
diff --git a/packages/integration-react/src/ReactAdapter.test.tsx b/packages/integration-react/src/ReactAdapter.test.tsx
index 23290a83..eea6d335 100644
--- a/packages/integration-react/src/ReactAdapter.test.tsx
+++ b/packages/integration-react/src/ReactAdapter.test.tsx
@@ -3,8 +3,8 @@ import { render, screen, act } from '@testing-library/react';
 import { OpenFeatureClient, ProviderEvents } from '@openfeature/web-sdk';
 import {
   useStringValue,
-  ClientManager,
-  ClientManagerContext,
+  FeaturesStore,
+  FeatureStoreContext,
   useNumberValue,
   useObjectValue,
   useBooleanValue,
@@ -28,20 +28,27 @@ const fakeClient = {
 } as jest.MockedObject<OpenFeatureClient>;
 
 describe('hooks', () => {
-  let fireReady: Function | undefined;
-  let fireError: Function | undefined;
-  let fireStale: Function | undefined;
+  const handlerMap = new Map<ProviderEvents, Set<Function>>();
+  const fire = (event: ProviderEvents) => {
+    const handlers = handlerMap.get(event);
+    if (handlers) {
+      for (const handler of handlers) {
+        handler();
+      }
+    }
+  };
 
   beforeEach(() => {
     fakeClient.addHandler.mockImplementation((event, cb) => {
-      if (event === ProviderEvents.Error) {
-        fireError = cb;
-      }
-      if (event === ProviderEvents.Stale) {
-        fireStale = cb;
-      }
-      if (event === ProviderEvents.Ready) {
-        fireReady = cb;
+      let handlers = handlerMap.get(event);
+      if (!handlers) handlerMap.set(event, (handlers = new Set()));
+      handlers.add(cb);
+    });
+    fakeClient.removeHandler.mockImplementation((event, cb) => {
+      const handlers = handlerMap.get(event);
+      if (handlers) {
+        handlers.delete(cb);
+        if (handlers.size === 0) handlerMap.delete(event);
       }
     });
   });
@@ -66,31 +73,25 @@ describe('hooks', () => {
           return <p>{JSON.stringify(val)}</p>;
         };
 
-        const manager = new ClientManager(fakeClient);
-        manager.ref();
+        const featureStore = FeaturesStore.forClient(fakeClient);
 
         const { unmount } = render(
-          <ClientManagerContext.Provider value={() => manager}>
+          <FeatureStoreContext.Provider value={featureStore}>
             <React.Suspense fallback={<p>suspense</p>}>
               <TestComponent />
             </React.Suspense>
-          </ClientManagerContext.Provider>,
+          </FeatureStoreContext.Provider>,
         );
 
         expect(screen.getByText('suspense')).toBeInTheDocument();
 
-        await act(async () => {
-          fireReady!();
-          await manager.promise;
-        });
+        fire(ProviderEvents.Ready);
 
-        expect(screen.getByText(JSON.stringify(mockValue))).toBeInTheDocument();
+        expect(await screen.findByText(JSON.stringify(mockValue))).toBeInTheDocument();
 
         unmount();
 
-        manager.unref();
-        expect(fakeClient.addHandler).toBeCalledTimes(3);
-        expect(fakeClient.removeHandler).toBeCalledTimes(3);
+        expect(handlerMap.size).toBe(0);
       },
     );
   });
@@ -109,44 +110,39 @@ describe('hooks', () => {
     `(
       `($name): should show suspense until the value is resolved and reload when the context is changed, show new value and cleanup`,
       async ({ hookUnderTest, defaultValue, mockValue1, mockValue2, mockFunc }) => {
-        mockFunc.mockReturnValueOnce(mockValue1).mockReturnValueOnce(mockValue2);
+        mockFunc.mockReturnValue(mockValue1);
         const TestComponent = () => {
           const val = hookUnderTest('flag', defaultValue);
           return <p>{JSON.stringify(val)}</p>;
         };
 
-        const manager = new ClientManager(fakeClient);
-        manager.ref();
+        const featureStore = FeaturesStore.forClient(fakeClient);
 
         const { unmount } = render(
-          <ClientManagerContext.Provider value={() => manager}>
+          <FeatureStoreContext.Provider value={featureStore}>
             <React.Suspense fallback={<p>suspense</p>}>
               <TestComponent />
             </React.Suspense>
-          </ClientManagerContext.Provider>,
+          </FeatureStoreContext.Provider>,
         );
 
         expect(screen.getByText('suspense')).toBeInTheDocument();
 
-        await act(async () => {
-          fireReady!();
-          await manager.promise;
-        });
+        fire(ProviderEvents.Ready);
 
-        expect(screen.getByText(JSON.stringify(mockValue1))).toBeInTheDocument();
+        expect(await screen.findByText(JSON.stringify(mockValue1))).toBeInTheDocument();
 
-        await act(async () => {
-          fireStale!();
-          fireReady!();
+        act(() => {
+          fire(ProviderEvents.Stale);
+          mockFunc.mockReturnValue(mockValue2);
+          fire(ProviderEvents.Ready);
         });
 
-        expect(screen.getByText(JSON.stringify(mockValue2))).toBeInTheDocument();
+        expect(await screen.findByText(JSON.stringify(mockValue2))).toBeInTheDocument();
 
         unmount();
 
-        manager.unref();
-        expect(fakeClient.addHandler).toBeCalledTimes(3);
-        expect(fakeClient.removeHandler).toBeCalledTimes(3);
+        expect(handlerMap.size).toBe(0);
       },
     );
   });
@@ -169,31 +165,27 @@ describe('hooks', () => {
         return <p>{JSON.stringify(val)}</p>;
       };
 
-      const manager = new ClientManager(fakeClient);
-      manager.ref();
+      const featureStore = FeaturesStore.forClient(fakeClient);
 
       const { unmount } = render(
-        <ClientManagerContext.Provider value={() => manager}>
+        <FeatureStoreContext.Provider value={featureStore}>
           <React.Suspense fallback={<p>suspense</p>}>
             <TestComponent />
           </React.Suspense>
-        </ClientManagerContext.Provider>,
+        </FeatureStoreContext.Provider>,
       );
 
       expect(screen.getByText('suspense')).toBeInTheDocument();
 
-      await act(async () => {
-        fireError!();
-        await manager.promise;
+      act(() => {
+        fire(ProviderEvents.Error);
       });
 
-      expect(screen.getByText(JSON.stringify(defaultValue))).toBeInTheDocument();
+      expect(await screen.findByText(JSON.stringify(defaultValue))).toBeInTheDocument();
 
       unmount();
 
-      manager.unref();
-      expect(fakeClient.addHandler).toBeCalledTimes(3);
-      expect(fakeClient.removeHandler).toBeCalledTimes(3);
+      expect(handlerMap.size).toBe(0);
     });
   });
 });
diff --git a/packages/integration-react/src/ReactAdaptor.tsx b/packages/integration-react/src/ReactAdaptor.tsx
index 03835478..8dfcece9 100644
--- a/packages/integration-react/src/ReactAdaptor.tsx
+++ b/packages/integration-react/src/ReactAdaptor.tsx
@@ -1,206 +1,134 @@
-import React, { EffectCallback, createContext, useContext, useEffect, useState, useMemo } from 'react';
+import React, { createContext, useContext, useMemo } from 'react';
+import { useSyncExternalStore } from 'use-sync-external-store/shim';
 import {
   Client,
   EvaluationDetails,
-  EventHandler,
+  Features,
   JsonValue,
   OpenFeature,
   ProviderEvents,
   ResolutionDetails,
 } from '@openfeature/web-sdk';
 
-type SetState = () => void;
-
-export class ClientManager {
-  private readonly onReady: EventHandler<ProviderEvents>;
-  private readonly onStale: EventHandler<ProviderEvents>;
-  private readonly onError: EventHandler<ProviderEvents>;
-  public readonly effect: EffectCallback;
-  private references: number = 0;
-  promise?: Promise<void>;
-  private onChangeHandlers: Set<SetState> = new Set();
-
-  constructor(private readonly instance: Client) {
-    let resolve = this.createPromise();
-
-    this.onReady = this.onError = () => {
-      this.promise = undefined;
-      resolve();
-      for (const cb of this.onChangeHandlers) {
-        try {
-          cb();
-        } catch (e) {
-          // do nothing
-        }
+export interface FeaturesStore {
+  subscribe: (onChange: () => void) => () => void;
+  getSnapshot: () => Features;
+}
+
+export namespace FeaturesStore {
+  export function forClient(client: Client): FeaturesStore {
+    let [readyPromise, resolveReady] = createPromise();
+    const changeHandlers = new Set<() => void>();
+    let isReady = false;
+
+    const onReadyOrError = () => {
+      isReady = true;
+      resolveReady();
+      for (const handler of changeHandlers) {
+        handler();
       }
-      this.onChangeHandlers.clear();
     };
-    this.onStale = () => {
-      resolve = this.createPromise();
+
+    const onStale = () => {
+      isReady = false;
+      [readyPromise, resolveReady] = createPromise();
     };
 
-    this.effect = () => {
-      this.ref();
+    const subscribe = (onChange: () => void) => {
+      if (changeHandlers.has(onChange)) {
+        throw new Error('Already subscribed');
+      }
+      changeHandlers.add(onChange);
+      if (changeHandlers.size === 1) {
+        client.addHandler(ProviderEvents.Ready, onReadyOrError);
+        client.addHandler(ProviderEvents.Error, onReadyOrError);
+        client.addHandler(ProviderEvents.Stale, onStale);
+      }
+
       return () => {
-        this.unref();
+        if (!changeHandlers.has(onChange)) {
+          throw new Error('Already unsubscribed');
+        }
+        changeHandlers.delete(onChange);
+        if (changeHandlers.size === 0) {
+          client.removeHandler(ProviderEvents.Ready, onReadyOrError);
+          client.removeHandler(ProviderEvents.Error, onReadyOrError);
+          client.removeHandler(ProviderEvents.Stale, onStale);
+        }
       };
     };
-  }
 
-  onceOnChange(callback: SetState) {
-    this.onChangeHandlers.add(callback);
-  }
-
-  getClient(): Client {
-    if (this.promise) {
-      throw this.promise;
-    }
-    return this.instance;
-  }
-
-  public ref() {
-    if (this.references++ === 0) {
-      this.instance.addHandler(ProviderEvents.Error, this.onError);
-      this.instance.addHandler(ProviderEvents.Ready, this.onReady);
-      this.instance.addHandler(ProviderEvents.Stale, this.onStale);
-    }
-  }
-
-  public unref() {
-    if (--this.references === 0) {
-      this.instance.removeHandler(ProviderEvents.Error, this.onError);
-      this.instance.removeHandler(ProviderEvents.Ready, this.onReady);
-      this.instance.removeHandler(ProviderEvents.Stale, this.onStale);
-    }
-  }
-
-  private createPromise(): () => void {
-    let resolver: () => void;
-    this.promise = new Promise(resolve => {
-      resolver = resolve;
-    });
-    return () => {
-      resolver();
+    const getSnapshot = () => {
+      if (!isReady) {
+        readyPromise.then(subscribe(() => {}));
+        throw readyPromise;
+      }
+      return client;
     };
+
+    return { subscribe, getSnapshot };
   }
 }
 
-let defaultClientManager: ClientManager | undefined;
+const defaultFeatureStore = FeaturesStore.forClient(OpenFeature.getClient());
 
-// exported for tests
-export const ClientManagerContext = createContext<() => ClientManager>(() => {
-  if (!defaultClientManager) {
-    defaultClientManager = new ClientManager(OpenFeature.getClient());
-    defaultClientManager.ref();
-  }
-  return defaultClientManager;
-});
+export const FeatureStoreContext = createContext(defaultFeatureStore);
 
 type OpenFeatureContextProviderProps = {
-  name: string;
+  client?: Client;
 };
-export const OpenFeatureContextProvider: React.FC<React.PropsWithChildren<OpenFeatureContextProviderProps>> = ({
-  children,
-  name,
-}) => {
-  const manager = useMemo(() => new ClientManager(OpenFeature.getClient(name)), [name]);
-  useEffect(() => manager.effect(), [manager]);
-
-  return <ClientManagerContext.Provider value={() => manager}>{children}</ClientManagerContext.Provider>;
+export const OpenFeatureContextProvider: React.FC<React.PropsWithChildren<OpenFeatureContextProviderProps>> = props => {
+  const { client = OpenFeature.getClient(), children } = props;
+  const featuresStore = useMemo(() => FeaturesStore.forClient(client), [client]);
+
+  return <FeatureStoreContext.Provider value={featuresStore}>{children}</FeatureStoreContext.Provider>;
 };
 
 export function useStringValue(flagKey: string, defaultValue: string): string {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getStringValue(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getStringValue(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getStringValue(flagKey, defaultValue));
 }
 
 export function useStringDetails(flagKey: string, defaultValue: string): ResolutionDetails<string> {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getStringDetails(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getStringDetails(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getStringDetails(flagKey, defaultValue));
 }
 
 export function useBooleanValue(flagKey: string, defaultValue: boolean): boolean {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getBooleanValue(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getBooleanValue(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getBooleanValue(flagKey, defaultValue));
 }
 
 export function useBooleanDetails(flagKey: string, defaultValue: boolean): ResolutionDetails<boolean> {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getBooleanDetails(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getBooleanDetails(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getBooleanDetails(flagKey, defaultValue));
 }
 
 export function useNumberValue(flagKey: string, defaultValue: number): number {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getNumberValue(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getNumberValue(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getNumberValue(flagKey, defaultValue));
 }
 
 export function useNumberDetails(flagKey: string, defaultValue: number): EvaluationDetails<number> {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getNumberDetails(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getNumberDetails(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getNumberDetails(flagKey, defaultValue));
 }
 
 export function useObjectValue<T extends JsonValue>(flagKey: string, defaultValue: T): T {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getObjectValue(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getObjectValue(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
-
-  return value as T;
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getObjectValue<T>(flagKey, defaultValue));
 }
 
 export function useObjectDetails<T extends JsonValue>(flagKey: string, defaultValue: T): EvaluationDetails<T> {
-  const manager = useContext(ClientManagerContext)();
-  const client = manager.getClient();
-  const [value, setValue] = useState(() => client.getObjectDetails(flagKey, defaultValue));
-
-  useEffect(() => {
-    manager.onceOnChange(() => setValue(client.getObjectDetails(flagKey, defaultValue)));
-  }, [defaultValue, flagKey, manager, client]);
+  const store = useContext(FeatureStoreContext);
+  return useSyncExternalStore(store.subscribe, () => store.getSnapshot().getObjectDetails<T>(flagKey, defaultValue));
+}
 
-  return value as EvaluationDetails<T>;
+function createPromise<T = void>(): [Promise<T>, (value: T) => void, (reason: any) => void] {
+  let resolve: (value: T) => void;
+  let reject: (reason: any) => void;
+  const promise = new Promise<T>((...args) => {
+    [resolve, reject] = args;
+  });
+  return [promise, resolve!, reject!];
 }