(cherrypick) [Fizz][Float] Do not write after closing the stream (#27541 ) (#31881 )

Float methods can hang on to a reference to a Request after the request is closed due to AsyncLocalStorage. If a Float method is called at this point we do not want to attempt to flush anything. This change updates the closing logic to also call `stopFlowing` which will ensure that any checks against the destination properly reflect that we cannot do any writes. In addition it updates the enqueueFlush logic to existence check the destination inside the work function since it can change across the work scheduling gap if it is async.
(cherrypick) [Flight Reply] Reject any new Chunks not yet discovered at the time of reportGlobalError (#31840 ) (#31872 )
2024-12-20 14:44:43 -08:00 · 2024-12-20 09:01:12 -08:00
4 changed files with 106 additions and 2 deletions
--- a/packages/react-dom/src/tests/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/tests/ReactDOMFizzServer-test.js
@@ -3624,6 +3624,46 @@ describe('ReactDOMFizzServer', () => {
    );
  });

+  // https://github.com/facebook/react/issues/27540
+  // This test is not actually asserting much because there is possibly a bug in the closeing logic for the
+  // Node implementation of Fizz. The close leads to an abort which sets the destination to null before the Float
+  // method has an opportunity to schedule a write. We should fix this probably and once we do this test will start
+  // to fail if the underyling issue of writing after stream completion isn't fixed
+  it('does not try to write to the stream after it has been closed', async () => {
+    async function preloadLate() {
+      await 1;
+      ReactDOM.preconnect('foo');
+    }
+
+    function Preload() {
+      preloadLate();
+      return null;
+    }
+
+    function App() {
+      return (
+        <html>
+          <body>
+            <main>hello</main>
+            <Preload />
+          </body>
+        </html>
+      );
+    }
+    await act(() => {
+      renderToPipeableStream(<App />).pipe(writable);
+    });
+
+    expect(getVisibleChildren(document)).toEqual(
+      <html>
+        <head />
+        <body>
+          <main>hello</main>
+        </body>
+      </html>,
+    );
+  });
+
  describe('error escaping', () => {
    it('escapes error hash, message, and component stack values in directly flushed errors (html escaping)', async () => {
      window.__outlet = {};
--- a/packages/react-dom/src/tests/ReactDOMFizzServerBrowser-test.js
+++ b/packages/react-dom/src/tests/ReactDOMFizzServerBrowser-test.js
@@ -15,6 +15,7 @@ global.ReadableStream =
 global.TextEncoder = require('util').TextEncoder;

 let React;
+let ReactDOM;
 let ReactDOMFizzServer;
 let Suspense;

@@ -22,6 +23,7 @@ describe('ReactDOMFizzServerBrowser', () => {
  beforeEach(() => {
    jest.resetModules();
    React = require('react');
+    ReactDOM = require('react-dom');
    ReactDOMFizzServer = require('react-dom/server.browser');
    Suspense = React.Suspense;
  });
@@ -547,4 +549,37 @@ describe('ReactDOMFizzServerBrowser', () => {
    // However, it does error the shell.
    expect(caughtError.message).toEqual('testing postpone');
  });
+
+  // https://github.com/facebook/react/issues/27540
+  // This test is not actually asserting much because in our test environment the Float method cannot find the request after
+  // an await and thus is a noop. If we fix our test environment to support AsyncLocalStorage we can assert that the
+  // stream does not write after closing.
+  it('does not try to write to the stream after it has been closed', async () => {
+    async function preloadLate() {
+      await 1;
+      ReactDOM.preconnect('foo');
+    }
+
+    function Preload() {
+      preloadLate();
+      return null;
+    }
+
+    function App() {
+      return (
+        <html>
+          <body>
+            <main>hello</main>
+            <Preload />
+          </body>
+        </html>
+      );
+    }
+    const stream = await ReactDOMFizzServer.renderToReadableStream(<App />);
+    const result = await readResult(stream);
+
+    expect(result).toMatchInlineSnapshot(
+      `"<!DOCTYPE html><html><head></head><body><main>hello</main></body></html>"`,
+    );
+  });
 });
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -4041,6 +4041,9 @@ function flushCompletedQueues(
      }
      // We're done.
      close(destination);
+      // We need to stop flowing now because we do not want any async contexts which might call
+      // float methods to initiate any flushes after this point
+      stopFlowing(request);
    } else {
      completeWriting(destination);
      flushBuffered(destination);
@@ -4066,9 +4069,17 @@ function enqueueFlush(request: Request): void {
    // happen when we start flowing again
    request.destination !== null
  ) {
-    const destination = request.destination;
    request.flushScheduled = true;
-    scheduleWork(() => flushCompletedQueues(request, destination));
+    scheduleWork(() => {
+      // We need to existence check destination again here because it might go away
+      // in between the enqueueFlush call and the work execution
+      const destination = request.destination;
+      if (destination) {
+        flushCompletedQueues(request, destination);
+      } else {
+        request.flushScheduled = false;
+      }
+    });
  }
 }

--- a/packages/react-server/src/ReactFlightReplyServer.js
+++ b/packages/react-server/src/ReactFlightReplyServer.js
@@ -135,6 +135,8 @@ export type Response = {
  _formData: FormData,
  _chunks: Map<number, SomeChunk<any>>,
  _fromJSON: (key: string, value: JSONValue) => any,
+  _closed: boolean,
+  _closedReason: mixed,
 };

 export function getRoot<T>(response: Response): Thenable<T> {
@@ -198,6 +200,14 @@ function createResolvedModelChunk<T>(
  return new Chunk(RESOLVED_MODEL, value, null, response);
 }

+function createErroredChunk<T>(
+  response: Response,
+  reason: mixed,
+): ErroredChunk<T> {
+  // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
+  return new Chunk(ERRORED, null, reason, response);
+}
+
 function resolveModelChunk<T>(chunk: SomeChunk<T>, value: string): void {
  if (chunk.status !== PENDING) {
    // We already resolved. We didn't expect to see this.
@@ -297,6 +307,8 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
 // 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(response: Response, error: Error): void {
+  response._closed = true;
+  response._closedReason = error;
  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
@@ -318,6 +330,10 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
    if (backingEntry != null) {
      // We assume that this is a string entry for now.
      chunk = createResolvedModelChunk(response, (backingEntry: any));
+    } else if (response._closed) {
+      // We have already errored the response and we're not going to get
+      // anything more streaming in so this will immediately error.
+      chunk = createErroredChunk(response, response._closedReason);
    } else {
      // We're still waiting on this entry to stream in.
      chunk = createPendingChunk(response);
@@ -519,6 +535,8 @@ export function createResponse(
      }
      return value;
    },
+    _closed: false,
+    _closedReason: null,
  };
  return response;
 }