experimental_use(context) for server components and ssr (#25226)

implements the experimental use(context) API for the server components (Flight) and SSR (Fizz) runtimes

Author: mofeiZ

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

@@ -5292,6 +5292,54 @@ describe('ReactDOMFizzServer', () => {
       expect(getVisibleChildren(container)).toEqual('ABC');
     });
 
+    // @gate enableUseHook
+    it('basic use(context)', async () => {
+      const ContextA = React.createContext('default');
+      const ContextB = React.createContext('B');
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default',
+      );
+      function Client() {
+        return use(ContextA) + use(ContextB);
+      }
+      function ServerComponent() {
+        return use(ServerContext);
+      }
+      function Server() {
+        return (
+          <ServerContext.Provider value="C">
+            <ServerComponent />
+          </ServerContext.Provider>
+        );
+      }
+      function App() {
+        return (
+          <>
+            <ContextA.Provider value="A">
+              <Client />
+            </ContextA.Provider>
+            <Server />
+          </>
+        );
+      }
+
+      await act(async () => {
+        const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />);
+        pipe(writable);
+      });
+      expect(getVisibleChildren(container)).toEqual(['AB', 'C']);
+
+      // Hydration uses a different renderer runtime (Fiber instead of Fizz).
+      // We reset _currentRenderer here to not trigger a warning about multiple
+      // renderers concurrently using these contexts
+      ContextA._currentRenderer = null;
+      ServerContext._currentRenderer = null;
+      ReactDOMClient.hydrateRoot(container, <App />);
+      expect(Scheduler).toFlushAndYield([]);
+      expect(getVisibleChildren(container)).toEqual(['AB', 'C']);
+    });
+
     // @gate enableUseHook
     it('use(promise) in multiple components', async () => {
       const promiseA = Promise.resolve('A');

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

@@ -36,7 +36,10 @@ import {
   enableUseHook,
   enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
-import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
+import {
+  REACT_CONTEXT_TYPE,
+  REACT_SERVER_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
 
 import {NoMode, ConcurrentMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
@@ -771,8 +774,8 @@ function use<T>(usable: Usable<T>): T {
         }
       }
     } else if (
-      usable.$$typeof != null &&
-      usable.$$typeof === REACT_CONTEXT_TYPE
+      usable.$$typeof === REACT_CONTEXT_TYPE ||
+      usable.$$typeof === REACT_SERVER_CONTEXT_TYPE
     ) {
       const context: ReactContext<T> = (usable: any);
       return readContext(context);

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

@@ -36,7 +36,10 @@ import {
   enableUseHook,
   enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
-import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
+import {
+  REACT_CONTEXT_TYPE,
+  REACT_SERVER_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
 
 import {NoMode, ConcurrentMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
@@ -771,8 +774,8 @@ function use<T>(usable: Usable<T>): T {
         }
       }
     } else if (
-      usable.$$typeof != null &&
-      usable.$$typeof === REACT_CONTEXT_TYPE
+      usable.$$typeof === REACT_CONTEXT_TYPE ||
+      usable.$$typeof === REACT_SERVER_CONTEXT_TYPE
     ) {
       const context: ReactContext<T> = (usable: any);
       return readContext(context);

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js

@@ -594,6 +594,40 @@ describe('ReactFlightDOMBrowser', () => {
     expect(container.innerHTML).toBe('ABC');
   });
 
+  // @gate enableUseHook
+  it('basic use(context)', async () => {
+    const ContextA = React.createServerContext('ContextA', '');
+    const ContextB = React.createServerContext('ContextB', 'B');
+
+    function ServerComponent() {
+      return use(ContextA) + use(ContextB);
+    }
+    function Server() {
+      return (
+        <ContextA.Provider value="A">
+          <ServerComponent />
+        </ContextA.Provider>
+      );
+    }
+    const stream = ReactServerDOMWriter.renderToReadableStream(<Server />);
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    function Client() {
+      return response.readRoot();
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      // Client uses a different renderer.
+      // We reset _currentRenderer here to not trigger a warning about multiple
+      // renderers concurrently using this context
+      ContextA._currentRenderer = null;
+      root.render(<Client />);
+    });
+    expect(container.innerHTML).toBe('AB');
+  });
+
   // @gate enableUseHook
   it('use(promise) in multiple components', async () => {
     function Child({prefix}) {

packages/react-server/src/ReactFizzHooks.js

@@ -39,6 +39,10 @@ import {
   enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
 import is from 'shared/objectIs';
+import {
+  REACT_SERVER_CONTEXT_TYPE,
+  REACT_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
 
 type BasicStateAction<S> = (S => S) | S;
 type Dispatch<A> = A => void;
@@ -616,8 +620,12 @@ function use<T>(usable: Usable<T>): T {
           }
         }
       }
-    } else {
-      // TODO: Add support for Context
+    } else if (
+      usable.$$typeof === REACT_CONTEXT_TYPE ||
+      usable.$$typeof === REACT_SERVER_CONTEXT_TYPE
+    ) {
+      const context: ReactContext<T> = (usable: any);
+      return readContext(context);
     }
   }
 

packages/react-server/src/ReactFlightHooks.js

@@ -200,8 +200,9 @@ function use<T>(usable: Usable<T>): T {
           }
         }
       }
-    } else {
-      // TODO: Add support for Context
+    } else if (usable.$$typeof === REACT_SERVER_CONTEXT_TYPE) {
+      const context: ReactServerContext<T> = (usable: any);
+      return readContext(context);
     }
   }
 

packages/shared/ReactTypes.js

@@ -213,5 +213,4 @@ export type StartTransitionOptions = {
   name?: string,
 };
 
-// TODO: Add Context support
 export type Usable<T> = Thenable<T> | ReactContext<T>;