[Flight] Prevent serialized size leaking across requests (facebook#33121)

Author: eps1lon

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -692,6 +692,44 @@ describe('ReactFlightDOMEdge', () => {
     expect(html).toBe(html2);
   });
 
+  it('regression: should not leak serialized size', async () => {
+    const MAX_ROW_SIZE = 3200;
+    // This test case is a bit convoluted and may no longer trigger the original bug.
+    // Originally, the size of `promisedText` was not cleaned up so the sync portion
+    // ended up being deferred immediately when we called `renderToReadableStream` again
+    // i.e. `result2.syncText` became a Lazy element on the second request.
+    const longText = 'd'.repeat(MAX_ROW_SIZE);
+    const promisedText = Promise.resolve(longText);
+    const model = {syncText: <p>{longText}</p>, promisedText};
+
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(model),
+    );
+
+    const result = await ReactServerDOMClient.createFromReadableStream(stream, {
+      serverConsumerManifest: {
+        moduleMap: null,
+        moduleLoading: null,
+      },
+    });
+
+    const stream2 = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(model),
+    );
+
+    const result2 = await ReactServerDOMClient.createFromReadableStream(
+      stream2,
+      {
+        serverConsumerManifest: {
+          moduleMap: null,
+          moduleLoading: null,
+        },
+      },
+    );
+
+    expect(result2.syncText).toEqual(result.syncText);
+  });
+
   it('should be able to serialize any kind of typed array', async () => {
     const buffer = new Uint8Array([
       123, 4, 10, 5, 100, 255, 244, 45, 56, 67, 43, 124, 67, 89, 100, 20,

packages/react-server/src/ReactFlightServer.js

@@ -3927,18 +3927,9 @@ function emitChunk(
     return;
   }
   // For anything else we need to try to serialize it using JSON.
-  // We stash the outer parent size so we can restore it when we exit.
-  const parentSerializedSize = serializedSize;
-  // We don't reset the serialized size counter from reentry because that indicates that we
-  // are outlining a model and we actually want to include that size into the parent since
-  // it will still block the parent row. It only restores to zero at the top of the stack.
-  try {
-    // $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
-    const json: string = stringify(value, task.toJSON);
-    emitModelChunk(request, task.id, json);
-  } finally {
-    serializedSize = parentSerializedSize;
-  }
+  // $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
+  const json: string = stringify(value, task.toJSON);
+  emitModelChunk(request, task.id, json);
 }
 
 function erroredTask(request: Request, task: Task, error: mixed): void {
@@ -3976,6 +3967,11 @@ function retryTask(request: Request, task: Task): void {
   const prevDebugID = debugID;
   task.status = RENDERING;
 
+  // We stash the outer parent size so we can restore it when we exit.
+  const parentSerializedSize = serializedSize;
+  // We don't reset the serialized size counter from reentry because that indicates that we
+  // are outlining a model and we actually want to include that size into the parent since
+  // it will still block the parent row. It only restores to zero at the top of the stack.
   try {
     // Track the root so we know that we have to emit this object even though it
     // already has an ID. This is needed because we might see this object twice
@@ -4087,6 +4083,7 @@ function retryTask(request: Request, task: Task): void {
     if (__DEV__) {
       debugID = prevDebugID;
     }
+    serializedSize = parentSerializedSize;
   }
 }
 
@@ -4099,9 +4096,11 @@ function tryStreamTask(request: Request, task: Task): void {
     // so that we instead outline the row to get a new debugID if needed.
     debugID = null;
   }
+  const parentSerializedSize = serializedSize;
   try {
     emitChunk(request, task, task.model);
   } finally {
+    serializedSize = parentSerializedSize;
     if (__DEV__) {
       debugID = prevDebugID;
     }