Better fix for StrictMode

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/core/__tests__/Recoil_RecoilRoot-test.js

@@ -100,7 +100,9 @@ describe('initializeState', () => {
     expect(container.textContent).toEqual('"INITIALIZE""INITIALIZE"');
   });
 
-  testRecoil('Atom Effects run with global initialization', () => {
+  testRecoil('Atom Effects run with global initialization', ({strictMode}) => {
+    const sm = strictMode ? 2 : 1;
+
     let effectRan = 0;
     const myAtom = atom<string>({
       key: 'RecoilRoot - initializeState - atom effects',
@@ -132,7 +134,7 @@ describe('initializeState', () => {
     );
     // Effects are run when initialized with initializeState, even if not read.
     expect(container1.textContent).toEqual('NO READ');
-    expect(effectRan).toEqual(1);
+    expect(effectRan).toEqual(1 * sm);
 
     const container2 = renderElements(
       <RecoilRoot initializeState={initializeState}>
@@ -142,7 +144,7 @@ describe('initializeState', () => {
 
     // Effects are run first, initializeState() takes precedence
     expect(container2.textContent).toEqual('"INITIALIZE"');
-    expect(effectRan).toEqual(2);
+    expect(effectRan).toEqual(2 * sm);
   });
 
   testRecoil('initialize with nested store', () => {
@@ -170,6 +172,56 @@ describe('initializeState', () => {
 
     expect(container.textContent).toEqual('NESTED_ROOT/ROOT');
   });
+
+  testRecoil('initializeState is only called once', ({strictMode}) => {
+    if (strictMode) {
+      return;
+    }
+    const myAtom = atom({
+      key: 'RecoilRoot/override/atom',
+      default: 'DEFAULT',
+    });
+    const [ReadsWritesAtom, setAtom] = componentThatReadsAndWritesAtom(myAtom);
+
+    const initializeState = jest.fn(({set}) => set(myAtom, 'INIT'));
+    let forceUpdate = () => {
+      throw new Error('not rendered');
+    };
+    let setRootKey = _ => {
+      throw new Error('');
+    };
+    function MyRoot() {
+      const [counter, setCounter] = useState(0);
+      forceUpdate = () => setCounter(counter + 1);
+      const [key, setKey] = useState(0);
+      setRootKey = setKey;
+      return (
+        <RecoilRoot key={key} initializeState={initializeState}>
+          {counter}
+          <ReadsWritesAtom />
+        </RecoilRoot>
+      );
+    }
+    const container = renderElements(<MyRoot />);
+
+    expect(container.textContent).toEqual('0"INIT"');
+
+    act(forceUpdate);
+    expect(initializeState).toHaveBeenCalledTimes(1);
+    expect(container.textContent).toEqual('1"INIT"');
+
+    act(() => setAtom('SET'));
+    expect(initializeState).toHaveBeenCalledTimes(1);
+    expect(container.textContent).toEqual('1"SET"');
+
+    act(forceUpdate);
+    expect(initializeState).toHaveBeenCalledTimes(1);
+    expect(container.textContent).toEqual('2"SET"');
+
+    act(() => setRootKey(1));
+    expect(initializeState).toHaveBeenCalledTimes(2);
+    expect(container.textContent).toEqual('2"INIT"');
+  });
 });
 
 testRecoil(
@@ -332,51 +384,4 @@ describe('override prop', () => {
       expect(container.textContent).toEqual('"SET"');
     },
   );
-
-  testRecoil('initializeState is only called once', () => {
-    const myAtom = atom({
-      key: 'RecoilRoot/override/atom',
-      default: 'DEFAULT',
-    });
-    const [ReadsWritesAtom, setAtom] = componentThatReadsAndWritesAtom(myAtom);
-
-    const initializeState = jest.fn(({set}) => set(myAtom, 'INIT'));
-    let forceUpdate = () => {
-      throw new Error('not rendered');
-    };
-    let setRootKey = _ => {
-      throw new Error('');
-    };
-    function MyRoot() {
-      const [counter, setCounter] = useState(0);
-      forceUpdate = () => setCounter(counter + 1);
-      const [key, setKey] = useState(0);
-      setRootKey = setKey;
-      return (
-        <RecoilRoot key={key} initializeState={initializeState}>
-          {counter}
-          <ReadsWritesAtom />
-        </RecoilRoot>
-      );
-    }
-    const container = renderElements(<MyRoot />);
-
-    expect(container.textContent).toEqual('0"INIT"');
-
-    act(forceUpdate);
-    expect(initializeState).toHaveBeenCalledTimes(1);
-    expect(container.textContent).toEqual('1"INIT"');
-
-    act(() => setAtom('SET'));
-    expect(initializeState).toHaveBeenCalledTimes(1);
-    expect(container.textContent).toEqual('1"SET"');
-
-    act(forceUpdate);
-    expect(initializeState).toHaveBeenCalledTimes(1);
-    expect(container.textContent).toEqual('2"SET"');
-
-    act(() => setRootKey(1));
-    expect(initializeState).toHaveBeenCalledTimes(2);
-    expect(container.textContent).toEqual('2"INIT"');
-  });
 });