[Flight] Use AsyncLocalStorage to extend the scope of the cache to micro tasks (#25542)

sebmarkbage · web-flow · commit cce18e3504fc · 2022-10-23T01:06:58.000-04:00
This extends the scope of the cache and fetch instrumentation using
AsyncLocalStorage for microtasks. This is an intermediate step. It sets
up the dispatcher only once. This is unique to RSC because it uses the
react.shared-subset module for its shared state.

Ideally we should support multiple renderers. We should also have this
take over from an outer SSR's instrumented fetch. We should also be able
to have a fallback to global state per request where AsyncLocalStorage
doesn't exist and then the whole client-side solutions. I'm still
figuring out the right wiring for that so this is a temporary hack.
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
@@ -43,6 +43,8 @@ import {
   resolveModuleMetaData,
   getModuleKey,
   isModuleReference,
+  supportsRequestStorage,
+  requestStorage,
 } from './ReactFlightServerConfig';
 
 import {
@@ -157,6 +159,16 @@ export function createRequest(
   context?: Array<[string, ServerContextJSONValue]>,
   identifierPrefix?: string,
 ): Request {
+  if (
+    ReactCurrentCache.current !== null &&
+    ReactCurrentCache.current !== DefaultCacheDispatcher
+  ) {
+    throw new Error(
+      'Currently React only supports one RSC renderer at a time.',
+    );
+  }
+  ReactCurrentCache.current = DefaultCacheDispatcher;
+
   const abortSet: Set<Task> = new Set();
   const pingedTasks = [];
   const request = {
@@ -1155,10 +1167,8 @@ function retryTask(request: Request, task: Task): void {
 
 function performWork(request: Request): void {
   const prevDispatcher = ReactCurrentDispatcher.current;
-  const prevCacheDispatcher = ReactCurrentCache.current;
   const prevCache = getCurrentCache();
   ReactCurrentDispatcher.current = HooksDispatcher;
-  ReactCurrentCache.current = DefaultCacheDispatcher;
   setCurrentCache(request.cache);
   prepareToUseHooksForRequest(request);
 
@@ -1177,7 +1187,6 @@ function performWork(request: Request): void {
     fatalError(request, error);
   } finally {
     ReactCurrentDispatcher.current = prevDispatcher;
-    ReactCurrentCache.current = prevCacheDispatcher;
     setCurrentCache(prevCache);
     resetHooksForRequest();
   }
@@ -1254,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/ReactFlightServerConfigStream.js b/packages/react-server/src/ReactFlightServerConfigStream.js
@@ -72,6 +72,11 @@ import type {Chunk} from './ReactServerStreamConfig';
 
 export type {Destination, Chunk} from './ReactServerStreamConfig';
 
+export {
+  supportsRequestStorage,
+  requestStorage,
+} from './ReactServerStreamConfig';
+
 const stringify = JSON.stringify;
 
 function serializeRowHeader(tag: string, id: number) {
diff --git a/packages/react-server/src/ReactServerStreamConfigBrowser.js b/packages/react-server/src/ReactServerStreamConfigBrowser.js
@@ -21,6 +21,13 @@ export function flushBuffered(destination: Destination) {
   // transform streams. https://github.com/whatwg/streams/issues/960
 }
 
+// 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: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
+
 const VIEW_SIZE = 512;
 let currentView = null;
 let writtenBytes = 0;
diff --git a/packages/react-server/src/ReactServerStreamConfigNode.js b/packages/react-server/src/ReactServerStreamConfigNode.js
@@ -9,6 +9,7 @@
 
 import type {Writable} from 'stream';
 import {TextEncoder} from 'util';
+import {AsyncLocalStorage} from 'async_hooks';
 
 interface MightBeFlushable {
   flush?: () => void;
@@ -33,6 +34,11 @@ export function flushBuffered(destination: Destination) {
   }
 }
 
+export const supportsRequestStorage = true;
+export const requestStorage: AsyncLocalStorage<
+  Map<Function, mixed>,
+> = new AsyncLocalStorage();
+
 const VIEW_SIZE = 2048;
 let currentView = null;
 let writtenBytes = 0;
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) {
@@ -76,6 +78,24 @@ describe('ReactFetch', () => {
     expect(fetchCount).toBe(1);
   });
 
+  // @gate enableFetchInstrumentation && enableCache
+  it('can dedupe fetches in micro tasks', async () => {
+    async function getData() {
+      const r1 = await fetch('hello');
+      const t1 = await r1.text();
+      const r2 = await fetch('world');
+      const t2 = await r2.text();
+      return t1 + ' ' + t2;
+    }
+    function Component() {
+      return use(getData());
+    }
+    expect(await render(Component)).toMatchInlineSnapshot(
+      `"GET hello [] GET world []"`,
+    );
+    expect(fetchCount).toBe(2);
+  });
+
   // @gate enableFetchInstrumentation && enableCache
   it('can dedupe fetches using Request and not', async () => {
     function Component() {
diff --git a/scripts/error-codes/codes.json b/scripts/error-codes/codes.json
@@ -442,5 +442,6 @@
   "454": "React expected a <body> element (document.body) to exist in the Document but one was not found. React never removes the body for any Document it renders into so the cause is likely in some other script running on this page.",
   "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."
+  "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."
 }
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', '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', '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',
