[ServerContext] Flight support for ServerContext

packages/react-client/src/ReactFlightClient.js

@@ -24,7 +24,13 @@ import {
   parseModel,
 } from './ReactFlightClientHostConfig';
 
-import {REACT_LAZY_TYPE, REACT_ELEMENT_TYPE} from 'shared/ReactSymbols';
+import {getOrCreateServerContext} from 'react/src/ReactServerContext';
+
+import {
+  REACT_LAZY_TYPE,
+  REACT_ELEMENT_TYPE,
+  REACT_PROVIDER_TYPE,
+} from 'shared/ReactSymbols';
 
 export type JSONValue =
   | number
@@ -318,6 +324,11 @@ export function parseModelString(
       // When passed into React, we'll know how to suspend on this.
       return createLazyChunkWrapper(chunk);
     }
+    case '!': {
+      if (value === '!') {
+        return REACT_PROVIDER_TYPE;
+      }
+    }
   }
   return value;
 }
@@ -327,10 +338,18 @@ export function parseModelTuple(
   value: {+[key: string]: JSONValue} | $ReadOnlyArray<JSONValue>,
 ): any {
   const tuple: [mixed, mixed, mixed, mixed] = (value: any);
-  if (tuple[0] === REACT_ELEMENT_TYPE) {
-    // TODO: Consider having React just directly accept these arrays as elements.
-    // Or even change the ReactElement type to be an array.
-    return createElement(tuple[1], tuple[2], tuple[3]);
+
+  switch (tuple[0]) {
+    case REACT_ELEMENT_TYPE:
+      // TODO: Consider having React just directly accept these arrays as elements.
+      // Or even change the ReactElement type to be an array.
+      return createElement(tuple[1], tuple[2], tuple[3]);
+    case REACT_PROVIDER_TYPE:
+      return createElement(
+        getOrCreateServerContext((tuple[1]: any)).Provider,
+        tuple[2],
+        tuple[3],
+      );
   }
   return value;
 }

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -17,6 +17,7 @@ let ReactNoopFlightServer;
 let ReactNoopFlightClient;
 let ErrorBoundary;
 let NoErrorExpected;
+let Scheduler;
 
 describe('ReactFlight', () => {
   beforeEach(() => {
@@ -27,6 +28,7 @@ describe('ReactFlight', () => {
     ReactNoopFlightServer = require('react-noop-renderer/flight-server');
     ReactNoopFlightClient = require('react-noop-renderer/flight-client');
     act = require('jest-react').act;
+    Scheduler = require('scheduler');
 
     ErrorBoundary = class extends React.Component {
       state = {hasError: false, error: null};
@@ -302,4 +304,233 @@ describe('ReactFlight', () => {
       {withoutStack: true},
     );
   });
+
+  describe('ServerContext', () => {
+    // @gate enableServerContext
+    it('supports basic createServerContext usage', () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'hello from server',
+      );
+      function Foo() {
+        const context = React.useServerContext(ServerContext);
+        return <div>{context}</div>;
+      }
+
+      const transport = ReactNoopFlightServer.render(<Foo />);
+      act(() => {
+        ServerContext._currentRenderer = null;
+        ServerContext._currentRenderer2 = null;
+        ReactNoop.render(ReactNoopFlightClient.read(transport));
+      });
+
+      expect(ReactNoop).toMatchRenderedOutput(<div>hello from server</div>);
+    });
+
+    // @gate enableServerContext
+    it('propagates ServerContext providers in flight', () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default hello from server',
+      );
+
+      function Foo() {
+        return (
+          <div>
+            <ServerContext.Provider value="hi this is server">
+              <Bar />
+            </ServerContext.Provider>
+          </div>
+        );
+      }
+      function Bar() {
+        const context = React.useServerContext(ServerContext);
+        return context;
+      }
+
+      const transport = ReactNoopFlightServer.render(<Foo />);
+      act(() => {
+        ServerContext._currentRenderer = null;
+        ServerContext._currentRenderer2 = null;
+        ReactNoop.render(ReactNoopFlightClient.read(transport));
+      });
+
+      expect(ReactNoop).toMatchRenderedOutput(<div>hi this is server</div>);
+    });
+
+    // @gate enableServerContext
+    it('propagates ServerContext and cleansup providers in flight', () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default hello from server',
+      );
+
+      function Foo() {
+        return (
+          <>
+            <ServerContext.Provider value="hi this is server outer">
+              <ServerContext.Provider value="hi this is server">
+                <Bar />
+              </ServerContext.Provider>
+              <ServerContext.Provider value="hi this is server2">
+                <Bar />
+              </ServerContext.Provider>
+              <Bar />
+            </ServerContext.Provider>
+            <ServerContext.Provider value="hi this is server outer2">
+              <Bar />
+            </ServerContext.Provider>
+            <Bar />
+          </>
+        );
+      }
+      function Bar() {
+        const context = React.useServerContext(ServerContext);
+        return <span>{context}</span>;
+      }
+
+      const transport = ReactNoopFlightServer.render(<Foo />);
+      act(() => {
+        ServerContext._currentRenderer = null;
+        ServerContext._currentRenderer2 = null;
+        ReactNoop.render(ReactNoopFlightClient.read(transport));
+      });
+
+      expect(ReactNoop).toMatchRenderedOutput(
+        <>
+          <span>hi this is server</span>
+          <span>hi this is server2</span>
+          <span>hi this is server outer</span>
+          <span>hi this is server outer2</span>
+          <span>default hello from server</span>
+        </>,
+      );
+    });
+
+    // @gate enableServerContext
+    it('propagates ServerContext providers in flight after suspending', async () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default hello from server',
+      );
+
+      function Foo() {
+        return (
+          <div>
+            <ServerContext.Provider value="hi this is server">
+              <React.Suspense fallback={'Loading'}>
+                <Bar />
+              </React.Suspense>
+            </ServerContext.Provider>
+          </div>
+        );
+      }
+
+      let resolve;
+      const promise = new Promise(res => {
+        resolve = () => {
+          promise.unsuspend = true;
+          res();
+        };
+      });
+
+      function Bar() {
+        if (!promise.unsuspend) {
+          Scheduler.unstable_yieldValue('suspended');
+          throw promise;
+        }
+        Scheduler.unstable_yieldValue('rendered');
+        const context = React.useServerContext(ServerContext);
+        return context;
+      }
+
+      const transport = ReactNoopFlightServer.render(<Foo />);
+
+      expect(Scheduler).toHaveYielded(['suspended']);
+
+      await act(async () => {
+        resolve();
+        await promise;
+        jest.runAllImmediates();
+      });
+
+      expect(Scheduler).toHaveYielded(['rendered']);
+
+      act(() => {
+        ServerContext._currentRenderer = null;
+        ServerContext._currentRenderer2 = null;
+        ReactNoop.render(ReactNoopFlightClient.read(transport));
+      });
+
+      expect(ReactNoop).toMatchRenderedOutput(<div>hi this is server</div>);
+    });
+
+    // @gate enableServerContext
+    it('serializes ServerContext to client', async () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default hello from server',
+      );
+
+      function ClientBar() {
+        Scheduler.unstable_yieldValue('ClientBar');
+        const context = React.useServerContext(ServerContext);
+        return <span>{context}</span>;
+      }
+
+      const Bar = moduleReference(ClientBar);
+
+      function Foo() {
+        return (
+          <ServerContext.Provider value="hi this is server">
+            <Bar />
+          </ServerContext.Provider>
+        );
+      }
+
+      const model = {
+        foo: <Foo />,
+      };
+
+      const transport = ReactNoopFlightServer.render(model);
+
+      expect(Scheduler).toHaveYielded([]);
+
+      act(() => {
+        ServerContext._currentRenderer = null;
+        ServerContext._currentRenderer2 = null;
+        const flightModel = ReactNoopFlightClient.read(transport);
+        ReactNoop.render(flightModel.foo);
+      });
+
+      expect(Scheduler).toHaveYielded(['ClientBar']);
+      expect(ReactNoop).toMatchRenderedOutput(<span>hi this is server</span>);
+    });
+
+    // @gate enableServerContext
+    it('takes ServerContext from client for refetching usecases', async () => {
+      const ServerContext = React.createServerContext(
+        'ServerContext',
+        'default hello from server',
+      );
+      function Bar() {
+        return <span>{React.useServerContext(ServerContext)}</span>;
+      }
+      const transport = ReactNoopFlightServer.render(<Bar />, {
+        context: [
+          {
+            name: 'ServerContext',
+            value: 'Override',
+          },
+        ],
+      });
+
+      act(() => {
+        const flightModel = ReactNoopFlightClient.read(transport);
+        ReactNoop.render(flightModel);
+      });
+
+      expect(ReactNoop).toMatchRenderedOutput(<span>Override</span>);
+    });
+  });
 });

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

@@ -12,7 +12,9 @@ import type {
   MutableSourceGetSnapshotFn,
   MutableSourceSubscribeFn,
   ReactContext,
+  ReactServerContext,
   ReactProviderType,
+  ServerContextJSONValue,
 } from 'shared/ReactTypes';
 import type {
   Fiber,
@@ -103,7 +105,9 @@ function getCacheForType<T>(resourceType: () => T): T {
   throw new Error('Not implemented.');
 }
 
-function readContext<T>(context: ReactContext<T>): T {
+function readContext<T: any>(
+  context: ReactContext<T> | ReactServerContext<T>,
+): T {
   // For now we don't expose readContext usage in the hooks debugging info.
   return context._currentValue;
 }
@@ -117,6 +121,17 @@ function useContext<T>(context: ReactContext<T>): T {
   return context._currentValue;
 }
 
+function useServerContext<T: ServerContextJSONValue>(
+  context: ReactServerContext<T>,
+): T {
+  hookLog.push({
+    primitive: 'ServerContext',
+    stackError: new Error(),
+    value: context._currentValue,
+  });
+  return context._currentValue;
+}
+
 function useState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
@@ -331,7 +346,8 @@ function useId(): string {
 
 const Dispatcher: DispatcherType = {
   getCacheForType,
-  readContext,
+  // TODO: figure out why flow is complaining here
+  readContext: (readContext: any),
   useCacheRefresh,
   useCallback,
   useContext,
@@ -343,6 +359,7 @@ const Dispatcher: DispatcherType = {
   useMemo,
   useReducer,
   useRef,
+  useServerContext,
   useState,
   useTransition,
   useMutableSource,
@@ -673,12 +690,16 @@ export function inspectHooks<Props>(
   return buildTree(rootStack, readHookLog, includeHooksSource);
 }
 
-function setupContexts(contextMap: Map<ReactContext<any>, any>, fiber: Fiber) {
+function setupContexts(
+  contextMap: Map<ReactContext<any> | ReactServerContext<any>, any>,
+  fiber: Fiber,
+) {
   let current = fiber;
   while (current) {
     if (current.tag === ContextProvider) {
       const providerType: ReactProviderType<any> = current.type;
-      const context: ReactContext<any> = providerType._context;
+      const context: ReactContext<any> | ReactServerContext<any> =
+        providerType._context;
       if (!contextMap.has(context)) {
         // Store the current value that we're going to restore later.
         contextMap.set(context, context._currentValue);
@@ -690,7 +711,9 @@ function setupContexts(contextMap: Map<ReactContext<any>, any>, fiber: Fiber) {
   }
 }
 
-function restoreContexts(contextMap: Map<ReactContext<any>, any>) {
+function restoreContexts(
+  contextMap: Map<ReactContext<any> | ReactServerContext<any>, any>,
+) {
   contextMap.forEach((value, context) => (context._currentValue = value));
 }