[pull] main from react:main

packages/react-client/src/ReactFlightReplyClient.js

@@ -39,6 +39,12 @@ import getPrototypeOf from 'shared/getPrototypeOf';
 
 const ObjectPrototype = Object.prototype;
 
+// Passed to replyLifetimeController.abort(). Nothing reads the reason, but a
+// call to abort() without one constructs an AbortError DOMException. Capturing
+// the stack trace dominates that cost, and the cost grows with the depth of the
+// stack.
+const REPLY_ENDED = 'The reply ended.';
+
 import {
   usedWithSSR,
   checkEvalAvailabilityOnceDev,
@@ -184,14 +190,58 @@ export function processReply(
   root: ReactServerValue,
   formFieldPrefix: string,
   temporaryReferences: void | TemporaryReferenceSet,
-  resolve: (string | FormData) => void,
-  reject: (error: mixed) => void,
-): (reason: mixed) => void {
+  onResolve: (string | FormData) => void,
+  onReject: (error: mixed) => void,
+  signal: void | AbortSignal,
+): void {
   let nextPartId = 1;
   let pendingParts = 0;
   let formData: null | FormData = null;
   const writtenObjects: WeakMap<Reference, string> = new WeakMap();
   let modelRoot: null | ReactServerValue = root;
+  let settled = false;
+  // Bounds the abort listener that attachAbortSignal attaches to the caller's
+  // signal. Null until a signal is attached, so a reply that gets no signal
+  // never creates a controller.
+  let replyLifetimeController: null | AbortController = null;
+
+  // Ending the lifetime makes the runtime remove the caller's abort listener.
+  // Without that, the listener keeps everything this reply serialized reachable
+  // for as long as the caller's signal lives, and a composite signal from
+  // AbortSignal.any() is itself retained by the runtime while it has any abort
+  // listener attached.
+  function endReplyLifetime(): void {
+    if (replyLifetimeController !== null) {
+      replyLifetimeController.abort(REPLY_ENDED);
+    }
+  }
+
+  function resolve(value: string | FormData): void {
+    settled = true;
+    endReplyLifetime();
+    onResolve(value);
+  }
+
+  function reject(error: mixed): void {
+    settled = true;
+    endReplyLifetime();
+    onReject(error);
+  }
+
+  function attachAbortSignal(abortSignal: AbortSignal): void {
+    if (abortSignal.aborted) {
+      abort(abortSignal.reason);
+      return;
+    }
+    replyLifetimeController = new AbortController();
+    abortSignal.addEventListener(
+      'abort',
+      () => {
+        abort(abortSignal.reason);
+      },
+      {signal: replyLifetimeController.signal},
+    );
+  }
 
   if (__DEV__) {
     // We use eval to create fake function stacks which includes Component stacks.
@@ -894,6 +944,9 @@ export function processReply(
   }
 
   function abort(reason: mixed): void {
+    // Nothing can make the reply pending again from here, so the caller's
+    // signal has no further effect on it.
+    endReplyLifetime();
     if (pendingParts > 0) {
       pendingParts = 0; // Don't resolve again later.
       // Resolve with what we have so far, which may have holes at this point.
@@ -920,7 +973,17 @@ export function processReply(
     }
   }
 
-  return abort;
+  // Wired up after serializing: abort() reads `json` and resolves with the
+  // parts that finished, so it must not be reachable before then. A reply that
+  // already settled gets no listener, since aborting it would be a no-op and
+  // the lifetime that removes the listener has already ended.
+  //
+  // TODO: Skip serializing when the signal is already aborted, the way the
+  // server entry points abort before rendering starts. Needs a decision on what
+  // to resolve with, since abort() resolves with the parts that finished.
+  if (signal !== undefined && !settled) {
+    attachAbortSignal(signal);
+  }
 }
 
 const boundCache: WeakMap<

packages/react-dom/src/__tests__/ReactDOMFizzStaticBrowser-test.js

@@ -1441,4 +1441,162 @@ describe('ReactDOMFizzStaticBrowser', () => {
 
     expect(getVisibleChildren(container)).toEqual(<div>Hi</div>);
   });
+
+  describe('abort signal lifetime', () => {
+    // Collects the lifetime signal that React bounds each abort listener with.
+    // React passes that signal to addEventListener instead of calling
+    // removeEventListener, so the runtime performs the removal and nothing here
+    // observes it directly. An aborted lifetime is what shows the listener is
+    // gone.
+    function trackAbortListenerLifetimes(signal) {
+      const lifetimes = [];
+      const add = signal.addEventListener.bind(signal);
+      signal.addEventListener = (type, listener, options) => {
+        if (type === 'abort') {
+          lifetimes.push(options.signal);
+        }
+        return add(type, listener, options);
+      };
+      return lifetimes;
+    }
+
+    it('detaches the listener when a prerender completes', async () => {
+      const controller = new AbortController();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      const result = await serverAct(() =>
+        ReactDOMFizzStatic.prerender(<div>hello world</div>, {
+          signal: controller.signal,
+        }),
+      );
+      expect(lifetimes).toHaveLength(1);
+      expect(lifetimes[0].aborted).toBe(false);
+
+      await readContent(result.prelude);
+      expect(lifetimes[0].aborted).toBe(true);
+    });
+
+    it('detaches the listener when a render completes', async () => {
+      const controller = new AbortController();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      const stream = await serverAct(() =>
+        ReactDOMFizzServer.renderToReadableStream(<div>hello world</div>, {
+          signal: controller.signal,
+        }),
+      );
+      expect(lifetimes).toHaveLength(1);
+      expect(lifetimes[0].aborted).toBe(false);
+
+      await readContent(stream);
+      expect(lifetimes[0].aborted).toBe(true);
+    });
+
+    it('detaches the listener when the signal aborts mid-render', async () => {
+      let hasLoaded = false;
+      let resolve;
+      const promise = new Promise(r => (resolve = r));
+      function Wait() {
+        if (!hasLoaded) {
+          throw promise;
+        }
+        return 'Done';
+      }
+
+      const controller = new AbortController();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      const resultPromise = ReactDOMFizzStatic.prerender(
+        <div>
+          <Suspense fallback="Loading">
+            <Wait />
+          </Suspense>
+        </div>,
+        {signal: controller.signal, onError() {}},
+      );
+      await jest.runAllTimers();
+      expect(lifetimes).toHaveLength(1);
+      expect(lifetimes[0].aborted).toBe(false);
+
+      controller.abort();
+      hasLoaded = true;
+      resolve();
+      await serverAct(() => resultPromise);
+
+      expect(lifetimes[0].aborted).toBe(true);
+    });
+
+    it('detaches the listener when the stream is cancelled', async () => {
+      let hasLoaded = false;
+      let resolve;
+      const promise = new Promise(r => (resolve = r));
+      function Wait() {
+        if (!hasLoaded) {
+          throw promise;
+        }
+        return 'Done';
+      }
+
+      const controller = new AbortController();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      const stream = await serverAct(() =>
+        ReactDOMFizzServer.renderToReadableStream(
+          <div>
+            <Suspense fallback="Loading">
+              <Wait />
+            </Suspense>
+          </div>,
+          {signal: controller.signal, onError() {}},
+        ),
+      );
+      expect(lifetimes).toHaveLength(1);
+      expect(lifetimes[0].aborted).toBe(false);
+
+      await serverAct(() => stream.cancel());
+      hasLoaded = true;
+      resolve();
+
+      expect(lifetimes[0].aborted).toBe(true);
+    });
+
+    it('detaches the listener when the shell errors', async () => {
+      // A shell error rejects before the caller ever receives a stream, so
+      // nothing consumes the request and it never closes. The listener has to
+      // come off at the fatal error itself.
+      function Boom() {
+        throw new Error('boom');
+      }
+
+      const controller = new AbortController();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      await expect(
+        serverAct(() =>
+          ReactDOMFizzServer.renderToReadableStream(<Boom />, {
+            signal: controller.signal,
+            onError() {},
+          }),
+        ),
+      ).rejects.toThrow('boom');
+
+      expect(lifetimes).toHaveLength(1);
+      expect(lifetimes[0].aborted).toBe(true);
+    });
+
+    it('attaches no listener when the signal is already aborted', async () => {
+      const controller = new AbortController();
+      controller.abort();
+      const lifetimes = trackAbortListenerLifetimes(controller.signal);
+
+      await serverAct(() =>
+        ReactDOMFizzStatic.prerender(<div>hello world</div>, {
+          signal: controller.signal,
+          onError() {},
+        }),
+      );
+
+      expect(lifetimes).toHaveLength(0);
+    });
+  });
 });

packages/react-dom/src/server/ReactDOMFizzServerBrowser.js

@@ -24,6 +24,7 @@ import {
   startFlowing,
   stopFlowing,
   abort,
+  attachAbortSignal,
 } from 'react-server/src/ReactFizzServer';
 
 import {
@@ -151,16 +152,7 @@ function renderToReadableStream(
       options ? options.formState : undefined,
     );
     if (options && options.signal) {
-      const signal = options.signal;
-      if (signal.aborted) {
-        abort(request, (signal as any).reason);
-      } else {
-        const listener = () => {
-          abort(request, (signal as any).reason);
-          signal.removeEventListener('abort', listener);
-        };
-        signal.addEventListener('abort', listener);
-      }
+      attachAbortSignal(request, options.signal);
     }
     startWork(request);
   });
@@ -221,16 +213,7 @@ function resume(
       onFatalError,
     );
     if (options && options.signal) {
-      const signal = options.signal;
-      if (signal.aborted) {
-        abort(request, (signal as any).reason);
-      } else {
-        const listener = () => {
-          abort(request, (signal as any).reason);
-          signal.removeEventListener('abort', listener);
-        };
-        signal.addEventListener('abort', listener);
-      }
+      attachAbortSignal(request, options.signal);
     }
     startWork(request);
   });

packages/react-dom/src/server/ReactDOMFizzServerBun.js

@@ -23,6 +23,7 @@ import {
   startFlowing,
   stopFlowing,
   abort,
+  attachAbortSignal,
 } from 'react-server/src/ReactFizzServer';
 
 import {
@@ -140,16 +141,7 @@ function renderToReadableStream(
       options ? options.formState : undefined,
     );
     if (options && options.signal) {
-      const signal = options.signal;
-      if (signal.aborted) {
-        abort(request, (signal as any).reason);
-      } else {
-        const listener = () => {
-          abort(request, (signal as any).reason);
-          signal.removeEventListener('abort', listener);
-        };
-        signal.addEventListener('abort', listener);
-      }
+      attachAbortSignal(request, options.signal);
     }
     startWork(request);
   });

packages/react-dom/src/server/ReactDOMFizzServerEdge.js

@@ -24,6 +24,7 @@ import {
   startFlowing,
   stopFlowing,
   abort,
+  attachAbortSignal,
 } from 'react-server/src/ReactFizzServer';
 
 import {
@@ -151,16 +152,7 @@ function renderToReadableStream(
       options ? options.formState : undefined,
     );
     if (options && options.signal) {
-      const signal = options.signal;
-      if (signal.aborted) {
-        abort(request, (signal as any).reason);
-      } else {
-        const listener = () => {
-          abort(request, (signal as any).reason);
-          signal.removeEventListener('abort', listener);
-        };
-        signal.addEventListener('abort', listener);
-      }
+      attachAbortSignal(request, options.signal);
     }
     startWork(request);
   });
@@ -221,16 +213,7 @@ function resume(
       onFatalError,
     );
     if (options && options.signal) {
-      const signal = options.signal;
-      if (signal.aborted) {
-        abort(request, (signal as any).reason);
-      } else {
-        const listener = () => {
-          abort(request, (signal as any).reason);
-          signal.removeEventListener('abort', listener);
-        };
-        signal.addEventListener('abort', listener);
-      }
+      attachAbortSignal(request, options.signal);
     }
     startWork(request);
   });