From 64b707d9182de777fde51c9e326ccf0f52663867 Mon Sep 17 00:00:00 2001
From: Yuanchao Zhu <yczhu@meta.com>
Date: Tue, 18 Jul 2023 13:11:00 -0700
Subject: [PATCH] Set the cache as missing state when missed updates

Differential Revision: D47029820

fbshipit-source-id: a6e014f96b4398440a519a4392e6a3eafdff7bce
---
 .../relay-hooks/FragmentResource.js           | 48 +++++++++++++++----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/packages/react-relay/relay-hooks/FragmentResource.js b/packages/react-relay/relay-hooks/FragmentResource.js
index 5b21a436d5af1..58a3e6091ddad 100644
--- a/packages/react-relay/relay-hooks/FragmentResource.js
+++ b/packages/react-relay/relay-hooks/FragmentResource.js
@@ -740,10 +740,26 @@ class FragmentResourceImpl {
       // Only update the cache when the data is changed to avoid
       // returning different `data` instances
       if (didMissUpdates) {
-        this._cache.set(cacheKey, {
-          kind: 'done',
-          result: getFragmentResult(cacheKey, currentSnapshots, storeEpoch),
-        });
+        const result = getFragmentResult(
+          cacheKey,
+          currentSnapshots,
+          storeEpoch,
+        );
+        if (
+          RelayFeatureFlags.ENABLE_RELAY_OPERATION_TRACKER_SUSPENSE &&
+          result.isMissingData
+        ) {
+          this._cache.set(cacheKey, {
+            kind: 'missing',
+            result,
+            snapshot: currentSnapshots,
+          });
+        } else {
+          this._cache.set(cacheKey, {
+            kind: 'done',
+            result,
+          });
+        }
       }
       return [didMissUpdates, currentSnapshots];
     }
@@ -762,10 +778,26 @@ class FragmentResourceImpl {
       relayResolverErrors: currentSnapshot.relayResolverErrors,
     };
     if (updatedData !== renderData) {
-      this._cache.set(cacheKey, {
-        kind: 'done',
-        result: getFragmentResult(cacheKey, updatedCurrentSnapshot, storeEpoch),
-      });
+      const result = getFragmentResult(
+        cacheKey,
+        updatedCurrentSnapshot,
+        storeEpoch,
+      );
+      if (
+        RelayFeatureFlags.ENABLE_RELAY_OPERATION_TRACKER_SUSPENSE &&
+        result.isMissingData
+      ) {
+        this._cache.set(cacheKey, {
+          kind: 'missing',
+          result: result,
+          snapshot: updatedCurrentSnapshot,
+        });
+      } else {
+        this._cache.set(cacheKey, {
+          kind: 'done',
+          result,
+        });
+      }
     }
     return [updatedData !== renderData, updatedCurrentSnapshot];
   }