Use an explicit temporary reference set on the server

This ensures that objects can only be passed by reference if they're part of the same pair of renders. This also ensures that you can avoid this short cut when it's not possible such as in SSR passes.
facebook · May 9, 2024 · 7c69449 · 7c69449
1 parent 73098e2
commit 7c69449
Show file tree

Hide file tree

Showing 13 changed files with 305 additions and 61 deletions.
diff --git a/packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js b/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;

diff --git a/packages/react-server-dom-turbopack/src/ReactFlightDOMServerBrowser.js b/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};
diff --git a/packages/react-server-dom-turbopack/src/ReactFlightDOMServerEdge.js b/packages/react-server-dom-turbopack/src/ReactFlightDOMServerEdge.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};
diff --git a/packages/react-server-dom-turbopack/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-turbopack/src/ReactFlightDOMServerNode.js
@@ -22,6 +22,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -36,23 +37,41 @@ 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};
+
 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 = {
@@ -72,6 +91,7 @@ function renderToPipeableStream(
     options ? options.identifierPrefix : undefined,
     options ? options.onPostpone : undefined,
     options ? options.environmentName : undefined,
+    options ? options.temporaryReferences : undefined,
   );
   let hasStartedFlowing = false;
   startWork(request);
@@ -85,6 +105,17 @@ 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) {
@@ -151,13 +182,19 @@ function decodeReplyFromBusboy<T>(
 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;
@@ -168,4 +205,5 @@ export {
   decodeReplyFromBusboy,
   decodeReply,
   decodeAction,
+  decodeFormState,
 };
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js
@@ -37,10 +37,17 @@ export {
   createClientModuleProxy,
 } from './ReactFlightWebpackReferences';
 
+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,
 };
@@ -57,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;
@@ -93,13 +101,19 @@ function renderToReadableStream(
 function decodeReply<T>(
   body: string | FormData,
   webpackMap: ServerManifest,
+  options?: {temporaryReferences?: TemporaryReferenceSet},
 ): Thenable<T> {
   if (typeof body === 'string') {
     const form = new FormData();
     form.append('0', body);
     body = form;
   }
-  const response = createResponse(webpackMap, '', body);
+  const response = createResponse(
+    webpackMap,
+    '',
+    options ? options.temporaryReferences : undefined,
+    body,
+  );
   const root = getRoot<T>(response);
   close(response);
   return root;