[Flight Reply] Encode Objects Returned to the Client by Reference (#29010)

Stacked on #28997.

We can use the technique of referencing an object by its row + property
name path for temporary references - like we do for deduping. That way
we don't need to generate an ID for temporary references. Instead, they
can just be an opaque marker in the slot and it has the implicit ID of
the row + path.

Then we can stash all objects, even the ones that are actually available
to read on the server, as temporary references. Without adding anything
to the payload since the IDs are implicit. If the same object is
returned to the client, it can be referenced by reference instead of
serializing it back to the client. This also helps preserve object
identity.

We assume that the objects are immutable when they pass the boundary.

I'm not sure if this is worth it but with this mechanism, if you return
the `FormData` payload from a `useActionState` it doesn't have to be
serialized on the way back to the client. This is a common pattern for
having access to the last submission as "default value" to the form
fields. However you can still control it by replacing it with another
object if you want. In MPA mode, the temporary references are not
configured and so it needs to be serialized in that case. That's
required anyway for hydration purposes.

I'm not sure if people will actually use this in practice though or if
FormData will always be destructured into some other object like with a
library that turns it into typed data, and back. If so, the object
identity is lost.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -915,15 +915,15 @@ function parseModelString(
       }
       case 'T': {
         // Temporary Reference
-        const id = parseInt(value.slice(2), 16);
+        const reference = '$' + value.slice(2);
         const temporaryReferences = response._tempRefs;
         if (temporaryReferences == null) {
           throw new Error(
             'Missing a temporary reference set but the RSC response returned a temporary reference. ' +
               'Pass a temporaryReference option with the set that was used with the reply.',
           );
         }
-        return readTemporaryReference(temporaryReferences, id);
+        return readTemporaryReference(temporaryReferences, reference);
       }
       case 'Q': {
         // Map

packages/react-client/src/ReactFlightReplyClient.js

@@ -109,8 +109,8 @@ function serializeServerReferenceID(id: number): string {
   return '$F' + id.toString(16);
 }
 
-function serializeTemporaryReferenceID(id: number): string {
-  return '$T' + id.toString(16);
+function serializeTemporaryReferenceMarker(): string {
+  return '$T';
 }
 
 function serializeFormDataReference(id: number): string {
@@ -405,15 +405,22 @@ export function processReply(
     if (typeof value === 'object') {
       switch ((value: any).$$typeof) {
         case REACT_ELEMENT_TYPE: {
-          if (temporaryReferences === undefined) {
-            throw new Error(
-              'React Element cannot be passed to Server Functions from the Client without a ' +
-                'temporary reference set. Pass a TemporaryReferenceSet to the options.' +
-                (__DEV__ ? describeObjectForErrorMessage(parent, key) : ''),
-            );
+          if (temporaryReferences !== undefined && key.indexOf(':') === -1) {
+            // TODO: If the property name contains a colon, we don't dedupe. Escape instead.
+            const parentReference = writtenObjects.get(parent);
+            if (parentReference !== undefined) {
+              // If the parent has a reference, we can refer to this object indirectly
+              // through the property name inside that parent.
+              const reference = parentReference + ':' + key;
+              // Store this object so that the server can refer to it later in responses.
+              writeTemporaryReference(temporaryReferences, reference, value);
+              return serializeTemporaryReferenceMarker();
+            }
           }
-          return serializeTemporaryReferenceID(
-            writeTemporaryReference(temporaryReferences, value),
+          throw new Error(
+            'React Element cannot be passed to Server Functions from the Client without a ' +
+              'temporary reference set. Pass a TemporaryReferenceSet to the options.' +
+              (__DEV__ ? describeObjectForErrorMessage(parent, key) : ''),
           );
         }
         case REACT_LAZY_TYPE: {
@@ -529,7 +536,12 @@ export function processReply(
         if (parentReference !== undefined) {
           // If the parent has a reference, we can refer to this object indirectly
           // through the property name inside that parent.
-          writtenObjects.set(value, parentReference + ':' + key);
+          const reference = parentReference + ':' + key;
+          writtenObjects.set(value, reference);
+          if (temporaryReferences !== undefined) {
+            // Store this object so that the server can refer to it later in responses.
+            writeTemporaryReference(temporaryReferences, reference, value);
+          }
         }
       }
 
@@ -693,10 +705,9 @@ export function processReply(
               'Classes or null prototypes are not supported.',
           );
         }
-        // We can serialize class instances as temporary references.
-        return serializeTemporaryReferenceID(
-          writeTemporaryReference(temporaryReferences, value),
-        );
+        // We will have written this object to the temporary reference set above
+        // so we can replace it with a marker to refer to this slot later.
+        return serializeTemporaryReferenceMarker();
       }
       if (__DEV__) {
         if (
@@ -777,27 +788,41 @@ export function processReply(
         formData.set(formFieldPrefix + refId, metaDataJSON);
         return serializeServerReferenceID(refId);
       }
-      if (temporaryReferences === undefined) {
-        throw new Error(
-          'Client Functions cannot be passed directly to Server Functions. ' +
-            'Only Functions passed from the Server can be passed back again.',
-        );
+      if (temporaryReferences !== undefined && key.indexOf(':') === -1) {
+        // TODO: If the property name contains a colon, we don't dedupe. Escape instead.
+        const parentReference = writtenObjects.get(parent);
+        if (parentReference !== undefined) {
+          // If the parent has a reference, we can refer to this object indirectly
+          // through the property name inside that parent.
+          const reference = parentReference + ':' + key;
+          // Store this object so that the server can refer to it later in responses.
+          writeTemporaryReference(temporaryReferences, reference, value);
+          return serializeTemporaryReferenceMarker();
+        }
       }
-      return serializeTemporaryReferenceID(
-        writeTemporaryReference(temporaryReferences, value),
+      throw new Error(
+        'Client Functions cannot be passed directly to Server Functions. ' +
+          'Only Functions passed from the Server can be passed back again.',
       );
     }
 
     if (typeof value === 'symbol') {
-      if (temporaryReferences === undefined) {
-        throw new Error(
-          'Symbols cannot be passed to a Server Function without a ' +
-            'temporary reference set. Pass a TemporaryReferenceSet to the options.' +
-            (__DEV__ ? describeObjectForErrorMessage(parent, key) : ''),
-        );
+      if (temporaryReferences !== undefined && key.indexOf(':') === -1) {
+        // TODO: If the property name contains a colon, we don't dedupe. Escape instead.
+        const parentReference = writtenObjects.get(parent);
+        if (parentReference !== undefined) {
+          // If the parent has a reference, we can refer to this object indirectly
+          // through the property name inside that parent.
+          const reference = parentReference + ':' + key;
+          // Store this object so that the server can refer to it later in responses.
+          writeTemporaryReference(temporaryReferences, reference, value);
+          return serializeTemporaryReferenceMarker();
+        }
       }
-      return serializeTemporaryReferenceID(
-        writeTemporaryReference(temporaryReferences, value),
+      throw new Error(
+        'Symbols cannot be passed to a Server Function without a ' +
+          'temporary reference set. Pass a TemporaryReferenceSet to the options.' +
+          (__DEV__ ? describeObjectForErrorMessage(parent, key) : ''),
       );
     }
 
@@ -812,7 +837,12 @@ export function processReply(
 
   function serializeModel(model: ReactServerValue, id: number): string {
     if (typeof model === 'object' && model !== null) {
-      writtenObjects.set(model, serializeByValueID(id));
+      const reference = serializeByValueID(id);
+      writtenObjects.set(model, reference);
+      if (temporaryReferences !== undefined) {
+        // Store this object so that the server can refer to it later in responses.
+        writeTemporaryReference(temporaryReferences, reference, model);
+      }
     }
     modelRoot = model;
     // $FlowFixMe[incompatible-return] it's not going to be undefined because we'll encode it.

packages/react-client/src/ReactFlightTemporaryReferences.js

@@ -9,33 +9,23 @@
 
 interface Reference {}
 
-export opaque type TemporaryReferenceSet = Array<Reference | symbol>;
+export opaque type TemporaryReferenceSet = Map<string, Reference | symbol>;
 
 export function createTemporaryReferenceSet(): TemporaryReferenceSet {
-  return [];
+  return new Map();
 }
 
 export function writeTemporaryReference(
   set: TemporaryReferenceSet,
+  reference: string,
   object: Reference | symbol,
-): number {
-  // We always create a new entry regardless if we've already written the same
-  // object. This ensures that we always generate a deterministic encoding of
-  // each slot in the reply for cacheability.
-  const newId = set.length;
-  set.push(object);
-  return newId;
+): void {
+  set.set(reference, object);
 }
 
 export function readTemporaryReference<T>(
   set: TemporaryReferenceSet,
-  id: number,
+  reference: string,
 ): T {
-  if (id < 0 || id >= set.length) {
-    throw new Error(
-      "The RSC response contained a reference that doesn't exist in the temporary reference set. " +
-        'Always pass the matching set that was used to create the reply when parsing its response.',
-    );
-  }
-  return (set[id]: any);
+  return (set.get(reference): any);
 }

packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js

@@ -47,15 +47,30 @@ export {
   registerClientReference,
 } from './ReactFlightESMReferences';
 
+import type {TemporaryReferenceSet} from 'react-server/src/ReactFlightServerTemporaryReferences';
+
+export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTemporaryReferences';
+
+export type {TemporaryReferenceSet};
+
 function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }
 
+function createCancelHandler(request: Request, reason: string) {
+  return () => {
+    stopFlowing(request);
+    // eslint-disable-next-line react-internal/prod-error-codes
+    abort(request, new Error(reason));
+  };
+}
+
 type Options = {
   environmentName?: string,
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
   identifierPrefix?: string,
+  temporaryReferences?: TemporaryReferenceSet,
 };
 
 type PipeableStream = {
@@ -75,6 +90,7 @@ function renderToPipeableStream(
     options ? options.identifierPrefix : undefined,
     options ? options.onPostpone : undefined,
     options ? options.environmentName : undefined,
+    options ? options.temporaryReferences : undefined,
   );
   let hasStartedFlowing = false;
   startWork(request);
@@ -88,10 +104,20 @@ function renderToPipeableStream(
       hasStartedFlowing = true;
       startFlowing(request, destination);
       destination.on('drain', createDrainHandler(destination, request));
+      destination.on(
+        'error',
+        createCancelHandler(
+          request,
+          'The destination stream errored while writing data.',
+        ),
+      );
+      destination.on(
+        'close',
+        createCancelHandler(request, 'The destination stream closed early.'),
+      );
       return destination;
     },
     abort(reason: mixed) {
-      stopFlowing(request);
       abort(request, reason);
     },
   };
@@ -155,13 +181,19 @@ function decodeReplyFromBusboy<T>(
 function decodeReply<T>(
   body: string | FormData,
   moduleBasePath: ServerManifest,
+  options?: {temporaryReferences?: TemporaryReferenceSet},
 ): Thenable<T> {
   if (typeof body === 'string') {
     const form = new FormData();
     form.append('0', body);
     body = form;
   }
-  const response = createResponse(moduleBasePath, '', body);
+  const response = createResponse(
+    moduleBasePath,
+    '',
+    options ? options.temporaryReferences : undefined,
+    body,
+  );
   const root = getRoot<T>(response);
   close(response);
   return root;

packages/react-server-dom-turbopack/src/ReactFlightDOMServerBrowser.js

@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -25,18 +26,28 @@ import {
   getRoot,
 } from 'react-server/src/ReactFlightReplyServer';
 
-import {decodeAction} from 'react-server/src/ReactFlightActionServer';
+import {
+  decodeAction,
+  decodeFormState,
+} from 'react-server/src/ReactFlightActionServer';
 
 export {
   registerServerReference,
   registerClientReference,
   createClientModuleProxy,
 } from './ReactFlightTurbopackReferences';
 
+import type {TemporaryReferenceSet} from 'react-server/src/ReactFlightServerTemporaryReferences';
+
+export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTemporaryReferences';
+
+export type {TemporaryReferenceSet};
+
 type Options = {
   environmentName?: string,
   identifierPrefix?: string,
   signal?: AbortSignal,
+  temporaryReferences?: TemporaryReferenceSet,
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
 };
@@ -53,6 +64,7 @@ function renderToReadableStream(
     options ? options.identifierPrefix : undefined,
     options ? options.onPostpone : undefined,
     options ? options.environmentName : undefined,
+    options ? options.temporaryReferences : undefined,
   );
   if (options && options.signal) {
     const signal = options.signal;
@@ -75,7 +87,10 @@ function renderToReadableStream(
       pull: (controller): ?Promise<void> => {
         startFlowing(request, controller);
       },
-      cancel: (reason): ?Promise<void> => {},
+      cancel: (reason): ?Promise<void> => {
+        stopFlowing(request);
+        abort(request, reason);
+      },
     },
     // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
     {highWaterMark: 0},
@@ -86,16 +101,22 @@ function renderToReadableStream(
 function decodeReply<T>(
   body: string | FormData,
   turbopackMap: ServerManifest,
+  options?: {temporaryReferences?: TemporaryReferenceSet},
 ): Thenable<T> {
   if (typeof body === 'string') {
     const form = new FormData();
     form.append('0', body);
     body = form;
   }
-  const response = createResponse(turbopackMap, '', body);
+  const response = createResponse(
+    turbopackMap,
+    '',
+    options ? options.temporaryReferences : undefined,
+    body,
+  );
   const root = getRoot<T>(response);
   close(response);
   return root;
 }
 
-export {renderToReadableStream, decodeReply, decodeAction};
+export {renderToReadableStream, decodeReply, decodeAction, decodeFormState};