Better fix for StrictMode (#1473)

Summary:
Pull Request resolved: #1473

A better fix for React StrictMode.  Instead of maintaining a ref in `<RecoilRoot>` for the set of atoms to re-initialize we will just keep the atoms in `knownAtoms` even when they are cleaned up.  This way we maintain the set of known atoms we need to re-initialize.  However, we have to be careful to ensure the we actually re-initialize the atom when the effect runs again, even if it is already in `knownAtoms`.  We also have to be careful when using an initialized atom in a `Snapshot` that it does not try to re-initialize.

Reviewed By: davidmccabe

Differential Revision: D32751941

fbshipit-source-id: 7b36a6c6ed344cc899b48d649ea9dafff5af004c

packages/recoil/core/Recoil_FunctionalCore.js

@@ -101,10 +101,20 @@ function initializeNodeIfNewToStore(
   });
 }
 
-function initNode(store: Store, key: NodeKey): void {
+function initializeNode(store: Store, key: NodeKey): void {
   initializeNodeIfNewToStore(store, store.getState().currentTree, key, 'get');
 }
 
+function reinitializeNode(store: Store, key: NodeKey): void {
+  const storeState = store.getState();
+  // If this atom was previously initialized (set in knownAtoms), but was
+  // cleaned up (not set in nodeCleanupFunctions), then re-initialize it.
+  if (!storeState.nodeCleanupFunctions.has(key)) {
+    storeState.knownAtoms.delete(key); // Force atom to re-initialize
+  }
+  initializeNodeIfNewToStore(store, storeState.currentTree, key, 'get');
+}
+
 function cleanUpNode(store: Store, key: NodeKey) {
   const state = store.getState();
   state.nodeCleanupFunctions.get(key)?.();
@@ -251,7 +261,8 @@ module.exports = {
   getNodeLoadable,
   peekNodeLoadable,
   setNodeValue,
-  initNode,
+  initializeNode,
+  reinitializeNode,
   cleanUpNode,
   setUnvalidatedAtomValue_DEPRECATED,
   peekNodeInfo,

packages/recoil/core/Recoil_RecoilRoot.react.js

@@ -10,7 +10,7 @@
  */
 'use strict';
 
-import type {NodeKey, StoreID} from './Recoil_Keys';
+import type {StoreID} from './Recoil_Keys';
 import type {RecoilValue} from './Recoil_RecoilValue';
 import type {MutableSnapshot} from './Recoil_Snapshot';
 import type {Store, StoreRef, StoreState, TreeState} from './Recoil_State';
@@ -27,7 +27,7 @@ const {
 const {
   cleanUpNode,
   getDownstreamNodes,
-  initNode,
+  reinitializeNode,
   setNodeValue,
   setUnvalidatedAtomValue_DEPRECATED,
 } = require('./Recoil_FunctionalCore');
@@ -483,32 +483,22 @@ function RecoilRoot_INTERNAL({
   );
 
   // Cleanup when the <RecoilRoot> is unmounted
-  const cleanedUpNodesRef = useRefInitOnce<Map<StoreID, Set<NodeKey>>>(
-    () => new Map(),
-  );
   useEffect(() => {
     // React is free to call effect cleanup handlers and effects at will, the
     // deps array is only an optimization.  For example, React strict mode
     // will execute each effect twice for testing.  Therefore, we need symmetry
     // to re-initialize all known atoms after they were cleaned up.
     const store = storeRef.current;
-    for (const atomKey of cleanedUpNodesRef.current.get(store.storeID) ?? []) {
-      initNode(store, atomKey);
+    for (const atomKey of new Set(store.getState().knownAtoms)) {
+      reinitializeNode(store, atomKey);
     }
-    cleanedUpNodesRef.current.delete(store.storeID);
 
     return () => {
-      const {knownAtoms} = store.getState();
-      // Save the set of known atoms from this cleanup in case the effect is just
-      // being cleaned up and then executed again so we can re-initialize the
-      // atoms above.
-      // eslint-disable-next-line fb-www/react-hooks-deps
-      cleanedUpNodesRef.current.set(store.storeID, new Set(knownAtoms));
-      for (const atomKey of knownAtoms) {
-        cleanUpNode(storeRef.current, atomKey);
+      for (const atomKey of store.getState().knownAtoms) {
+        cleanUpNode(store, atomKey);
       }
     };
-  }, [storeRef, cleanedUpNodesRef]);
+  }, [storeRef]);
 
   return (
     <AppContext.Provider value={storeRef}>

packages/recoil/core/Recoil_Snapshot.js

@@ -22,7 +22,7 @@ import type {RecoilState, RecoilValue} from './Recoil_RecoilValue';
 import type {StateID, Store, StoreState, TreeState} from './Recoil_State';
 
 const {batchUpdates} = require('./Recoil_Batching');
-const {initNode, peekNodeInfo} = require('./Recoil_FunctionalCore');
+const {initializeNode, peekNodeInfo} = require('./Recoil_FunctionalCore');
 const {graph} = require('./Recoil_Graph');
 const {getNextStoreID} = require('./Recoil_Keys');
 const {
@@ -96,7 +96,7 @@ class Snapshot {
     // Initialize any nodes that are live in the parent store (primarily so that this
     // snapshot gets counted towards the node's live stores count).
     for (const nodeKey of this._store.getState().knownAtoms) {
-      initNode(this._store, nodeKey);
+      initializeNode(this._store, nodeKey);
       updateRetainCount(this._store, nodeKey, 1);
     }
     this.retain();

packages/recoil/recoil_values/Recoil_atom.js

@@ -229,9 +229,17 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
     initState: TreeState,
     trigger: Trigger,
   ): () => void {
-    liveStoresCount++;
-    const alreadyKnown = store.getState().knownAtoms.has(key);
+    const cleanupAtom = () => {
+      liveStoresCount--;
+      cleanupEffectsByStore.get(store)?.forEach(cleanup => cleanup());
+      cleanupEffectsByStore.delete(store);
+    };
+
+    if (store.getState().knownAtoms.has(key)) {
+      return cleanupAtom;
+    }
     store.getState().knownAtoms.add(key);
+    liveStoresCount++;
 
     // Setup async defaults to notify subscribers when they resolve
     if (defaultLoadable.state === 'loading') {
@@ -246,17 +254,19 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
         .catch(notifyDefaultSubscribers);
     }
 
+    ///////////////////
     // Run Atom Effects
+    ///////////////////
 
-    // This state is scoped by Store, since this is in the initAtom() closure
-    let initValue: NewValue<T> = DEFAULT_VALUE;
-    let pendingSetSelf: ?{
-      effect: AtomEffect<T>,
-      value: T | DefaultValue,
-    } = null;
-
-    if (options.effects_UNSTABLE != null && !alreadyKnown) {
+    if (options.effects_UNSTABLE != null) {
+      // This state is scoped by Store, since this is in the initAtom() closure
       let duringInit = true;
+      let initValue: NewValue<T> = DEFAULT_VALUE;
+      let isInitError: boolean = false;
+      let pendingSetSelf: ?{
+        effect: AtomEffect<T>,
+        value: T | DefaultValue,
+      } = null;
 
       function getLoadable<S>(recoilValue: RecoilValue<S>): Loadable<S> {
         // Normally we can just get the current value of another atom.
@@ -394,49 +404,51 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
         };
 
       for (const effect of options.effects_UNSTABLE ?? []) {
-        const cleanup = effect({
-          node,
-          storeID: store.storeID,
-          trigger,
-          setSelf: setSelf(effect),
-          resetSelf: resetSelf(effect),
-          onSet: onSet(effect),
-          getPromise,
-          getLoadable,
-          getInfo_UNSTABLE,
-        });
-        if (cleanup != null) {
-          cleanupEffectsByStore.set(store, [
-            ...(cleanupEffectsByStore.get(store) ?? []),
-            cleanup,
-          ]);
+        try {
+          const cleanup = effect({
+            node,
+            storeID: store.storeID,
+            trigger,
+            setSelf: setSelf(effect),
+            resetSelf: resetSelf(effect),
+            onSet: onSet(effect),
+            getPromise,
+            getLoadable,
+            getInfo_UNSTABLE,
+          });
+          if (cleanup != null) {
+            cleanupEffectsByStore.set(store, [
+              ...(cleanupEffectsByStore.get(store) ?? []),
+              cleanup,
+            ]);
+          }
+        } catch (error) {
+          initValue = error;
+          isInitError = true;
         }
       }
 
       duringInit = false;
-    }
 
-    // Mutate initial state in place since we know there are no other subscribers
-    // since we are the ones initializing on first use.
-    if (!(initValue instanceof DefaultValue)) {
-      const frozenInitValue = maybeFreezeValueOrPromise(initValue);
-      const initLoadable = isPromise(frozenInitValue)
-        ? loadableWithPromise(wrapPendingPromise(store, frozenInitValue))
-        : loadableWithValue(frozenInitValue);
-      initState.atomValues.set(key, initLoadable);
-
-      // If there is a pending transaction, then also mutate the next state tree.
-      // This could happen if the atom was first initialized in an action that
-      // also updated some other atom's state.
-      store.getState().nextTree?.atomValues.set(key, initLoadable);
+      // Mutate initial state in place since we know there are no other subscribers
+      // since we are the ones initializing on first use.
+      if (!(initValue instanceof DefaultValue)) {
+        const frozenInitValue = maybeFreezeValueOrPromise(initValue);
+        const initLoadable = isInitError
+          ? loadableWithError(initValue)
+          : isPromise(frozenInitValue)
+          ? loadableWithPromise(wrapPendingPromise(store, frozenInitValue))
+          : loadableWithValue(frozenInitValue);
+        initState.atomValues.set(key, initLoadable);
+
+        // If there is a pending transaction, then also mutate the next state tree.
+        // This could happen if the atom was first initialized in an action that
+        // also updated some other atom's state.
+        store.getState().nextTree?.atomValues.set(key, initLoadable);
+      }
     }
 
-    return () => {
-      liveStoresCount--;
-      cleanupEffectsByStore.get(store)?.forEach(cleanup => cleanup());
-      cleanupEffectsByStore.delete(store);
-      store.getState().knownAtoms.delete(key);
-    };
+    return cleanupAtom;
   }
 
   function peekAtom(_store, state: TreeState): ?Loadable<T> {

packages/recoil/recoil_values/__tests__/Recoil_atom-test.js

@@ -928,6 +928,16 @@ describe('Effects', () => {
     expect(c.textContent).toBe('');
     expect(refCountsA).toEqual([0, 0]);
     expect(refCountsB).toEqual([0, 0]);
+
+    act(() => setNumRoots(1));
+    expect(c.textContent).toBe('"A""B"');
+    expect(refCountsA).toEqual([1, 1]);
+    expect(refCountsB).toEqual([1, 1]);
+
+    act(() => setNumRoots(0));
+    expect(c.textContent).toBe('');
+    expect(refCountsA).toEqual([0, 0]);
+    expect(refCountsB).toEqual([0, 0]);
   });
 
   // Test that effects can initialize state when an atom is first used after an