[ServerContext] Flight support for ServerContext

packages/react-client/src/ReactFlightClient.js

@@ -24,7 +24,7 @@ import {
   parseModel,
 } from './ReactFlightClientHostConfig';
 
-import {getOrCreateServerContext} from 'react/src/ReactServerContext';
+import {getOrCreateServerContext} from './ReactServerContext';
 
 import {
   REACT_LAZY_TYPE,

packages/react-client/src/ReactFlightClientStream.js

@@ -46,6 +46,7 @@ function processFullRow(response: Response, row: string): void {
       return;
     }
     case 'M': {
+      debugger;
       resolveModule(response, id, text);
       return;
     }

packages/react-client/src/ReactServerContext.js

@@ -0,0 +1,99 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import {
+  REACT_PROVIDER_TYPE,
+  REACT_SERVER_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
+
+import type {
+  ReactServerContext,
+  ServerContextJSONValue,
+} from 'shared/ReactTypes';
+
+import ReactSharedInternals from 'shared/ReactSharedInternals';
+
+import {enableServerContext} from 'shared/ReactFeatureFlags';
+
+const globalServerContextRegistry =
+  ReactSharedInternals.globalServerContextRegistry;
+
+export function createServerContext<T: ServerContextJSONValue>(
+  globalName: string,
+  defaultValue: T,
+): ReactServerContext<T> {
+  if (!enableServerContext) {
+    throw new Error('Not implemented.');
+  }
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      defaultValue,
+    );
+  }
+  const context = globalServerContextRegistry[globalName];
+  if (!context._definitionLoaded) {
+    context._currentValue = defaultValue;
+    context._currentValue2 = defaultValue;
+    context._definitionLoaded = true;
+  } else {
+    throw new Error(`ServerContext: ${globalName} already defined`);
+  }
+  return context;
+}
+
+function _createServerContext<T: ServerContextJSONValue>(
+  globalName: string,
+  defaultValue?: T,
+): ReactServerContext<T> {
+  const context: ReactServerContext<T> = {
+    $$typeof: REACT_SERVER_CONTEXT_TYPE,
+    // As a workaround to support multiple concurrent renderers, we categorize
+    // some renderers as primary and others as secondary. We only expect
+    // there to be two concurrent renderers at most: React Native (primary) and
+    // Fabric (secondary); React DOM (primary) and React ART (secondary).
+    // Secondary renderers store their context values on separate fields.
+    _currentValue: (defaultValue: any),
+    _currentValue2: (defaultValue: any),
+    // Used to track how many concurrent renderers this context currently
+    // supports within in a single renderer. Such as parallel server rendering.
+    _threadCount: 0,
+    _definitionLoaded: false,
+    // These are circular
+    Provider: (null: any),
+    displayName: globalName,
+  };
+
+  context.Provider = {
+    $$typeof: REACT_PROVIDER_TYPE,
+    _context: context,
+  };
+
+  if (__DEV__) {
+    context._currentRenderer = null;
+    context._currentRenderer2 = null;
+  }
+  globalServerContextRegistry[globalName] = context;
+  return context;
+}
+
+// This function is called by FlightClient to create a server context sent from
+// the server. Its possible that FlightClient is creating it before the
+// definition is loaded on the server. We'll create it with a null default value
+// if thats the case and when the definition loads it will  set the correct
+// default value.
+export function getOrCreateServerContext(globalName: string) {
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      undefined,
+    );
+  }
+  return globalServerContextRegistry[globalName];
+}

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

@@ -505,6 +505,13 @@ describe('ReactFlight', () => {
 
       expect(Scheduler).toHaveYielded(['ClientBar']);
       expect(ReactNoop).toMatchRenderedOutput(<span>hi this is server</span>);
+
+      expect(() => {
+        const ServerContext2 = React.createServerContext(
+          'ServerContext',
+          'default hello from server',
+        );
+      }).toThrow('ServerContext: ServerContext already defined');
     });
 
     // @gate enableServerContext

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

@@ -1531,6 +1531,7 @@ function mountLazyComponent(
       hint = ' Did you wrap a component in React.lazy() more than once?';
     }
   }
+  debugger;
 
   // This message intentionally doesn't mention ForwardRef or MemoComponent
   // because the fact that it's a separate type of work is an

packages/react-server/src/ReactFlightServer.js

@@ -38,7 +38,7 @@ import {
   isModuleReference,
 } from './ReactFlightServerConfig';
 
-import {getOrCreateServerContext} from 'react/src/ReactServerContext';
+import {getOrCreateServerContext} from './ReactServerContext';
 import {
   Dispatcher,
   getCurrentCache,

packages/react-server/src/ReactServerContext.js

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import {
+  REACT_PROVIDER_TYPE,
+  REACT_SERVER_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
+
+import type {
+  ReactServerContext,
+  ServerContextJSONValue,
+} from 'shared/ReactTypes';
+
+import {enableServerContext} from 'shared/ReactFeatureFlags';
+import ReactSharedInternals from 'shared/ReactSharedInternals';
+
+const globalServerContextRegistry =
+  ReactSharedInternals.globalServerContextRegistry;
+
+export function createServerContext<T: ServerContextJSONValue>(
+  globalName: string,
+  defaultValue: T,
+): ReactServerContext<T> {
+  if (!enableServerContext) {
+    throw new Error('Not implemented.');
+  }
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      defaultValue,
+    );
+  }
+  const context = globalServerContextRegistry[globalName];
+  if (!context._definitionLoaded) {
+    context._currentValue = defaultValue;
+    context._currentValue2 = defaultValue;
+    context._definitionLoaded = true;
+  } else {
+    throw new Error(`ServerContext: ${globalName} already defined`);
+  }
+  return context;
+}
+
+function _createServerContext<T: ServerContextJSONValue>(
+  globalName: string,
+  defaultValue?: T,
+): ReactServerContext<T> {
+  const context: ReactServerContext<T> = {
+    $$typeof: REACT_SERVER_CONTEXT_TYPE,
+    // As a workaround to support multiple concurrent renderers, we categorize
+    // some renderers as primary and others as secondary. We only expect
+    // there to be two concurrent renderers at most: React Native (primary) and
+    // Fabric (secondary); React DOM (primary) and React ART (secondary).
+    // Secondary renderers store their context values on separate fields.
+    _currentValue: (defaultValue: any),
+    _currentValue2: (defaultValue: any),
+    // Used to track how many concurrent renderers this context currently
+    // supports within in a single renderer. Such as parallel server rendering.
+    _threadCount: 0,
+    _definitionLoaded: false,
+    // These are circular
+    Provider: (null: any),
+    displayName: globalName,
+  };
+
+  context.Provider = {
+    $$typeof: REACT_PROVIDER_TYPE,
+    _context: context,
+  };
+
+  if (__DEV__) {
+    context._currentRenderer = null;
+    context._currentRenderer2 = null;
+  }
+  globalServerContextRegistry[globalName] = context;
+  return context;
+}
+
+// This function is called by FlightClient to create a server context sent from
+// the server. Its possible that FlightClient is creating it before the
+// definition is loaded on the server. We'll create it with a null default value
+// if thats the case and when the definition loads it will  set the correct
+// default value.
+export function getOrCreateServerContext(globalName: string) {
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      undefined,
+    );
+  }
+  return globalServerContextRegistry[globalName];
+}

packages/react/src/ReactServerContext.js

@@ -17,11 +17,12 @@ import type {
   ServerContextJSONValue,
 } from 'shared/ReactTypes';
 
+import ReactSharedInternals from 'shared/ReactSharedInternals';
+
 import {enableServerContext} from 'shared/ReactFeatureFlags';
 
-const globalRegistry: {
-  [globalName: string]: ReactServerContext<any>,
-} = {};
+const globalServerContextRegistry =
+  ReactSharedInternals.globalServerContextRegistry;
 
 export function createServerContext<T: ServerContextJSONValue>(
   globalName: string,
@@ -30,10 +31,13 @@ export function createServerContext<T: ServerContextJSONValue>(
   if (!enableServerContext) {
     throw new Error('Not implemented.');
   }
-  if (!globalRegistry[globalName]) {
-    globalRegistry[globalName] = _createServerContext(globalName, defaultValue);
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      defaultValue,
+    );
   }
-  const context = globalRegistry[globalName];
+  const context = globalServerContextRegistry[globalName];
   if (!context._definitionLoaded) {
     context._currentValue = defaultValue;
     context._currentValue2 = defaultValue;
@@ -75,7 +79,7 @@ function _createServerContext<T: ServerContextJSONValue>(
     context._currentRenderer = null;
     context._currentRenderer2 = null;
   }
-  globalRegistry[globalName] = context;
+  globalServerContextRegistry[globalName] = context;
   return context;
 }
 
@@ -85,8 +89,11 @@ function _createServerContext<T: ServerContextJSONValue>(
 // if thats the case and when the definition loads it will  set the correct
 // default value.
 export function getOrCreateServerContext(globalName: string) {
-  if (!globalRegistry[globalName]) {
-    globalRegistry[globalName] = _createServerContext(globalName, undefined);
+  if (!globalServerContextRegistry[globalName]) {
+    globalServerContextRegistry[globalName] = _createServerContext(
+      globalName,
+      undefined,
+    );
   }
-  return globalRegistry[globalName];
+  return globalServerContextRegistry[globalName];
 }

packages/react/src/ReactServerContextRegistry.js

@@ -0,0 +1,5 @@
+import type {ReactServerContext} from 'shared/ReactTypes';
+
+export const globalServerContextRegistry: {
+  [globalName: string]: ReactServerContext<any>,
+} = {};

packages/react/src/ReactSharedInternals.js

@@ -11,6 +11,8 @@ import ReactCurrentBatchConfig from './ReactCurrentBatchConfig';
 import ReactCurrentActQueue from './ReactCurrentActQueue';
 import ReactCurrentOwner from './ReactCurrentOwner';
 import ReactDebugCurrentFrame from './ReactDebugCurrentFrame';
+import {enableServerContext} from 'shared/ReactFeatureFlags';
+import {globalServerContextRegistry} from './ReactServerContextRegistry';
 
 const ReactSharedInternals = {
   ReactCurrentDispatcher,
@@ -25,4 +27,8 @@ if (__DEV__) {
   ReactSharedInternals.ReactCurrentActQueue = ReactCurrentActQueue;
 }
 
+if (enableServerContext) {
+  ReactSharedInternals.globalServerContextRegistry = globalServerContextRegistry;
+}
+
 export default ReactSharedInternals;

packages/shared/forks/ReactFeatureFlags.www.js

@@ -100,7 +100,7 @@ export const deletedTreeCleanUpLevel = 3;
 export const enablePersistentOffscreenHostContainer = false;
 
 export const consoleManagedByDevToolsDuringStrictMode = true;
-export const enableServerContext = false;
+export const enableServerContext = true;
 
 // Some www surfaces are still using this. Remove once they have been migrated.
 export const enableUseMutableSource = true;