From 86ae856a9563612138b224a6544d8b909176ed9a Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Wed, 9 Oct 2024 21:36:21 -0700
Subject: [PATCH] [Flight] don't chunks for rejected thenables after abort

When aborting we emit chunks for each pending task. However there was a bug where a thenable could also reject before we could flush and we end up with an extra chunk throwing off the pendingChunks bookeeping. When a task is retried we skip it if is is not in PENDING status because we understand it was completed some other way. We need to replciate this for the reject pathway on serialized thenables since aborting if effectively completing all pending tasks and not something we need to continue to do once the thenable rejects later.
---
 .../src/__tests__/ReactFlightDOM-test.js      | 50 +++++++++++++++++++
 .../react-server/src/ReactFlightServer.js     | 35 +++++++------
 2 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
index b59eb05c7b3fb..febb5faf4ccec 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
@@ -3084,4 +3084,54 @@ describe('ReactFlightDOM', () => {
       </div>,
     );
   });
+
+  it('rejecting a thenable after an abort before flush should not lead to a frozen readable', async () => {
+    const ClientComponent = clientExports(function (props: {
+      promise: Promise<void>,
+    }) {
+      return 'hello world';
+    });
+
+    let reject;
+    const promise = new Promise((_, re) => {
+      reject = re;
+    });
+
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="loading...">
+            <ClientComponent promise={promise} />
+          </Suspense>
+        </div>
+      );
+    }
+
+    const errors = [];
+    const {writable, readable} = getTestStream();
+    const {pipe, abort} = await serverAct(() =>
+      ReactServerDOMServer.renderToPipeableStream(<App />, webpackMap, {
+        onError(x) {
+          errors.push(x);
+        },
+      }),
+    );
+    await serverAct(() => {
+      abort('STOP');
+      reject('STOP');
+    });
+    pipe(writable);
+
+    const reader = readable.getReader();
+    while (true) {
+      const {done} = await reader.read();
+      if (done) {
+        break;
+      }
+    }
+
+    expect(errors).toEqual(['STOP']);
+
+    // We expect it to get to the end here rather than hang on the reader.
+  });
 });
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index 0b1a4d3c7fee1..6ac9cbc8fee03 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -693,22 +693,27 @@ function serializeThenable(
       pingTask(request, newTask);
     },
     reason => {
-      if (
-        enablePostpone &&
-        typeof reason === 'object' &&
-        reason !== null &&
-        (reason: any).$$typeof === REACT_POSTPONE_TYPE
-      ) {
-        const postponeInstance: Postpone = (reason: any);
-        logPostpone(request, postponeInstance.message, newTask);
-        emitPostponeChunk(request, newTask.id, postponeInstance);
-      } else {
-        const digest = logRecoverableError(request, reason, newTask);
-        emitErrorChunk(request, newTask.id, digest, reason);
+      if (newTask.status === PENDING) {
+        // We expect that the only status it might be otherwise is ABORTED.
+        // When we abort we emit chunks in each pending task slot and don't need
+        // to do so again here.
+        if (
+          enablePostpone &&
+          typeof reason === 'object' &&
+          reason !== null &&
+          (reason: any).$$typeof === REACT_POSTPONE_TYPE
+        ) {
+          const postponeInstance: Postpone = (reason: any);
+          logPostpone(request, postponeInstance.message, newTask);
+          emitPostponeChunk(request, newTask.id, postponeInstance);
+        } else {
+          const digest = logRecoverableError(request, reason, newTask);
+          emitErrorChunk(request, newTask.id, digest, reason);
+        }
+        newTask.status = ERRORED;
+        request.abortableTasks.delete(newTask);
+        enqueueFlush(request);
       }
-      newTask.status = ERRORED;
-      request.abortableTasks.delete(newTask);
-      enqueueFlush(request);
     },
   );