Use AsyncLocalStorage to extend the scope of the cache to micro tasks

Author: sebmarkbage

.eslintrc.js

@@ -276,5 +276,6 @@ module.exports = {
     gate: 'readonly',
     trustedTypes: 'readonly',
     IS_REACT_ACT_ENVIRONMENT: 'readonly',
+    AsyncLocalStorage: 'readonly',
   },
 };

packages/react-server/src/ReactFlightCache.js

@@ -9,16 +9,30 @@
 
 import type {CacheDispatcher} from 'react-reconciler/src/ReactInternalTypes';
 
+import {
+  supportsRequestStorage,
+  requestStorage,
+} from './ReactFlightServerConfig';
+
 function createSignal(): AbortSignal {
   return new AbortController().signal;
 }
 
+function resolveCache(): Map<Function, mixed> {
+  if (currentCache) return currentCache;
+  if (supportsRequestStorage) {
+    const cache = requestStorage.getStore();
+    if (cache) return cache;
+  }
+  // Since we override the dispatcher all the time, we're effectively always
+  // active and so to support cache() and fetch() outside of render, we yield
+  // an empty Map.
+  return new Map();
+}
+
 export const DefaultCacheDispatcher: CacheDispatcher = {
   getCacheSignal(): AbortSignal {
-    if (!currentCache) {
-      throw new Error('Reading the cache is only supported while rendering.');
-    }
-    let entry: AbortSignal | void = (currentCache.get(createSignal): any);
+    let entry: AbortSignal | void = (resolveCache().get(createSignal): any);
     if (entry === undefined) {
       entry = createSignal();
       // $FlowFixMe[incompatible-use] found when upgrading Flow
@@ -27,11 +41,7 @@ export const DefaultCacheDispatcher: CacheDispatcher = {
     return entry;
   },
   getCacheForType<T>(resourceType: () => T): T {
-    if (!currentCache) {
-      throw new Error('Reading the cache is only supported while rendering.');
-    }
-
-    let entry: T | void = (currentCache.get(resourceType): any);
+    let entry: T | void = (resolveCache().get(resourceType): any);
     if (entry === undefined) {
       entry = resourceType();
       // TODO: Warn if undefined?

packages/react-server/src/ReactFlightServer.js

@@ -1263,7 +1263,11 @@ function flushCompletedChunks(
 }
 
 export function startWork(request: Request): void {
-  scheduleWork(() => performWork(request));
+  if (supportsRequestStorage) {
+    scheduleWork(() => requestStorage.run(request.cache, performWork, request));
+  } else {
+    scheduleWork(() => performWork(request));
+  }
 }
 
 export function startFlowing(request: Request, destination: Destination): void {

packages/react-server/src/ReactServerStreamConfigBrowser.js

@@ -24,7 +24,9 @@ export function flushBuffered(destination: Destination) {
 // For now we support AsyncLocalStorage as a global for the "browser" builds
 // TODO: Move this to some special WinterCG build.
 export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
-export const requestStorage = new AsyncLocalStorage();
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
 
 const VIEW_SIZE = 512;
 let currentView = null;

packages/react-server/src/ReactServerStreamConfigNode.js

@@ -35,7 +35,9 @@ export function flushBuffered(destination: Destination) {
 }
 
 export const supportsRequestStorage = true;
-export const requestStorage = new AsyncLocalStorage();
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = new AsyncLocalStorage();
 
 const VIEW_SIZE = 2048;
 let currentView = null;

packages/react/src/__tests__/ReactFetch-test.js

@@ -16,6 +16,8 @@ global.TextDecoder = require('util').TextDecoder;
 global.Headers = require('node-fetch').Headers;
 global.Request = require('node-fetch').Request;
 global.Response = require('node-fetch').Response;
+// Patch for Browser environments to be able to polyfill AsyncLocalStorage
+global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
 
 let fetchCount = 0;
 async function fetchMock(resource, options) {
@@ -81,14 +83,16 @@ describe('ReactFetch', () => {
     async function getData() {
       const r1 = await fetch('hello');
       const t1 = await r1.text();
-      const r2 = await fetch('hello');
+      const r2 = await fetch('world');
       const t2 = await r2.text();
       return t1 + ' ' + t2;
     }
     function Component() {
       return use(getData());
     }
-    expect(await render(Component)).toMatchInlineSnapshot(`"GET world []"`);
+    expect(await render(Component)).toMatchInlineSnapshot(
+      `"GET hello [] GET world []"`,
+    );
     expect(fetchCount).toBe(2);
   });
 

scripts/error-codes/codes.json

@@ -443,5 +443,5 @@
   "455": "This CacheSignal was requested outside React which means that it is immediately aborted.",
   "456": "Calling Offscreen.detach before instance handle has been set.",
   "457": "acquireHeadResource encountered a resource type it did not expect: \"%s\". This is a bug in React.",
-  "458": "Currently React only supports one RSC renderer at a time"
+  "458": "Currently React only supports one RSC renderer at a time."
 }

scripts/flow/environment.js

@@ -157,3 +157,19 @@ declare module 'pg/lib/utils' {
     prepareValue(val: any): mixed,
   };
 }
+
+declare class AsyncLocalStorage<T> {
+  disable(): void;
+  getStore(): T | void;
+  run(store: T, callback: (...args: any[]) => void, ...args: any[]): void;
+  enterWith(store: T): void;
+}
+
+declare module 'async_hooks' {
+  declare class AsyncLocalStorage<T> {
+    disable(): void;
+    getStore(): T | void;
+    run(store: T, callback: (...args: any[]) => void, ...args: any[]): void;
+    enterWith(store: T): void;
+  }
+}

scripts/rollup/bundles.js

@@ -320,7 +320,7 @@ const bundles = [
     global: 'ReactDOMServer',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['react', 'util', 'async-hooks', 'react-dom'],
+    externals: ['react', 'util', 'async_hooks', 'react-dom'],
   },
   {
     bundleTypes: __EXPERIMENTAL__ ? [FB_WWW_DEV, FB_WWW_PROD] : [],
@@ -394,7 +394,7 @@ const bundles = [
     global: 'ReactServerDOMServer',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['react', 'util', 'async-hooks', 'react-dom'],
+    externals: ['react', 'util', 'async_hooks', 'react-dom'],
   },
 
   /******* React Server DOM Webpack Client *******/