Expose cacheSignal() alongside cache()

packages/react-noop-renderer/src/ReactNoopFlightServer.js

@@ -70,6 +70,7 @@ type Options = {
   environmentName?: string | (() => string),
   filterStackFrame?: (url: string, functionName: string) => boolean,
   identifierPrefix?: string,
+  signal?: AbortSignal,
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
 };
@@ -87,6 +88,18 @@ function render(model: ReactClientValue, options?: Options): Destination {
     __DEV__ && options ? options.environmentName : undefined,
     __DEV__ && options ? options.filterStackFrame : undefined,
   );
+  const signal = options ? options.signal : undefined;
+  if (signal) {
+    if (signal.aborted) {
+      ReactNoopFlightServer.abort(request, (signal: any).reason);
+    } else {
+      const listener = () => {
+        ReactNoopFlightServer.abort(request, (signal: any).reason);
+        signal.removeEventListener('abort', listener);
+      };
+      signal.addEventListener('abort', listener);
+    }
+  }
   ReactNoopFlightServer.startWork(request);
   ReactNoopFlightServer.startFlowing(request, destination);
   return destination;

packages/react-reconciler/src/ReactFiberAsyncDispatcher.js

@@ -25,8 +25,14 @@ function getCacheForType<T>(resourceType: () => T): T {
   return cacheForType;
 }
 
+function cacheSignal(): null | AbortSignal {
+  const cache: Cache = readContext(CacheContext);
+  return cache.controller.signal;
+}
+
 export const DefaultAsyncDispatcher: AsyncDispatcher = ({
   getCacheForType,
+  cacheSignal,
 }: any);
 
 if (__DEV__) {

packages/react-reconciler/src/ReactInternalTypes.js

@@ -459,6 +459,7 @@ export type Dispatcher = {
 
 export type AsyncDispatcher = {
   getCacheForType: <T>(resourceType: () => T) => T,
+  cacheSignal: () => null | AbortSignal,
   // DEV-only
   getOwner: () => null | Fiber | ReactComponentInfo | ComponentStackNode,
 };

packages/react-reconciler/src/__tests__/ReactCache-test.js

@@ -14,6 +14,7 @@ let React;
 let ReactNoopFlightServer;
 let ReactNoopFlightClient;
 let cache;
+let cacheSignal;
 
 describe('ReactCache', () => {
   beforeEach(() => {
@@ -25,6 +26,7 @@ describe('ReactCache', () => {
     ReactNoopFlightClient = require('react-noop-renderer/flight-client');
 
     cache = React.cache;
+    cacheSignal = React.cacheSignal;
 
     jest.resetModules();
     __unmockReact();
@@ -220,4 +222,86 @@ describe('ReactCache', () => {
     expect(cachedFoo.length).toBe(0);
     expect(cachedFoo.displayName).toBe(undefined);
   });
+
+  it('cacheSignal() returns null outside a render', async () => {
+    expect(cacheSignal()).toBe(null);
+  });
+
+  it('cacheSignal() aborts when the render finishes normally', async () => {
+    let renderedCacheSignal = null;
+
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+
+    async function Test() {
+      renderedCacheSignal = cacheSignal();
+      await promise;
+      return 'Hi';
+    }
+
+    const controller = new AbortController();
+    const errors = [];
+    const result = ReactNoopFlightServer.render(<Test />, {
+      signal: controller.signal,
+      onError(x) {
+        errors.push(x);
+      },
+    });
+    expect(errors).toEqual([]);
+    expect(renderedCacheSignal).not.toBe(controller.signal); // In the future we might make these the same
+    expect(renderedCacheSignal.aborted).toBe(false);
+    await resolve();
+    await 0;
+    await 0;
+
+    expect(await ReactNoopFlightClient.read(result)).toBe('Hi');
+
+    expect(errors).toEqual([]);
+    expect(renderedCacheSignal.aborted).toBe(true);
+    expect(renderedCacheSignal.reason.message).toContain(
+      'This render completed successfully.',
+    );
+  });
+
+  it('cacheSignal() aborts when the render is aborted', async () => {
+    let renderedCacheSignal = null;
+
+    const promise = new Promise(() => {});
+
+    async function Test() {
+      renderedCacheSignal = cacheSignal();
+      await promise;
+      return 'Hi';
+    }
+
+    const controller = new AbortController();
+    const errors = [];
+    const result = ReactNoopFlightServer.render(<Test />, {
+      signal: controller.signal,
+      onError(x) {
+        errors.push(x);
+        return 'hi';
+      },
+    });
+    expect(errors).toEqual([]);
+    expect(renderedCacheSignal).not.toBe(controller.signal); // In the future we might make these the same
+    expect(renderedCacheSignal.aborted).toBe(false);
+    const reason = new Error('Timed out');
+    controller.abort(reason);
+    expect(errors).toEqual([reason]);
+    expect(renderedCacheSignal.aborted).toBe(true);
+    expect(renderedCacheSignal.reason).toBe(reason);
+
+    let clientError = null;
+    try {
+      await ReactNoopFlightClient.read(result);
+    } catch (x) {
+      clientError = x;
+    }
+    expect(clientError).not.toBe(null);
+    if (__DEV__) {
+      expect(clientError.message).toBe('Timed out');
+    }
+    expect(clientError.digest).toBe('hi');
+  });
 });

packages/react-server/src/ReactFizzAsyncDispatcher.js

@@ -16,8 +16,13 @@ function getCacheForType<T>(resourceType: () => T): T {
   throw new Error('Not implemented.');
 }
 
+function cacheSignal(): null | AbortSignal {
+  throw new Error('Not implemented.');
+}
+
 export const DefaultAsyncDispatcher: AsyncDispatcher = ({
   getCacheForType,
+  cacheSignal,
 }: any);
 
 if (__DEV__) {

packages/react-server/src/ReactFlightServer.js

@@ -419,6 +419,7 @@ export type Request = {
   destination: null | Destination,
   bundlerConfig: ClientManifest,
   cache: Map<Function, mixed>,
+  cacheController: AbortController,
   nextChunkId: number,
   pendingChunks: number,
   hints: Hints,
@@ -529,6 +530,7 @@ function RequestInstance(
   this.destination = null;
   this.bundlerConfig = bundlerConfig;
   this.cache = new Map();
+  this.cacheController = new AbortController();
   this.nextChunkId = 0;
   this.pendingChunks = 0;
   this.hints = hints;
@@ -604,7 +606,7 @@ export function createRequest(
   model: ReactClientValue,
   bundlerConfig: ClientManifest,
   onError: void | ((error: mixed) => ?string),
-  identifierPrefix?: string,
+  identifierPrefix: void | string,
   onPostpone: void | ((reason: string) => void),
   temporaryReferences: void | TemporaryReferenceSet,
   environmentName: void | string | (() => string), // DEV-only
@@ -636,7 +638,7 @@ export function createPrerenderRequest(
   onAllReady: () => void,
   onFatalError: () => void,
   onError: void | ((error: mixed) => ?string),
-  identifierPrefix?: string,
+  identifierPrefix: void | string,
   onPostpone: void | ((reason: string) => void),
   temporaryReferences: void | TemporaryReferenceSet,
   environmentName: void | string | (() => string), // DEV-only
@@ -3369,6 +3371,13 @@ function fatalError(request: Request, error: mixed): void {
     request.status = CLOSING;
     request.fatalError = error;
   }
+  const abortReason = new Error(
+    'The render was aborted due to a fatal error.',
+    {
+      cause: error,
+    },
+  );
+  request.cacheController.abort(abortReason);
 }
 
 function emitPostponeChunk(
@@ -4840,6 +4849,12 @@ function flushCompletedChunks(
     if (enableTaint) {
       cleanupTaintQueue(request);
     }
+    if (request.status < ABORTING) {
+      const abortReason = new Error(
+        'This render completed successfully. All cacheSignals are now aborted to allow clean up of any unused resources.',
+      );
+      request.cacheController.abort(abortReason);
+    }
     request.status = CLOSED;
     close(destination);
     request.destination = null;
@@ -4921,6 +4936,7 @@ export function abort(request: Request, reason: mixed): void {
     // We define any status below OPEN as OPEN equivalent
     if (request.status <= OPEN) {
       request.status = ABORTING;
+      request.cacheController.abort(reason);
     }
     const abortableTasks = request.abortableTasks;
     if (abortableTasks.size > 0) {

packages/react-server/src/flight/ReactFlightAsyncDispatcher.js

@@ -31,6 +31,13 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
     }
     return entry;
   },
+  cacheSignal(): null | AbortSignal {
+    const request = resolveRequest();
+    if (request) {
+      return request.cacheController.signal;
+    }
+    return null;
+  },
 }: any);
 
 if (__DEV__) {

packages/react-suspense-test-utils/src/ReactSuspenseTestUtils.js

@@ -22,6 +22,9 @@ export function waitForSuspense<T>(fn: () => T): Promise<T> {
       }
       return entry;
     },
+    cacheSignal(): null {
+      return null;
+    },
     getOwner(): null {
       return null;
     },

packages/react/index.development.js

@@ -44,6 +44,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   startTransition,
   unstable_LegacyHidden,
   unstable_Activity,

packages/react/index.experimental.development.js

@@ -27,6 +27,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   startTransition,
   unstable_Activity,
   unstable_postpone,

packages/react/index.experimental.js

@@ -27,6 +27,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   startTransition,
   unstable_Activity,
   unstable_postpone,

packages/react/index.fb.js

@@ -14,6 +14,7 @@ export {
   __COMPILER_RUNTIME,
   act,
   cache,
+  cacheSignal,
   Children,
   cloneElement,
   Component,

packages/react/index.js

@@ -44,6 +44,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   startTransition,
   unstable_LegacyHidden,
   unstable_Activity,

packages/react/index.stable.development.js

@@ -27,6 +27,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   unstable_useCacheRefresh,
   startTransition,
   useId,

packages/react/index.stable.js

@@ -27,6 +27,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   unstable_useCacheRefresh,
   startTransition,
   useId,

packages/react/src/ReactCacheClient.js

@@ -8,9 +8,12 @@
  */
 
 import {disableClientCache} from 'shared/ReactFeatureFlags';
-import {cache as cacheImpl} from './ReactCacheImpl';
+import {
+  cache as cacheImpl,
+  cacheSignal as cacheSignalImpl,
+} from './ReactCacheImpl';
 
-export function noopCache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
+function noopCache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
   // On the client (i.e. not a Server Components environment) `cache` has
   // no caching behavior. We just return the function as-is.
   //
@@ -32,3 +35,11 @@ export function noopCache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
 export const cache: typeof noopCache = disableClientCache
   ? noopCache
   : cacheImpl;
+
+function noopCacheSignal(): null | AbortSignal {
+  return null;
+}
+
+export const cacheSignal: () => null | AbortSignal = disableClientCache
+  ? noopCacheSignal
+  : cacheSignalImpl;

packages/react/src/ReactCacheImpl.js

@@ -126,3 +126,15 @@ export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
     }
   };
 }
+
+export function cacheSignal(): null | AbortSignal {
+  const dispatcher = ReactSharedInternals.A;
+  if (!dispatcher) {
+    // If there is no dispatcher, then we treat this as not having an AbortSignal
+    // since in the same context, a cached function will be allowed to be called
+    // but it won't be cached. So it's neither an infinite AbortSignal nor an
+    // already resolved one.
+    return null;
+  }
+  return dispatcher.cacheSignal();
+}

packages/react/src/ReactCacheServer.js

@@ -7,4 +7,4 @@
  * @flow
  */
 
-export {cache} from './ReactCacheImpl';
+export {cache, cacheSignal} from './ReactCacheImpl';

packages/react/src/ReactClient.js

@@ -33,7 +33,7 @@ import {createContext} from './ReactContext';
 import {lazy} from './ReactLazy';
 import {forwardRef} from './ReactForwardRef';
 import {memo} from './ReactMemo';
-import {cache} from './ReactCacheClient';
+import {cache, cacheSignal} from './ReactCacheClient';
 import {postpone} from './ReactPostpone';
 import {
   getCacheForType,
@@ -83,6 +83,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   postpone as unstable_postpone,
   useCallback,
   useContext,

packages/react/src/ReactServer.experimental.development.js

@@ -35,7 +35,7 @@ import {
 import {forwardRef} from './ReactForwardRef';
 import {lazy} from './ReactLazy';
 import {memo} from './ReactMemo';
-import {cache} from './ReactCacheServer';
+import {cache, cacheSignal} from './ReactCacheServer';
 import {startTransition} from './ReactStartTransition';
 import {postpone} from './ReactPostpone';
 import {captureOwnerStack} from './ReactOwnerStack';
@@ -70,6 +70,7 @@ export {
   lazy,
   memo,
   cache,
+  cacheSignal,
   startTransition,
   getCacheForType as unstable_getCacheForType,
   postpone as unstable_postpone,