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

sebmarkbage · sebmarkbage · commit e4de6114b1f7 · 2022-10-23T00:54:34.000-04:00
diff --git a/.eslintrc.js b/.eslintrc.js
@@ -276,5 +276,6 @@ module.exports = {
     gate: 'readonly',
     trustedTypes: 'readonly',
     IS_REACT_ACT_ENVIRONMENT: 'readonly',
+    AsyncLocalStorage: 'readonly',
   },
 };
diff --git a/packages/react-dom-bindings/src/server/ReactDOMLegacyServerStreamConfig.js b/packages/react-dom-bindings/src/server/ReactDOMLegacyServerStreamConfig.js
@@ -21,6 +21,9 @@ export function scheduleWork(callback: () => void) {
 
 export function flushBuffered(destination: Destination) {}
 
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<any> = (null: any);
+
 export function beginWriting(destination: Destination) {}
 
 export function writeChunk(
diff --git a/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js b/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
@@ -191,6 +191,11 @@ export function scheduleWork(callback: () => void) {
 
 export function flushBuffered(destination: Destination) {}
 
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = (null: any);
+
 export function beginWriting(destination: Destination) {}
 
 export function writeChunk(destination: Destination, chunk: Chunk): void {
diff --git a/packages/react-server-dom-relay/src/ReactServerStreamConfigFB.js b/packages/react-server-dom-relay/src/ReactServerStreamConfigFB.js
@@ -23,6 +23,11 @@ export function scheduleWork(callback: () => void) {
 
 export function flushBuffered(destination: Destination) {}
 
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = (null: any);
+
 export function beginWriting(destination: Destination) {}
 
 export function writeChunk(
diff --git a/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js b/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
@@ -186,6 +186,11 @@ export function scheduleWork(callback: () => void) {
 
 export function flushBuffered(destination: Destination) {}
 
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = (null: any);
+
 export function beginWriting(destination: Destination) {}
 
 export function writeChunk(destination: Destination, chunk: Chunk): void {
diff --git a/packages/react-server/src/ReactFlightCache.js b/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?
diff --git a/packages/react-server/src/ReactFlightServer.js b/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 {
diff --git a/packages/react-server/src/ReactServerStreamConfigBrowser.js b/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;
diff --git a/packages/react-server/src/ReactServerStreamConfigNode.js b/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;
diff --git a/packages/react-server/src/forks/ReactServerStreamConfig.custom.js b/packages/react-server/src/forks/ReactServerStreamConfig.custom.js
@@ -35,6 +35,8 @@ export const writeChunk = $$$hostConfig.writeChunk;
 export const writeChunkAndReturn = $$$hostConfig.writeChunkAndReturn;
 export const completeWriting = $$$hostConfig.completeWriting;
 export const flushBuffered = $$$hostConfig.flushBuffered;
+export const supportsRequestStorage = $$$hostConfig.supportsRequestStorage;
+export const requestStorage = $$$hostConfig.requestStorage;
 export const close = $$$hostConfig.close;
 export const closeWithError = $$$hostConfig.closeWithError;
 export const stringToChunk = $$$hostConfig.stringToChunk;
diff --git a/packages/react/src/__tests__/ReactFetch-test.js b/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);
   });
 
diff --git a/scripts/error-codes/codes.json b/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."
 }
diff --git a/scripts/flow/environment.js b/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;
+  }
+}
diff --git a/scripts/rollup/bundles.js b/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] : [],
@@ -350,7 +350,7 @@ const bundles = [
     global: 'ReactDOMStatic',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['react', 'util', 'stream', 'react-dom'],
+    externals: ['react', 'util', 'async_hooks', 'stream', 'react-dom'],
   },
 
   /******* React DOM Fizz Server External Runtime *******/
@@ -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 *******/
@@ -462,7 +462,7 @@ const bundles = [
     bundleTypes: [FB_WWW_DEV, FB_WWW_PROD],
     moduleType: RENDERER,
     entry: 'react-server-dom-relay',
-    global: 'ReactFlightDOMRelayClient', // TODO: Rename to Reader
+    global: 'ReactFlightDOMRelayClient',
     minifyWithProdErrorCodes: true,
     wrapWithModuleBoundaries: false,
     externals: [
@@ -477,7 +477,7 @@ const bundles = [
     bundleTypes: [RN_FB_DEV, RN_FB_PROD],
     moduleType: RENDERER,
     entry: 'react-server-native-relay/server',
-    global: 'ReactFlightNativeRelayServer', // TODO: Rename to Writer
+    global: 'ReactFlightNativeRelayServer',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
     externals: [
@@ -486,6 +486,7 @@ const bundles = [
       'JSResourceReferenceImpl',
       'ReactNativeInternalFeatureFlags',
       'util',
+      'async_hooks',
     ],
   },
 
diff --git a/scripts/rollup/validate/eslintrc.cjs.js b/scripts/rollup/validate/eslintrc.cjs.js
@@ -38,6 +38,9 @@ module.exports = {
     Uint8Array: 'readonly',
     Promise: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // Flight Webpack
     __webpack_chunk_load__: 'readonly',
     __webpack_require__: 'readonly',
diff --git a/scripts/rollup/validate/eslintrc.cjs2015.js b/scripts/rollup/validate/eslintrc.cjs2015.js
@@ -37,6 +37,9 @@ module.exports = {
     Uint8Array: 'readonly',
     Promise: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // Flight Webpack
     __webpack_chunk_load__: 'readonly',
     __webpack_require__: 'readonly',
diff --git a/scripts/rollup/validate/eslintrc.esm.js b/scripts/rollup/validate/eslintrc.esm.js
@@ -36,6 +36,9 @@ module.exports = {
     Uint8Array: 'readonly',
     Promise: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // Flight Webpack
     __webpack_chunk_load__: 'readonly',
     __webpack_require__: 'readonly',
diff --git a/scripts/rollup/validate/eslintrc.fb.js b/scripts/rollup/validate/eslintrc.fb.js
@@ -37,6 +37,9 @@ module.exports = {
     Uint8Array: 'readonly',
     Promise: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // jest
     jest: 'readonly',
 
diff --git a/scripts/rollup/validate/eslintrc.rn.js b/scripts/rollup/validate/eslintrc.rn.js
@@ -33,6 +33,9 @@ module.exports = {
     TaskController: 'readonly',
     reportError: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // jest
     jest: 'readonly',
 
diff --git a/scripts/rollup/validate/eslintrc.umd.js b/scripts/rollup/validate/eslintrc.umd.js
@@ -42,6 +42,9 @@ module.exports = {
     Uint8Array: 'readonly',
     Promise: 'readonly',
 
+    // Temp
+    AsyncLocalStorage: 'readonly',
+
     // Flight Webpack
     __webpack_chunk_load__: 'readonly',
     __webpack_require__: 'readonly',

Original file line number	Diff line number	Diff line change
`@@ -1263,7 +1263,11 @@ function flushCompletedChunks(`
`1263`	`1263`	`}`
`1264`	`1264`
`1265`	`1265`	`export function startWork(request: Request): void {`
`1266`		`- scheduleWork(() => performWork(request));`
	`1266`	`+ if (supportsRequestStorage) {`
	`1267`	`+ scheduleWork(() => requestStorage.run(request.cache, performWork, request));`
	`1268`	`+ } else {`
	`1269`	`+ scheduleWork(() => performWork(request));`
	`1270`	`+ }`
`1267`	`1271`	`}`
`1268`	`1272`
`1269`	`1273`	`export function startFlowing(request: Request, destination: Destination): void {`
Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,9 @@ export function flushBuffered(destination: Destination) {`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`export const supportsRequestStorage = true;`
`38`		`-export const requestStorage = new AsyncLocalStorage();`
	`38`	`+export const requestStorage: AsyncLocalStorage<`
	`39`	`+ Map<Function, mixed>,`
	`40`	`+> = new AsyncLocalStorage();`
`39`	`41`
`40`	`42`	`const VIEW_SIZE = 2048;`
`41`	`43`	`let currentView = null;`