useRecoilCallback() snapshot gets the latest state (#1610)

Summary:
Pull Request resolved: facebookexperimental/Recoil#1610

NOTE: This is a breaking change, but the previous approach had a bug.

Previously, `useRecoilCallback()` would attempt to provide a `Snapshot` based on the state that was currently rendered.  However, there were some bugs where this was not the case.  For example, a callback called from an effect handler may or may not have had the latest rendered state based on the timing of the setter.

The new approach `useRecoilCallback()` will always try to get a snapshot with the latest state.  This should also be more intuitive for users.

Addresses #1604

Reviewed By: habond

Differential Revision: D34236485

fbshipit-source-id: 0cb17cec40b50e8debeb57dc193b9f410d170fc5

CHANGELOG.md

@@ -1,13 +1,15 @@
 # Change Log
 
 ## UPCOMING
+**_Add new changes here as they land_**
 
+- `shouldNotBeFrozen` now works in JS environment without `Window` interface. (#1571)
 - Avoid spurious console errors from effects when calling `setSelf()` from `onSet()` handlers. (#1589)
 - Better error reporting when selectors provide inconsistent results (#1696)
 
-**_Add new changes here as they land_**
+### Breaking Changes
 
-- `shouldNotBeFrozen` now works in JS environment without `Window` interface. (#1571)
+- `useRecoilCallback()` now provides a snapshot for the latest state instead of the latest rendered state, which had bugs (#1610, #1604)
 
 ## 0.6.1 (2022-01-29)
 

packages/recoil/core/Recoil_RecoilValueInterface.js

@@ -35,6 +35,7 @@ const {
   RecoilValueReadOnly,
   isRecoilValue,
 } = require('./Recoil_RecoilValue');
+const {invalidateMemoizedSnapshot} = require('./Recoil_SnapshotCache');
 const nullthrows = require('recoil-shared/util/Recoil_nullthrows');
 const recoverableViolation = require('recoil-shared/util/Recoil_recoverableViolation');
 
@@ -195,6 +196,7 @@ function applyActionsToStore(store, actions) {
       applyAction(store, newState, action);
     }
     invalidateDownstreams(store, newState);
+    invalidateMemoizedSnapshot();
     return newState;
   });
 }
@@ -369,5 +371,4 @@ module.exports = {
   invalidateDownstreams,
   copyTreeState,
   refreshRecoilValue,
-  invalidateDownstreams_FOR_TESTING: invalidateDownstreams,
 };

packages/recoil/core/Recoil_Snapshot.js

@@ -37,6 +37,7 @@ const {
   setUnvalidatedRecoilValue,
 } = require('./Recoil_RecoilValueInterface');
 const {updateRetainCount} = require('./Recoil_Retention');
+const {setInvalidateMemoizedSnapshot} = require('./Recoil_SnapshotCache');
 const {
   getNextTreeStateVersion,
   makeEmptyStoreState,
@@ -332,21 +333,24 @@ const [memoizedCloneSnapshot, invalidateMemoizedSnapshot] =
     (store, version) => {
       const storeState = store.getState();
       const treeState =
-        version === 'current'
-          ? storeState.currentTree
+        version === 'latest'
+          ? storeState.nextTree ?? storeState.currentTree
           : nullthrows(storeState.previousTree);
       return new Snapshot(cloneStoreState(store, treeState));
     },
     (store, version) =>
       String(version) +
       String(store.storeID) +
+      String(store.getState().nextTree?.version) +
       String(store.getState().currentTree.version) +
       String(store.getState().previousTree?.version),
   );
+// Avoid circular dependencies
+setInvalidateMemoizedSnapshot(invalidateMemoizedSnapshot);
 
 function cloneSnapshot(
   store: Store,
-  version: 'current' | 'previous' = 'current',
+  version: 'latest' | 'previous' = 'latest',
 ): Snapshot {
   const snapshot = memoizedCloneSnapshot(store, version);
   if (!snapshot.isRetained()) {

packages/recoil/core/Recoil_SnapshotCache.js

@@ -0,0 +1,26 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails oncall+recoil
+ * @flow strict-local
+ * @format
+ */
+'use strict';
+
+let _invalidateMemoizedSnapshot: ?() => void = null;
+
+function setInvalidateMemoizedSnapshot(invalidate: () => void): void {
+  _invalidateMemoizedSnapshot = invalidate;
+}
+
+function invalidateMemoizedSnapshot(): void {
+  _invalidateMemoizedSnapshot?.();
+}
+
+module.exports = {
+  setInvalidateMemoizedSnapshot,
+  invalidateMemoizedSnapshot,
+};

packages/recoil/hooks/Recoil_SnapshotHooks.js

@@ -158,7 +158,7 @@ function useRecoilTransactionObserver(
   useTransactionSubscription(
     useCallback(
       store => {
-        const snapshot = cloneSnapshot(store, 'current');
+        const snapshot = cloneSnapshot(store, 'latest');
         const previousSnapshot = cloneSnapshot(store, 'previous');
         callback({
           snapshot,

packages/recoil/hooks/Recoil_useRecoilCallback.js

@@ -12,6 +12,7 @@
 
 import type {TransactionInterface} from '../core/Recoil_AtomicUpdates';
 import type {RecoilState, RecoilValue} from '../core/Recoil_RecoilValue';
+import type {Snapshot} from '../core/Recoil_Snapshot';
 import type {Store} from '../core/Recoil_State';
 
 const {atomicUpdater} = require('../core/Recoil_AtomicUpdates');
@@ -22,7 +23,7 @@ const {
   refreshRecoilValue,
   setRecoilValue,
 } = require('../core/Recoil_RecoilValueInterface');
-const {Snapshot, cloneSnapshot} = require('../core/Recoil_Snapshot');
+const {cloneSnapshot} = require('../core/Recoil_Snapshot');
 const {gotoSnapshot} = require('./Recoil_SnapshotHooks');
 const {useCallback} = require('react');
 const err = require('recoil-shared/util/Recoil_err');

packages/recoil/hooks/__tests__/Recoil_useRecoilCallback-test.js

@@ -17,6 +17,7 @@ const {
 let React,
   useRef,
   useState,
+  useEffect,
   act,
   useStoreRef,
   atom,
@@ -30,11 +31,12 @@ let React,
   ReadsAtom,
   flushPromisesAndTimers,
   renderElements,
+  stringAtom,
   invariant;
 
 const testRecoil = getRecoilTestFn(() => {
   React = require('react');
-  ({useRef, useState} = require('react'));
+  ({useRef, useState, useEffect} = require('react'));
   ({act} = require('ReactTestUtils'));
 
   ({useStoreRef} = require('../../core/Recoil_RecoilRoot'));
@@ -52,6 +54,7 @@ const testRecoil = getRecoilTestFn(() => {
     ReadsAtom,
     flushPromisesAndTimers,
     renderElements,
+    stringAtom,
   } = require('recoil-shared/__test_utils__/Recoil_TestingUtils'));
   invariant = require('recoil-shared/util/Recoil_invariant');
 });
@@ -235,16 +238,16 @@ testRecoil('Reads from a snapshot created at callback call time', async () => {
   expect(seenValue).toBe(345);
 });
 
-testRecoil('Setter updater sees current state', () => {
+testRecoil('Setter updater sees latest state', () => {
   const myAtom = atom({key: 'useRecoilCallback updater', default: 'DEFAULT'});
 
   let setAtom;
   let cb;
   function Component() {
     setAtom = useSetRecoilState(myAtom);
     cb = useRecoilCallback(({snapshot, set}) => prevValue => {
-      // snapshot sees the stable snapshot from batch at beginning of transaction
-      expect(snapshot.getLoadable(myAtom).contents).toEqual('DEFAULT');
+      // snapshot sees a snapshot with the latest set state
+      expect(snapshot.getLoadable(myAtom).contents).toEqual(prevValue);
 
       // Test that callback sees value updates from within the same transaction
       set(myAtom, value => {
@@ -268,7 +271,7 @@ testRecoil('Setter updater sees current state', () => {
 
   expect(c.textContent).toEqual('"DEFAULT"');
 
-  // Set then callback in the same transaction
+  // Set and callback in the same transaction
   act(() => {
     setAtom('SET');
     cb('SET');
@@ -277,6 +280,49 @@ testRecoil('Setter updater sees current state', () => {
   expect(c.textContent).toEqual('"UPDATE AGAIN"');
 });
 
+testRecoil('Snapshot from effect uses rendered state', () => {
+  const myAtom = stringAtom();
+  let setState,
+    actCallback,
+    effectCallback,
+    actCallbackValue,
+    effectCallbackValue,
+    effectValue;
+  function Component() {
+    setState = useSetRecoilState(myAtom);
+    const value = useRecoilValue(myAtom);
+    effectCallback = useRecoilCallback(
+      ({snapshot}) =>
+        () => {
+          effectCallbackValue = snapshot.getLoadable(myAtom).getValue();
+        },
+      [],
+    );
+    actCallback = useRecoilCallback(
+      ({snapshot}) =>
+        () => {
+          actCallbackValue = snapshot.getLoadable(myAtom).getValue();
+        },
+      [],
+    );
+
+    useEffect(() => {
+      effectValue = value;
+      effectCallback();
+    }, [value]);
+    return null;
+  }
+
+  renderElements(<Component />);
+  act(() => {
+    setState('SET');
+    actCallback();
+  });
+  expect(actCallbackValue).toBe('SET');
+  expect(effectValue).toBe('SET');
+  expect(effectCallbackValue).toBe('SET');
+});
+
 testRecoil('goes to snapshot', async () => {
   const myAtom = atom({
     key: 'Goto Snapshot From Callback',

packages/shared/__test_utils__/Recoil_TestingUtils.js

@@ -33,7 +33,7 @@ const {
   sendEndOfBatchNotifications_FOR_TESTING,
 } = require('../../recoil/core/Recoil_RecoilRoot');
 const {
-  invalidateDownstreams_FOR_TESTING,
+  invalidateDownstreams,
 } = require('../../recoil/core/Recoil_RecoilValueInterface');
 const {makeEmptyStoreState} = require('../../recoil/core/Recoil_State');
 const invariant = require('../util/Recoil_invariant');
@@ -65,7 +65,7 @@ function makeStore(): Store {
       const currentStoreState = store.getState();
       // FIXME: does not increment state version number
       currentStoreState.currentTree = replacer(currentStoreState.currentTree); // no batching so nextTree is never active
-      invalidateDownstreams_FOR_TESTING(store, currentStoreState.currentTree);
+      invalidateDownstreams(store, currentStoreState.currentTree);
       const {reactMode} = require('../../recoil/core/Recoil_ReactMode');
       if (reactMode().early) {
         notifyComponents_FOR_TESTING(