Implement useSyncExternalStore on server (facebook#22347)

Adds a third argument called `getServerSnapshot`.

On the server, React calls this one instead of the normal `getSnapshot`.
We also call it during hydration.

So it represents the snapshot that is used to generate the initial,
server-rendered HTML. The purpose is to avoid server-client mismatches.
What we render during hydration needs to match up exactly with what we
render on the server.

The pattern is for the server to send down a serialized copy of the
store that was used to generate the initial HTML. On the client, React
will call either `getSnapshot` or `getServerSnapshot` on the client as
appropriate, depending on whether it's currently hydrating.

The argument is optional for fully client rendered use cases. If the
user does attempt to omit `getServerSnapshot`, and the hook is called
on the server, React will abort that subtree on the server and
revert to client rendering, up to the nearest Suspense boundary.

For the userspace shim, we will need to use a heuristic (canUseDOM)
to determine whether we are in a server environment. I'll do that in
a follow up.

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -258,6 +258,7 @@ function useMemo<T>(
 function useSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
+  getServerSnapshot?: () => T,
 ): T {
   // useSyncExternalStore() composes multiple hooks internally.
   // Advance the current hook index the same number of times

packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js

@@ -17,6 +17,8 @@ let ReactDOM;
 let ReactDOMFizzServer;
 let Suspense;
 let SuspenseList;
+let useSyncExternalStore;
+let useSyncExternalStoreExtra;
 let PropTypes;
 let textCache;
 let document;
@@ -39,6 +41,9 @@ describe('ReactDOMFizzServer', () => {
     Stream = require('stream');
     Suspense = React.Suspense;
     SuspenseList = React.SuspenseList;
+    useSyncExternalStore = React.unstable_useSyncExternalStore;
+    useSyncExternalStoreExtra = require('use-sync-external-store/extra')
+      .useSyncExternalStoreExtra;
     PropTypes = require('prop-types');
 
     textCache = new Map();
@@ -1478,4 +1483,156 @@ describe('ReactDOMFizzServer', () => {
     // We should've been able to display the content without waiting for the rest of the fallback.
     expect(getVisibleChildren(container)).toEqual(<div>Hello</div>);
   });
+
+  // @gate supportsNativeUseSyncExternalStore
+  // @gate experimental
+  it('calls getServerSnapshot instead of getSnapshot', async () => {
+    const ref = React.createRef();
+
+    function getServerSnapshot() {
+      return 'server';
+    }
+
+    function getClientSnapshot() {
+      return 'client';
+    }
+
+    function subscribe() {
+      return () => {};
+    }
+
+    function Child({text}) {
+      Scheduler.unstable_yieldValue(text);
+      return text;
+    }
+
+    function App() {
+      const value = useSyncExternalStore(
+        subscribe,
+        getClientSnapshot,
+        getServerSnapshot,
+      );
+      return (
+        <div ref={ref}>
+          <Child text={value} />
+        </div>
+      );
+    }
+
+    const loggedErrors = [];
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <Suspense fallback="Loading...">
+          <App />
+        </Suspense>,
+        writable,
+        {
+          onError(x) {
+            loggedErrors.push(x);
+          },
+        },
+      );
+      startWriting();
+    });
+    expect(Scheduler).toHaveYielded(['server']);
+
+    const serverRenderedDiv = container.getElementsByTagName('div')[0];
+
+    ReactDOM.hydrateRoot(container, <App />);
+
+    // The first paint uses the server snapshot
+    expect(Scheduler).toFlushUntilNextPaint(['server']);
+    expect(getVisibleChildren(container)).toEqual(<div>server</div>);
+    // Hydration succeeded
+    expect(ref.current).toEqual(serverRenderedDiv);
+
+    // Asynchronously we detect that the store has changed on the client,
+    // and patch up the inconsistency
+    expect(Scheduler).toFlushUntilNextPaint(['client']);
+    expect(getVisibleChildren(container)).toEqual(<div>client</div>);
+    expect(ref.current).toEqual(serverRenderedDiv);
+  });
+
+  // The selector implementation uses the lazy ref initialization pattern
+  // @gate !(enableUseRefAccessWarning && __DEV__)
+  // @gate supportsNativeUseSyncExternalStore
+  // @gate experimental
+  it('calls getServerSnapshot instead of getSnapshot (with selector and isEqual)', async () => {
+    // Same as previous test, but with a selector that returns a complex object
+    // that is memoized with a custom `isEqual` function.
+    const ref = React.createRef();
+
+    function getServerSnapshot() {
+      return {env: 'server', other: 'unrelated'};
+    }
+
+    function getClientSnapshot() {
+      return {env: 'client', other: 'unrelated'};
+    }
+
+    function selector({env}) {
+      return {env};
+    }
+
+    function isEqual(a, b) {
+      return a.env === b.env;
+    }
+
+    function subscribe() {
+      return () => {};
+    }
+
+    function Child({text}) {
+      Scheduler.unstable_yieldValue(text);
+      return text;
+    }
+
+    function App() {
+      const {env} = useSyncExternalStoreExtra(
+        subscribe,
+        getClientSnapshot,
+        getServerSnapshot,
+        selector,
+        isEqual,
+      );
+      return (
+        <div ref={ref}>
+          <Child text={env} />
+        </div>
+      );
+    }
+
+    const loggedErrors = [];
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <Suspense fallback="Loading...">
+          <App />
+        </Suspense>,
+        writable,
+        {
+          onError(x) {
+            loggedErrors.push(x);
+          },
+        },
+      );
+      startWriting();
+    });
+    expect(Scheduler).toHaveYielded(['server']);
+
+    const serverRenderedDiv = container.getElementsByTagName('div')[0];
+
+    ReactDOM.hydrateRoot(container, <App />);
+
+    // The first paint uses the server snapshot
+    expect(Scheduler).toFlushUntilNextPaint(['server']);
+    expect(getVisibleChildren(container)).toEqual(<div>server</div>);
+    // Hydration succeeded
+    expect(ref.current).toEqual(serverRenderedDiv);
+
+    // Asynchronously we detect that the store has changed on the client,
+    // and patch up the inconsistency
+    expect(Scheduler).toFlushUntilNextPaint(['client']);
+    expect(getVisibleChildren(container)).toEqual(<div>client</div>);
+    expect(ref.current).toEqual(serverRenderedDiv);
+  });
 });

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -464,8 +464,16 @@ export function useCallback<T>(
 function useSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
+  getServerSnapshot?: () => T,
 ): T {
-  throw new Error('Not yet implemented');
+  if (getServerSnapshot === undefined) {
+    invariant(
+      false,
+      'Missing getServerSnapshot, which is required for ' +
+        'server-rendered content. Will revert to client rendering.',
+    );
+  }
+  return getServerSnapshot();
 }
 
 function useDeferredValue<T>(value: T): T {

packages/react-reconciler/src/ReactFiberHooks.new.js

@@ -938,23 +938,64 @@ function rerenderReducer<S, I, A>(
 function mountSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
+  getServerSnapshot?: () => T,
 ): T {
   const fiber = currentlyRenderingFiber;
   const hook = mountWorkInProgressHook();
-  // Read the current snapshot from the store on every render. This breaks the
-  // normal rules of React, and only works because store updates are
-  // always synchronous.
-  const nextSnapshot = getSnapshot();
-  if (__DEV__) {
-    if (!didWarnUncachedGetSnapshot) {
-      if (nextSnapshot !== getSnapshot()) {
-        console.error(
-          'The result of getSnapshot should be cached to avoid an infinite loop',
-        );
-        didWarnUncachedGetSnapshot = true;
+
+  let nextSnapshot;
+  const isHydrating = getIsHydrating();
+  if (isHydrating) {
+    if (getServerSnapshot === undefined) {
+      invariant(
+        false,
+        'Missing getServerSnapshot, which is required for ' +
+          'server-rendered content. Will revert to client rendering.',
+      );
+    }
+    nextSnapshot = getServerSnapshot();
+    if (__DEV__) {
+      if (!didWarnUncachedGetSnapshot) {
+        if (nextSnapshot !== getServerSnapshot()) {
+          console.error(
+            'The result of getServerSnapshot should be cached to avoid an infinite loop',
+          );
+          didWarnUncachedGetSnapshot = true;
+        }
       }
     }
+  } else {
+    nextSnapshot = getSnapshot();
+    if (__DEV__) {
+      if (!didWarnUncachedGetSnapshot) {
+        if (nextSnapshot !== getSnapshot()) {
+          console.error(
+            'The result of getSnapshot should be cached to avoid an infinite loop',
+          );
+          didWarnUncachedGetSnapshot = true;
+        }
+      }
+    }
+    // Unless we're rendering a blocking lane, schedule a consistency check.
+    // Right before committing, we will walk the tree and check if any of the
+    // stores were mutated.
+    //
+    // We won't do this if we're hydrating server-rendered content, because if
+    // the content is stale, it's already visible anyway. Instead we'll patch
+    // it up in a passive effect.
+    const root: FiberRoot | null = getWorkInProgressRoot();
+    invariant(
+      root !== null,
+      'Expected a work-in-progress root. This is a bug in React. Please file an issue.',
+    );
+    if (!includesBlockingLane(root, renderLanes)) {
+      pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
+    }
   }
+
+  // Read the current snapshot from the store on every render. This breaks the
+  // normal rules of React, and only works because store updates are
+  // always synchronous.
   hook.memoizedState = nextSnapshot;
   const inst: StoreInstance<T> = {
     value: nextSnapshot,
@@ -980,24 +1021,13 @@ function mountSyncExternalStore<T>(
     null,
   );
 
-  // Unless we're rendering a blocking lane, schedule a consistency check. Right
-  // before committing, we will walk the tree and check if any of the stores
-  // were mutated.
-  const root: FiberRoot | null = getWorkInProgressRoot();
-  invariant(
-    root !== null,
-    'Expected a work-in-progress root. This is a bug in React. Please file an issue.',
-  );
-  if (!includesBlockingLane(root, renderLanes)) {
-    pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
-  }
-
   return nextSnapshot;
 }
 
 function updateSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
+  getServerSnapshot?: () => T,
 ): T {
   const fiber = currentlyRenderingFiber;
   const hook = updateWorkInProgressHook();
@@ -2235,10 +2265,11 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       mountHookTypesDev();
-      return mountSyncExternalStore(subscribe, getSnapshot);
+      return mountSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -2366,10 +2397,11 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       updateHookTypesDev();
-      return mountSyncExternalStore(subscribe, getSnapshot);
+      return mountSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -2497,10 +2529,11 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot);
+      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -2629,10 +2662,11 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot);
+      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -2774,11 +2808,12 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       warnInvalidHookAccess();
       mountHookTypesDev();
-      return mountSyncExternalStore(subscribe, getSnapshot);
+      return mountSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -2921,11 +2956,12 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       warnInvalidHookAccess();
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot);
+      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';
@@ -3069,11 +3105,12 @@ if (__DEV__) {
     useSyncExternalStore<T>(
       subscribe: (() => void) => () => void,
       getSnapshot: () => T,
+      getServerSnapshot?: () => T,
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       warnInvalidHookAccess();
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot);
+      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
     },
     useOpaqueIdentifier(): OpaqueIDType | void {
       currentHookNameInDev = 'useOpaqueIdentifier';