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

gnoff · gnoff · commit 5bbb04e9869f · 2024-12-20T14:43:38.000-08:00
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. fixes: facebook#27540
diff --git 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 = {};
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
@@ -15,13 +15,15 @@ global.ReadableStream =
 global.TextEncoder = require('util').TextEncoder;
 
 let React;
+let ReactDOM;
 let ReactDOMFizzServer;
 let Suspense;
 
 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>"`,
+    );
+  });
 });
diff --git 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;
+      }
+    });
   }
 }
 
