From 3873809c2df83990af4adc9e25136a44fc6f75f1 Mon Sep 17 00:00:00 2001
From: Andrey Lunyov <alunyov@meta.com>
Date: Thu, 4 Jan 2024 12:22:59 -0800
Subject: [PATCH] Handle missed updates from the low-priority state changes in
 the new hooks implementation

Reviewed By: josephsavona

Differential Revision: D52541289

fbshipit-source-id: 6e2d35814e4b9138cf9c9027e6314fe2cc483109
---
 .../__tests__/useFragmentNode-test.js         | 52 +++++++++++++++++++
 .../useFragmentInternal_EXPERIMENTAL.js       | 28 +++++++++-
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js b/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
index ef65e9a1fecac..961162f00d0e8 100644
--- a/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
+++ b/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
@@ -1118,6 +1118,58 @@ describe.each([
       });
     });
 
+    it('should return the latest data when the hi-priority update happens at the same time as the low-priority store update', () => {
+      const startTransition = React.startTransition;
+      if (startTransition != null) {
+        internalAct(() => {
+          renderSingularFragment({
+            isConcurrent: true,
+          });
+        });
+        assertFragmentResults([
+          {
+            data: {
+              id: '1',
+              name: 'Alice',
+              profile_picture: null,
+              ...createFragmentRef('1', singularQuery),
+            },
+          },
+        ]);
+
+        internalAct(() => {
+          // Trigger store update with the lower priority
+          startTransition(() => {
+            environment.commitUpdate(store => {
+              store.get('1')?.setValue('Alice Updated Name', 'name');
+            });
+          });
+          // Trigger a hi-pri update with the higher priority, that should force component to re-render
+          forceSingularUpdate();
+        });
+
+        // Assert that the component re-renders twice, both times with the latest data
+        assertFragmentResults([
+          {
+            data: {
+              id: '1',
+              name: 'Alice Updated Name',
+              profile_picture: null,
+              ...createFragmentRef('1', singularQuery),
+            },
+          },
+          {
+            data: {
+              id: '1',
+              name: 'Alice Updated Name',
+              profile_picture: null,
+              ...createFragmentRef('1', singularQuery),
+            },
+          },
+        ]);
+      }
+    });
+
     it('should re-read and resubscribe to fragment when variables change', () => {
       renderSingularFragment();
       assertFragmentResults([
diff --git a/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js b/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
index 9318dcab05c93..c8f034420e577 100644
--- a/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
+++ b/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
@@ -244,6 +244,7 @@ function subscribeToSnapshot(
   environment: IEnvironment,
   state: FragmentState,
   setState: StateUpdaterFunction<FragmentState>,
+  hasPendingStateChanges: {current: boolean},
 ): () => void {
   if (state.kind === 'bailout') {
     return () => {};
@@ -264,11 +265,14 @@ function subscribeToSnapshot(
               name: 'useFragment.subscription.missedUpdates',
               hasDataChanges: dataChanged,
             });
+            hasPendingStateChanges.current = dataChanged;
             return dataChanged ? nextState : prevState;
           } else {
             return prevState;
           }
         }
+
+        hasPendingStateChanges.current = true;
         return {
           kind: 'singular',
           snapshot: latestSnapshot,
@@ -297,6 +301,8 @@ function subscribeToSnapshot(
                 name: 'useFragment.subscription.missedUpdates',
                 hasDataChanges: dataChanged,
               });
+              hasPendingStateChanges.current =
+                hasPendingStateChanges.current || dataChanged;
               return dataChanged ? nextState : prevState;
             } else {
               return prevState;
@@ -304,6 +310,7 @@ function subscribeToSnapshot(
           }
           const updated = [...prevState.snapshots];
           updated[index] = latestSnapshot;
+          hasPendingStateChanges.current = true;
           return {
             kind: 'plural',
             snapshots: updated,
@@ -529,6 +536,8 @@ function useFragmentInternal_EXPERIMENTAL(
   // they're missing even though we are out of options for possibly fetching them:
   handlePotentialSnapshotErrorsForState(environment, state);
 
+  const hasPendingStateChanges = useRef<boolean>(false);
+
   useEffect(() => {
     // Check for updates since the state was rendered
     let currentState = subscribedState;
@@ -547,9 +556,26 @@ function useFragmentInternal_EXPERIMENTAL(
       }
       currentState = updatedState;
     }
-    return subscribeToSnapshot(environment, currentState, setState);
+    return subscribeToSnapshot(
+      environment,
+      currentState,
+      setState,
+      hasPendingStateChanges,
+    );
   }, [environment, subscribedState]);
 
+  if (hasPendingStateChanges.current) {
+    const updates = handleMissedUpdates(environment, state);
+    if (updates != null) {
+      const [hasStateUpdates, updatedState] = updates;
+      if (hasStateUpdates) {
+        setState(updatedState);
+        state = updatedState;
+      }
+    }
+    hasPendingStateChanges.current = false;
+  }
+
   let data: ?SelectorData | Array<?SelectorData>;
   if (isPlural) {
     // Plural fragments require allocating an array of the snapshot data values,