Fix suspense for client-only fragments that contains suspensible live…

… resolvers

Reviewed By: rbalicki2

Differential Revision: D39219336

fbshipit-source-id: 2d23cfc2b8141f901a0e48af38f31600927d8222

packages/react-relay/__tests__/LiveResolvers-test.js

@@ -1192,6 +1192,78 @@ describe.each([
 
     expect(renderer.toJSON()).toBe('Loading...');
   });
+
+  describe('client-only fragments', () => {
+    const LiveResolversTestCounterUserFragment = graphql`
+      fragment LiveResolversTestCounterUserFragment on User {
+        counter_suspends_when_odd
+      }
+    `;
+
+    const LiveResolversTestLiveResolverSuspenseQuery = graphql`
+      query LiveResolversTestLiveResolverSuspenseQuery($id: ID!) {
+        node(id: $id) {
+          ...LiveResolversTestCounterUserFragment
+        }
+      }
+    `;
+
+    function Environment({
+      children,
+      environment,
+    }: {
+      children: React.Node,
+      environment: RelayModernEnvironment,
+    }) {
+      return (
+        <RelayEnvironmentProvider environment={environment}>
+          <React.Suspense fallback="Loading...">{children}</React.Suspense>
+        </RelayEnvironmentProvider>
+      );
+    }
+
+    function TestComponent(props: {id: string}) {
+      const queryData = useLazyLoadQuery(
+        LiveResolversTestLiveResolverSuspenseQuery,
+        {id: props.id},
+      );
+      const fragmentData = useFragment(
+        LiveResolversTestCounterUserFragment,
+        queryData.node,
+      );
+      return fragmentData?.counter_suspends_when_odd;
+    }
+
+    test('correctly suspend on fragments with client-only data', () => {
+      const environment = new RelayModernEnvironment({
+        network: RelayNetwork.create(jest.fn()),
+        store: new LiveResolverStore(RelayRecordSource.create()),
+      });
+      environment.commitPayload(
+        createOperationDescriptor(LiveResolversTestLiveResolverSuspenseQuery, {
+          id: '1',
+        }),
+        {
+          node: {id: '1', __typename: 'User'},
+        },
+      );
+      const renderer = TestRenderer.create(
+        <Environment environment={environment}>
+          <TestComponent id="1" />
+        </Environment>,
+      );
+      expect(renderer.toJSON()).toEqual('0');
+      TestRenderer.act(() => {
+        GLOBAL_STORE.dispatch({type: 'INCREMENT'});
+      });
+      TestRenderer.act(() => jest.runAllImmediates());
+      expect(renderer.toJSON()).toEqual('Loading...');
+      TestRenderer.act(() => {
+        GLOBAL_STORE.dispatch({type: 'INCREMENT'});
+      });
+      expect(renderer.toJSON()).toEqual('2');
+    });
+  });
 });
 
 test('Errors when reading a @live resolver that does not return a LiveState object', () => {

packages/relay-runtime/store/RelayReader.js

@@ -421,10 +421,11 @@ class RelayReader {
           );
           // The only case where we want to suspend due to missing data off of
           // a client extension is if we reached a client edge that we might be
-          // able to fetch:
+          // able to fetch, or there is a missing data in one of the live resolvers.
           this._isMissingData =
             isMissingData ||
-            this._missingClientEdges.length > alreadyMissingClientEdges;
+            this._missingClientEdges.length > alreadyMissingClientEdges ||
+            this._missingLiveResolverFields.length > 0;
           if (RelayFeatureFlags.ENABLE_CLIENT_EDGES) {
             this._clientEdgeTraversalPath.pop();
           }

packages/relay-runtime/store/__tests__/resolvers/CounterSuspendsWhenOddOnUser.js

@@ -0,0 +1,48 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ * @flow strict-local
+ * @emails oncall+relay
+ */
+
+'use strict';
+
+import type {LiveState} from 'relay-runtime/store/experimental-live-resolvers/LiveResolverStore';
+
+const {GLOBAL_STORE, Selectors} = require('./ExampleExternalStateStore');
+const {
+  suspenseSentinel,
+} = require('relay-runtime/store/experimental-live-resolvers/LiveResolverSuspenseSentinel');
+
+/**
+ * @RelayResolver
+ * @fieldName counter_suspends_when_odd
+ * @onType User
+ * @live
+ *
+ * A Relay Resolver that returns an object implementing the External State
+ * Resolver interface.
+ */
+function CounterSuspendsWhenOddOnUser(): LiveState<number> {
+  return {
+    read() {
+      const number = Selectors.getNumber(GLOBAL_STORE.getState());
+      if (number % 2 !== 0) {
+        return suspenseSentinel();
+      } else {
+        return number;
+      }
+    },
+    subscribe(cb): () => void {
+      // Here we could try to run the selector and short-circuit if
+      // the value has not changed, but for now we'll over-notify.
+      return GLOBAL_STORE.subscribe(cb);
+    },
+  };
+}
+
+module.exports = CounterSuspendsWhenOddOnUser;