Retain snapshots for async callbacks

CHANGELOG.md

@@ -5,6 +5,7 @@
 
 - The `default` value is now optional for `atom()` and `atomFamily()`.  If not provided the atom will initialize to a pending state. (#1639)
 - Significant optimization for selector evaluations.  2x improvement with 100 dependencies, 4x with 1,000, and 40x with 10,000. (#1515, #914)
+- Automatically retain snapshots for the duration of async callbacks. (#1632)
 - `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)

packages/recoil/hooks/Recoil_useRecoilCallback.js

@@ -28,6 +28,7 @@ const {gotoSnapshot} = require('./Recoil_SnapshotHooks');
 const {useCallback} = require('react');
 const err = require('recoil-shared/util/Recoil_err');
 const invariant = require('recoil-shared/util/Recoil_invariant');
+const isPromise = require('recoil-shared/util/Recoil_isPromise');
 const lazyProxy = require('recoil-shared/util/Recoil_lazyProxy');
 
 export type RecoilCallbackInterface = $ReadOnly<{
@@ -49,6 +50,7 @@ function recoilCallback<Args: $ReadOnlyArray<mixed>, Return, ExtraInterface>(
   extraInterface?: ExtraInterface,
 ): Return {
   let ret: Return | Sentinel = SENTINEL;
+  let releaseSnapshot;
   batchUpdates(() => {
     const errMsg =
       'useRecoilCallback() expects a function that returns a function: ' +
@@ -79,7 +81,11 @@ function recoilCallback<Args: $ReadOnlyArray<mixed>, Return, ExtraInterface>(
         transact_UNSTABLE: transaction => atomicUpdater(store)(transaction),
       },
       {
-        snapshot: () => cloneSnapshot(store),
+        snapshot: () => {
+          const snapshot = cloneSnapshot(store);
+          releaseSnapshot = snapshot.retain();
+          return snapshot;
+        },
       },
     );
 
@@ -93,6 +99,13 @@ function recoilCallback<Args: $ReadOnlyArray<mixed>, Return, ExtraInterface>(
     !(ret instanceof Sentinel),
     'batchUpdates should return immediately',
   );
+  if (isPromise(ret)) {
+    ret.finally(() => {
+      releaseSnapshot?.();
+    });
+  } else {
+    releaseSnapshot?.();
+  }
   return (ret: Return);
 }
 

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

@@ -630,12 +630,66 @@ describe('Selector Cache', () => {
   });
 });
 
-describe('Snapshot cache', () => {
-  testRecoil('Snapshot is cached', () => {
-    const myAtom = atom({
-      key: 'useRecoilCallback snapshot cached',
-      default: 'DEFAULT',
+describe('Snapshot', () => {
+  testRecoil('Snapshot is retained for async callbacks', async ({gks}) => {
+    let callback,
+      callbackSnapshot,
+      resolveSelector,
+      resolveSelector2,
+      resolveCallback;
+
+    const myAtom = stringAtom();
+    const mySelector1 = selector({
+      key: 'useRecoilCallback snapshot retain 1',
+      get: async ({get}) => {
+        await new Promise(resolve => {
+          resolveSelector = resolve;
+        });
+        return get(myAtom);
+      },
     });
+    const mySelector2 = selector({
+      key: 'useRecoilCallback snapshot retain 2',
+      get: async ({get}) => {
+        await new Promise(resolve => {
+          resolveSelector2 = resolve;
+        });
+        return get(myAtom);
+      },
+    });
+
+    function Component() {
+      callback = useRecoilCallback(({snapshot}) => async () => {
+        callbackSnapshot = snapshot;
+        return new Promise(resolve => {
+          resolveCallback = resolve;
+        });
+      });
+    }
+
+    renderElements(<Component />);
+    callback?.();
+    const selector1Promise = callbackSnapshot?.getPromise(mySelector1);
+    const selector2Promise = callbackSnapshot?.getPromise(mySelector2);
+
+    // Wait to allow callback snapshot to auto-release after clock tick.
+    // It should still be retained for duration of callback, though.
+    await flushPromisesAndTimers();
+
+    // Selectors resolving before callback is resolved should not be canceled
+    act(() => resolveSelector());
+    await expect(selector1Promise).resolves.toBe('DEFAULT');
+
+    // Selectors resolving after callback is resolved should be canceled
+    if (gks.includes('recoil_memory_managament_2020')) {
+      act(() => resolveCallback());
+      act(() => resolveSelector2());
+      await expect(selector2Promise).rejects.toEqual({});
+    }
+  });
+
+  testRecoil('Snapshot is cached', () => {
+    const myAtom = stringAtom();
 
     let getSnapshot;
     let setMyAtom, resetMyAtom;
@@ -693,10 +747,7 @@ describe('Snapshot cache', () => {
   });
 
   testRecoil('cached snapshot is invalidated if not retained', async () => {
-    const myAtom = atom({
-      key: 'useRecoilCallback snapshot cache retained',
-      default: 'DEFAULT',
-    });
+    const myAtom = stringAtom();
 
     let getSnapshot;
     let setMyAtom;

packages/shared/__test_utils__/Recoil_TestingUtils.js

@@ -398,7 +398,13 @@ const testGKs =
   };
 
 const WWW_GKS_TO_TEST = QUICK_TEST
-  ? [['recoil_hamt_2020', 'recoil_sync_external_store']]
+  ? [
+      [
+        'recoil_hamt_2020',
+        'recoil_sync_external_store',
+        'recoil_memory_managament_2020',
+      ],
+    ]
   : [
       // OSS for React <18:
       ['recoil_hamt_2020', 'recoil_suppress_rerender_in_callback'], // Also enables early rendering