0.10.0 (3) switch rehydration mismatch protection to useSyncExternalStore (#207)

* add test for race condition

* switch to `useSyncExternalStore`

* add comments

* fix up tests

* move React functions out of test

* disable semgrep

* Update packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.test.tsx

Co-authored-by: Jerel Miller <[email protected]>

* remove some `await test` calls

---------

Co-authored-by: Jerel Miller <[email protected]>

Author: phryneas

packages/client-react-streaming/src/AccumulateMultipartResponsesLink.test.ts

@@ -14,7 +14,7 @@ import type { SubscriptionObserver } from "zen-observable-ts";
 const { DebounceMultipartResponsesLink: AccumulateMultipartResponsesLink } =
   await import("#bundled");
 
-await test("normal queries can resolve synchronously", () => {
+test("normal queries can resolve synchronously", () => {
   const query = gql`
     query {
       fastField
@@ -46,7 +46,7 @@ await test("normal queries can resolve synchronously", () => {
   });
 });
 
-await test("deferred query will complete synchonously if maxDelay is 0", () => {
+test("deferred query will complete synchonously if maxDelay is 0", () => {
   const query = gql`
     query {
       fastField
@@ -81,7 +81,7 @@ await test("deferred query will complete synchonously if maxDelay is 0", () => {
   });
 });
 
-await test("`next` call will be debounced and results will be merged together", () => {
+test("`next` call will be debounced and results will be merged together", () => {
   mock.timers.enable();
 
   const query = gql`

packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.test.tsx

@@ -1,4 +1,4 @@
-import React, { Suspense, useMemo } from "rehackt";
+import React, { Suspense, use, useMemo } from "rehackt";
 import { outsideOf } from "../util/runInConditions.js";
 import assert from "node:assert";
 import test, { afterEach, describe } from "node:test";
@@ -24,17 +24,21 @@ const {
   InMemoryCache,
   WrapApolloProvider,
   DataTransportContext,
+  resetApolloSingletons,
 } = await import("#bundled");
 
-await describe(
+describe(
   "tests with DOM access",
   { skip: outsideOf("node", "browser") },
   async () => {
     // @ts-expect-error seems to have a wrong type?
     await import("global-jsdom/register");
-    const { render, cleanup } = await import("@testing-library/react");
+    const { render, cleanup, getQueriesForElement } = await import(
+      "@testing-library/react"
+    );
 
     afterEach(cleanup);
+    afterEach(resetApolloSingletons);
 
     const QUERY_ME: TypedDocumentNode<{ me: string }> = gql`
       query {
@@ -215,9 +219,9 @@ await describe(
         await findByText("User");
 
         assert.ok(attemptedRenderCount > 0);
-        // one render to rehydrate the server value
+        // will try with server value and immediately restart with client value
         // one rerender with the actual client value (which is hopefull equal)
-        assert.equal(finishedRenderCount, 2);
+        assert.equal(finishedRenderCount, 1);
 
         assert.deepStrictEqual(JSON.parse(JSON.stringify(client.extract())), {
           ROOT_QUERY: {
@@ -227,10 +231,127 @@ await describe(
         });
       }
     );
+
+    test(
+      "race condition: client ahead of server renders without hydration mismatch",
+      { skip: outsideOf("browser") },
+      async () => {
+        const { $RC, $RS, setBody, hydrateBody, appendToBody } = await import(
+          "../util/hydrationTest.js"
+        );
+        // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-constraint
+        let useStaticValueRefStub = <T extends unknown>(): { current: T } => {
+          throw new Error("Should not be called yet!");
+        };
+
+        const client = new ApolloClient({
+          connectToDevTools: false,
+          cache: new InMemoryCache(),
+        });
+        const simulateRequestStart = client.onQueryStarted!;
+        const simulateRequestData = client.onQueryProgress!;
+
+        const Provider = WrapApolloProvider(({ children }) => {
+          return (
+            <DataTransportContext.Provider
+              value={useMemo(
+                () => ({
+                  useStaticValueRef() {
+                    return useStaticValueRefStub();
+                  },
+                }),
+                []
+              )}
+            >
+              {children}
+            </DataTransportContext.Provider>
+          );
+        });
+
+        const finishedRenders: any[] = [];
+
+        function Child() {
+          const { data } = useSuspenseQuery(QUERY_ME);
+          finishedRenders.push(data);
+          return <div id="user">{data.me}</div>;
+        }
+
+        const promise = Promise.resolve();
+        // suspends on the server, immediately resolved in browser
+        function ParallelSuspending() {
+          use(promise);
+          return <div id="parallel">suspending in parallel</div>;
+        }
+
+        const { findByText } = getQueriesForElement(document.body);
+
+        // server starts streaming
+        setBody`<!--$?--><template id="B:0"></template>Fallback<!--/$-->`;
+        // request started on the server
+        simulateRequestStart(EVENT_STARTED);
+
+        hydrateBody(
+          <Provider makeClient={() => client}>
+            <Suspense fallback={"Fallback"}>
+              <Child />
+              <ParallelSuspending />
+            </Suspense>
+          </Provider>
+        );
+
+        await findByText("Fallback");
+        // this is the div for the suspense boundary
+        appendToBody`<div hidden id="S:0"><template id="P:1"></template><template id="P:2"></template></div>`;
+        // request has finished on the server
+        simulateRequestData(EVENT_DATA);
+        simulateRequestData(EVENT_COMPLETE);
+        // `Child` component wants to transport data from SSR render to the browser
+        useStaticValueRefStub = () => ({ current: FIRST_HOOK_RESULT as any });
+        // `Child` finishes rendering on the server
+        appendToBody`<div hidden id="S:1"><div id="user">User</div></div>`;
+        $RS("S:1", "P:1");
+
+        // meanwhile, in the browser, the cache is modified
+        client.cache.writeQuery({
+          query: QUERY_ME,
+          data: {
+            me: "Future me.",
+          },
+        });
+
+        // `ParallelSuspending` finishes rendering
+        appendToBody`<div hidden id="S:2"><div id="parallel">suspending in parallel</div></div>`;
+        $RS("S:2", "P:2");
+
+        // everything in the suspense boundary finished rendering, so assemble HTML and take up React rendering again
+        $RC("B:0", "S:0");
+
+        // we expect the *new* value to appear after hydration finished, not the old value from the server
+        await findByText("Future me.");
+
+        // one render to rehydrate the server value
+        // one rerender with the actual client value (which is hopefull equal)
+        assert.deepStrictEqual(finishedRenders, [
+          { me: "User" },
+          { me: "Future me." },
+        ]);
+
+        assert.deepStrictEqual(JSON.parse(JSON.stringify(client.extract())), {
+          ROOT_QUERY: {
+            __typename: "Query",
+            me: "Future me.",
+          },
+        });
+        assert.equal(
+          document.body.innerHTML,
+          `<!--$--><div id="user">Future me.</div><div id="parallel">suspending in parallel</div><!--/$-->`
+        );
+      }
+    );
   }
 );
 
-await describe("document transforms are applied correctly", async () => {
+describe("document transforms are applied correctly", async () => {
   const untransformedQuery = gql`
     query Test {
       user {

packages/client-react-streaming/src/DataTransportAbstraction/useTransportValue.tsx

@@ -1,5 +1,5 @@
 "use client";
-import { useContext, useEffect, useState } from "react";
+import { useContext, useSyncExternalStore } from "react";
 import { DataTransportContext } from "./DataTransportAbstraction.js";
 
 /**
@@ -12,20 +12,24 @@ import { DataTransportContext } from "./DataTransportAbstraction.js";
  * the component can change to client-side values instead.
  */
 export function useTransportValue<T>(value: T): T {
-  const [isClient, setIsClient] = useState(false);
-  useEffect(() => setIsClient(true), []);
-
   const dataTransport = useContext(DataTransportContext);
   if (!dataTransport)
     throw new Error(
       "useTransportValue must be used within a streaming-specific ApolloProvider"
     );
   const valueRef = dataTransport.useStaticValueRef(value);
-  if (isClient) {
+
+  const retVal = useSyncExternalStore(
+    () => () => {},
+    () => value,
+    () => valueRef.current
+  );
+
+  if (retVal === value) {
     // @ts-expect-error this value will never be used again
     // so we can safely delete it
     valueRef.current = undefined;
   }
 
-  return isClient ? value : valueRef.current;
+  return retVal;
 }

packages/client-react-streaming/src/util/hydrationTest.ts

@@ -0,0 +1,33 @@
+import { hydrateRoot } from "react-dom/client";
+
+/* eslint-disable */
+// prettier-ignore
+/** React completeSegment function */
+// @ts-expect-error This is React code.
+export function $RS(a, b) { a = document.getElementById(a); b = document.getElementById(b); for (a.parentNode.removeChild(a); a.firstChild;)b.parentNode.insertBefore(a.firstChild, b); b.parentNode.removeChild(b) }
+// prettier-ignore
+/** React completeBoundary function */
+// @ts-expect-error This is React code.
+export function $RC(b, c, e = undefined) { c = document.getElementById(c); c.parentNode.removeChild(c); var a = document.getElementById(b); if (a) { b = a.previousSibling; if (e) b.data = "$!", a.setAttribute("data-dgst", e); else { e = b.parentNode; a = b.nextSibling; var f = 0; do { if (a && 8 === a.nodeType) { var d = a.data; if ("/$" === d) if (0 === f) break; else f--; else "$" !== d && "$?" !== d && "$!" !== d || f++ } d = a.nextSibling; e.removeChild(a); a = d } while (a); for (; c.firstChild;)e.insertBefore(c.firstChild, a); b.data = "$" } b._reactRetry && b._reactRetry() } }
+/* eslint-enable */
+
+export function hydrateBody(
+  initialChildren: Parameters<typeof hydrateRoot>[1],
+  options?: Parameters<typeof hydrateRoot>[2]
+) {
+  return hydrateRoot(document.body, initialChildren, options);
+}
+
+export function setBody(html: TemplateStringsArray) {
+  if (html.length !== 1)
+    throw new Error("Expected exactly one template string");
+  // nosemgrep
+  document.body.innerHTML = html[0];
+}
+
+export function appendToBody(html: TemplateStringsArray) {
+  if (html.length !== 1)
+    throw new Error("Expected exactly one template string");
+  // nosemgrep
+  document.body.insertAdjacentHTML("beforeend", html[0]);
+}