Stop flowing and then abort if a stream is cancelled (#27405)

sebmarkbage · web-flow · commit d9e00f795b77 · 2023-09-22T15:16:49.000-04:00
We currently abort a stream either it's explicitly told to abort (e.g.
by an abortsignal). In this case we still finish writing what we have as
well as instructions for the client about what happened so it can
trigger fallback cases and log appropriately.

We also abort a request if the stream itself cancels. E.g. if you can't
write anymore. In this case we should not write anything to the outgoing
stream since it's supposed to be closed already now. However, we should
still abort the request so that more work isn't performed and so that we
can log the reason for it to the onError callback.

We should also not do any work after aborting.

There we need to stop the "flow" of bytes - so I call stopFlowing in the
cancel case before aborting.

The tests were testing this case but we had changed the implementation
to only start flowing at initial read (pull) instead of start like we
used to. As a result, it was no longer covering this case. We have to
call reader.read() in the tests to start the flow so that we need to
cancel it.

We also were missing a final assertion on the error logs and since we
were tracking them explicitly the extra error was silenced.
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
@@ -342,7 +342,8 @@ describe('ReactDOMFizzServerBrowser', () => {
     expect(isComplete).toBe(false);
 
     const reader = stream.getReader();
-    reader.cancel();
+    await reader.read();
+    await reader.cancel();
 
     expect(errors).toEqual([
       'The render was aborted by the server without a reason.',
@@ -355,6 +356,10 @@ describe('ReactDOMFizzServerBrowser', () => {
 
     expect(rendered).toBe(false);
     expect(isComplete).toBe(true);
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
   });
 
   it('should stream large contents that might overlow individual buffers', async () => {
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
@@ -19,6 +19,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -78,6 +79,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
@@ -158,6 +160,7 @@ function resume(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBun.js b/packages/react-dom/src/server/ReactDOMFizzServerBun.js
@@ -17,6 +17,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -68,6 +69,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerEdge.js b/packages/react-dom/src/server/ReactDOMFizzServerEdge.js
@@ -19,6 +19,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -78,6 +79,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
@@ -158,6 +160,7 @@ function resume(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerNode.js b/packages/react-dom/src/server/ReactDOMFizzServerNode.js
@@ -21,6 +21,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -35,9 +36,12 @@ function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }
 
-function createAbortHandler(request: Request, reason: string) {
-  // eslint-disable-next-line react-internal/prod-error-codes
-  return () => abort(request, new Error(reason));
+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 = {
@@ -122,14 +126,14 @@ function renderToPipeableStream(
       destination.on('drain', createDrainHandler(destination, request));
       destination.on(
         'error',
-        createAbortHandler(
+        createCancelHandler(
           request,
           'The destination stream errored while writing data.',
         ),
       );
       destination.on(
         'close',
-        createAbortHandler(request, 'The destination stream closed early.'),
+        createCancelHandler(request, 'The destination stream closed early.'),
       );
       return destination;
     },
@@ -180,14 +184,14 @@ function resumeToPipeableStream(
       destination.on('drain', createDrainHandler(destination, request));
       destination.on(
         'error',
-        createAbortHandler(
+        createCancelHandler(
           request,
           'The destination stream errored while writing data.',
         ),
       );
       destination.on(
         'close',
-        createAbortHandler(request, 'The destination stream closed early.'),
+        createCancelHandler(request, 'The destination stream closed early.'),
       );
       return destination;
     },
diff --git a/packages/react-dom/src/server/ReactDOMFizzStaticBrowser.js b/packages/react-dom/src/server/ReactDOMFizzStaticBrowser.js
@@ -18,6 +18,7 @@ import {
   createPrerenderRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
   getPostponedState,
 } from 'react-server/src/ReactFizzServer';
@@ -61,6 +62,10 @@ function prerender(
           pull: (controller): ?Promise<void> => {
             startFlowing(request, controller);
           },
+          cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
+            abort(request);
+          },
         },
         // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
         {highWaterMark: 0},
diff --git a/packages/react-dom/src/server/ReactDOMFizzStaticEdge.js b/packages/react-dom/src/server/ReactDOMFizzStaticEdge.js
@@ -18,6 +18,7 @@ import {
   createPrerenderRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
   getPostponedState,
 } from 'react-server/src/ReactFizzServer';
@@ -61,6 +62,10 @@ function prerender(
           pull: (controller): ?Promise<void> => {
             startFlowing(request, controller);
           },
+          cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
+            abort(request);
+          },
         },
         // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
         {highWaterMark: 0},
diff --git a/packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js
@@ -22,6 +22,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -90,6 +91,7 @@ function renderToPipeableStream(
       return destination;
     },
     abort(reason: mixed) {
+      stopFlowing(request);
       abort(request, reason);
     },
   };
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js
@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -78,7 +79,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},
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerEdge.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerEdge.js
@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -78,7 +79,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},
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js
@@ -22,6 +22,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -51,6 +52,14 @@ 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 = {
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
@@ -88,6 +97,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) {
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
@@ -1222,4 +1222,53 @@ describe('ReactFlightDOMBrowser', () => {
 
     expect(postponed).toBe('testing postpone');
   });
+
+  it('should not continue rendering after the reader cancels', async () => {
+    let hasLoaded = false;
+    let resolve;
+    let rendered = false;
+    const promise = new Promise(r => (resolve = r));
+    function Wait() {
+      if (!hasLoaded) {
+        throw promise;
+      }
+      rendered = true;
+      return 'Done';
+    }
+    const errors = [];
+    const stream = await ReactServerDOMServer.renderToReadableStream(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <Wait />
+        </Suspense>
+      </div>,
+      null,
+      {
+        onError(x) {
+          errors.push(x.message);
+        },
+      },
+    );
+
+    expect(rendered).toBe(false);
+
+    const reader = stream.getReader();
+    await reader.read();
+    await reader.cancel();
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
+
+    hasLoaded = true;
+    resolve();
+
+    await jest.runAllTimers();
+
+    expect(rendered).toBe(false);
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
+  });
 });
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
@@ -3998,6 +3998,10 @@ export function startFlowing(request: Request, destination: Destination): void {
   }
 }
 
+export function stopFlowing(request: Request): void {
+  request.destination = null;
+}
+
 // This is called to early terminate a request. It puts all pending boundaries in client rendered state.
 export function abort(request: Request, reason: mixed): void {
   try {
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
@@ -359,6 +359,7 @@ function serializeThenable(request: Request, thenable: Thenable<any>): number {
     },
     reason => {
       newTask.status = ERRORED;
+      request.abortableTasks.delete(newTask);
       // TODO: We should ideally do this inside performWork so it's scheduled
       const digest = logRecoverableError(request, reason);
       emitErrorChunk(request, newTask.id, digest, reason);
@@ -1570,6 +1571,10 @@ export function startFlowing(request: Request, destination: Destination): void {
   }
 }
 
+export function stopFlowing(request: Request): void {
+  request.destination = null;
+}
+
 // This is called to early terminate a request. It creates an error at all pending tasks.
 export function abort(request: Request, reason: mixed): void {
   try {
