[Flight] Lazily parse models and allow any value to suspend (#18476)

* Lazily initialize models as they're read intead of eagerly when received

This ensures that we don't spend CPU cycles processing models that we're
not going to end up rendering.

This model will also allow us to suspend during this initialization if
data is not yet available to satisfy the model.

* Refactoring carefully to ensure bundles still compile to something optimal

* Remove generic from Response

The root model needs to be cast at one point or another same as othe
chunks. So we can parameterize the read instead of the whole Response.

* Read roots from the 0 key of the map

The special case to read the root isn't worth the field and code.

* Store response on each Chunk

Instead of storing it on the data tuple which is kind of dynamic, we store
it on each Chunk. This uses more memory. Especially compared to just making
initializeBlock a closure, but overall is simpler.

* Rename private fields to underscores

Response objects are exposed.

* Encode server components as delayed references

This allows us to stream in server components one after another over the
wire. It also allows parallelizing their fetches and resuming only the
server component instead of the whole parent block.

This doesn't yet allow us to suspend deeper while waiting on this content
because we don't have "lazy elements".

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -14,12 +14,15 @@ import type {LazyComponent} from 'react/src/ReactLazy';
 import type {
   ModuleReference,
   ModuleMetaData,
+  UninitializedModel,
+  Response,
 } from './ReactFlightClientHostConfig';
 
 import {
   resolveModuleReference,
   preloadModule,
   requireModule,
+  parseModel,
 } from './ReactFlightClientHostConfig';
 
 import {
@@ -33,33 +36,48 @@ export type JSONValue =
   | null
   | boolean
   | string
-  | {[key: string]: JSONValue}
-  | Array<JSONValue>;
+  | {+[key: string]: JSONValue}
+  | $ReadOnlyArray<JSONValue>;
 
 const PENDING = 0;
-const RESOLVED = 1;
-const ERRORED = 2;
+const RESOLVED_MODEL = 1;
+const INITIALIZED = 2;
+const ERRORED = 3;
 
 type PendingChunk = {
   _status: 0,
   _value: null | Array<() => mixed>,
+  _response: Response,
   then(resolve: () => mixed): void,
 };
-type ResolvedChunk<T> = {
+type ResolvedModelChunk = {
   _status: 1,
+  _value: UninitializedModel,
+  _response: Response,
+  then(resolve: () => mixed): void,
+};
+type InitializedChunk<T> = {
+  _status: 2,
   _value: T,
+  _response: Response,
   then(resolve: () => mixed): void,
 };
 type ErroredChunk = {
-  _status: 2,
+  _status: 3,
   _value: Error,
+  _response: Response,
   then(resolve: () => mixed): void,
 };
-type SomeChunk<T> = PendingChunk | ResolvedChunk<T> | ErroredChunk;
+type SomeChunk<T> =
+  | PendingChunk
+  | ResolvedModelChunk
+  | InitializedChunk<T>
+  | ErroredChunk;
 
-function Chunk(status: any, value: any) {
+function Chunk(status: any, value: any, response: Response) {
   this._status = status;
   this._value = value;
+  this._response = response;
 }
 Chunk.prototype.then = function<T>(resolve: () => mixed) {
   const chunk: SomeChunk<T> = this;
@@ -73,45 +91,40 @@ Chunk.prototype.then = function<T>(resolve: () => mixed) {
   }
 };
 
-export type Response<T> = {
-  partialRow: string,
-  rootChunk: SomeChunk<T>,
-  chunks: Map<number, SomeChunk<any>>,
-  readRoot(): T,
+export type ResponseBase = {
+  _chunks: Map<number, SomeChunk<any>>,
+  readRoot<T>(): T,
+  ...
 };
 
-function readRoot<T>(): T {
-  const response: Response<T> = this;
-  const rootChunk = response.rootChunk;
-  if (rootChunk._status === RESOLVED) {
-    return rootChunk._value;
-  } else if (rootChunk._status === PENDING) {
-    // eslint-disable-next-line no-throw-literal
-    throw (rootChunk: Wakeable);
-  } else {
-    throw rootChunk._value;
+export type {Response};
+
+function readChunk<T>(chunk: SomeChunk<T>): T {
+  switch (chunk._status) {
+    case INITIALIZED:
+      return chunk._value;
+    case RESOLVED_MODEL:
+      return initializeModelChunk(chunk);
+    case PENDING:
+      // eslint-disable-next-line no-throw-literal
+      throw (chunk: Wakeable);
+    default:
+      throw chunk._value;
   }
 }
 
-export function createResponse<T>(): Response<T> {
-  const rootChunk: SomeChunk<any> = createPendingChunk();
-  const chunks: Map<number, SomeChunk<any>> = new Map();
-  chunks.set(0, rootChunk);
-  const response = {
-    partialRow: '',
-    rootChunk,
-    chunks: chunks,
-    readRoot: readRoot,
-  };
-  return response;
+function readRoot<T>(): T {
+  const response: Response = this;
+  const chunk = getChunk(response, 0);
+  return readChunk(chunk);
 }
 
-function createPendingChunk(): PendingChunk {
-  return new Chunk(PENDING, null);
+function createPendingChunk(response: Response): PendingChunk {
+  return new Chunk(PENDING, null, response);
 }
 
-function createErrorChunk(error: Error): ErroredChunk {
-  return new Chunk(ERRORED, error);
+function createErrorChunk(response: Response, error: Error): ErroredChunk {
+  return new Chunk(ERRORED, error, response);
 }
 
 function wakeChunk(listeners: null | Array<() => mixed>) {
@@ -135,29 +148,40 @@ function triggerErrorOnChunk<T>(chunk: SomeChunk<T>, error: Error): void {
   wakeChunk(listeners);
 }
 
-function createResolvedChunk<T>(value: T): ResolvedChunk<T> {
-  return new Chunk(RESOLVED, value);
+function createResolvedModelChunk(
+  response: Response,
+  value: UninitializedModel,
+): ResolvedModelChunk {
+  return new Chunk(RESOLVED_MODEL, value, response);
 }
 
-function resolveChunk<T>(chunk: SomeChunk<T>, value: T): void {
+function resolveModelChunk<T>(
+  chunk: SomeChunk<T>,
+  value: UninitializedModel,
+): void {
   if (chunk._status !== PENDING) {
     // We already resolved. We didn't expect to see this.
     return;
   }
   const listeners = chunk._value;
-  const resolvedChunk: ResolvedChunk<T> = (chunk: any);
-  resolvedChunk._status = RESOLVED;
+  const resolvedChunk: ResolvedModelChunk = (chunk: any);
+  resolvedChunk._status = RESOLVED_MODEL;
   resolvedChunk._value = value;
   wakeChunk(listeners);
 }
 
+function initializeModelChunk<T>(chunk: ResolvedModelChunk): T {
+  const value: T = parseModel(chunk._response, chunk._value);
+  const initializedChunk: InitializedChunk<T> = (chunk: any);
+  initializedChunk._status = INITIALIZED;
+  initializedChunk._value = value;
+  return value;
+}
+
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
-export function reportGlobalError<T>(
-  response: Response<T>,
-  error: Error,
-): void {
-  response.chunks.forEach(chunk => {
+export function reportGlobalError(response: Response, error: Error): void {
+  response._chunks.forEach(chunk => {
     // If this chunk was already resolved or errored, it won't
     // trigger an error but if it wasn't then we need to
     // because we won't be getting any new data to resolve it.
@@ -171,14 +195,7 @@ function readMaybeChunk<T>(maybeChunk: SomeChunk<T> | T): T {
     return maybeChunk;
   }
   const chunk: SomeChunk<T> = (maybeChunk: any);
-  if (chunk._status === RESOLVED) {
-    return chunk._value;
-  } else if (chunk._status === PENDING) {
-    // eslint-disable-next-line no-throw-literal
-    throw (chunk: Wakeable);
-  } else {
-    throw chunk._value;
-  }
+  return readChunk(chunk);
 }
 
 function createElement(type, key, props): React$Element<any> {
@@ -226,6 +243,7 @@ type UninitializedBlockPayload<Data> = [
   mixed,
   ModuleMetaData | SomeChunk<ModuleMetaData>,
   Data | SomeChunk<Data>,
+  Response,
 ];
 
 function initializeBlock<Props, Data>(
@@ -267,83 +285,102 @@ function createLazyBlock<Props, Data>(
   return lazyType;
 }
 
-export function parseModelFromJSON<T>(
-  response: Response<T>,
-  targetObj: Object,
-  key: string,
-  value: JSONValue,
-): mixed {
-  if (typeof value === 'string') {
-    if (value[0] === '$') {
-      if (value === '$') {
-        return REACT_ELEMENT_TYPE;
-      } else if (value[1] === '$' || value[1] === '@') {
-        // This was an escaped string value.
-        return value.substring(1);
-      } else {
-        const id = parseInt(value.substring(1), 16);
-        const chunks = response.chunks;
-        let chunk = chunks.get(id);
-        if (!chunk) {
-          chunk = createPendingChunk();
-          chunks.set(id, chunk);
-        }
+function getChunk(response: Response, id: number): SomeChunk<any> {
+  const chunks = response._chunks;
+  let chunk = chunks.get(id);
+  if (!chunk) {
+    chunk = createPendingChunk(response);
+    chunks.set(id, chunk);
+  }
+  return chunk;
+}
+
+export function parseModelString(
+  response: Response,
+  parentObject: Object,
+  value: string,
+): any {
+  if (value[0] === '$') {
+    if (value === '$') {
+      return REACT_ELEMENT_TYPE;
+    } else if (value[1] === '$' || value[1] === '@') {
+      // This was an escaped string value.
+      return value.substring(1);
+    } else {
+      const id = parseInt(value.substring(1), 16);
+      const chunk = getChunk(response, id);
+      if (parentObject[0] === REACT_BLOCK_TYPE) {
+        // Block types know how to deal with lazy values.
         return chunk;
       }
-    }
-    if (value === '@') {
-      return REACT_BLOCK_TYPE;
+      // For anything else we must Suspend this block if
+      // we don't yet have the value.
+      return readChunk(chunk);
     }
   }
-  if (typeof value === 'object' && value !== null) {
-    const tuple: [mixed, mixed, mixed, mixed] = (value: any);
-    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_BLOCK_TYPE: {
-        // TODO: Consider having React just directly accept these arrays as blocks.
-        return createLazyBlock((tuple: any));
-      }
-    }
+  if (value === '@') {
+    return REACT_BLOCK_TYPE;
   }
   return value;
 }
 
-export function resolveModelChunk<T, M>(
-  response: Response<T>,
+export function parseModelTuple(
+  response: Response,
+  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]);
+  } else if (tuple[0] === REACT_BLOCK_TYPE) {
+    // TODO: Consider having React just directly accept these arrays as blocks.
+    return createLazyBlock((tuple: any));
+  }
+  return value;
+}
+
+export function createResponse(): ResponseBase {
+  const chunks: Map<number, SomeChunk<any>> = new Map();
+  const response = {
+    _chunks: chunks,
+    readRoot: readRoot,
+  };
+  return response;
+}
+
+export function resolveModel(
+  response: Response,
   id: number,
-  model: M,
+  model: UninitializedModel,
 ): void {
-  const chunks = response.chunks;
+  const chunks = response._chunks;
   const chunk = chunks.get(id);
   if (!chunk) {
-    chunks.set(id, createResolvedChunk(model));
+    chunks.set(id, createResolvedModelChunk(response, model));
   } else {
-    resolveChunk(chunk, model);
+    resolveModelChunk(chunk, model);
   }
 }
 
-export function resolveErrorChunk<T>(
-  response: Response<T>,
+export function resolveError(
+  response: Response,
   id: number,
   message: string,
   stack: string,
 ): void {
   const error = new Error(message);
   error.stack = stack;
-  const chunks = response.chunks;
+  const chunks = response._chunks;
   const chunk = chunks.get(id);
   if (!chunk) {
-    chunks.set(id, createErrorChunk(error));
+    chunks.set(id, createErrorChunk(response, error));
   } else {
     triggerErrorOnChunk(chunk, error);
   }
 }
 
-export function close<T>(response: Response<T>): void {
+export function close(response: Response): void {
   // In case there are any remaining unresolved chunks, they won't
   // be resolved now. So we need to issue an error to those.
   // Ideally we should be able to early bail out if we kept a

packages/react-client/src/ReactFlightClientHostConfigStream.js

@@ -0,0 +1,23 @@
+/**
+ * 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 type {ResponseBase} from './ReactFlightClient';
+import type {StringDecoder} from './ReactFlightClientHostConfig';
+
+export type Response = ResponseBase & {
+  _partialRow: string,
+  _fromJSON: (key: string, value: JSONValue) => any,
+  _stringDecoder: StringDecoder,
+};
+
+export type UninitializedModel = string;
+
+export function parseModel<T>(response: Response, json: UninitializedModel): T {
+  return JSON.parse(json, response._fromJSON);
+}