@samitouri / QOS-React / commits / cce18e3504

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

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.

Sebastian Markbåge committed Oct 23, 2022 at 01:06 UTC cce18e3504fca651051da72cf9c2d126c16eb013
21 files changed +137 -19
.eslintrc.js
+1
@@ -276,5 +276,6 @@ module.exports = {
276 gate: 'readonly',
277 trustedTypes: 'readonly',
278 IS_REACT_ACT_ENVIRONMENT: 'readonly',
279 + AsyncLocalStorage: 'readonly',
280 },
281 };
packages/react-dom-bindings/src/server/ReactDOMLegacyServerStreamConfig.js
+3
@@ -21,6 +21,9 @@ export function scheduleWork(callback: () => void) {
21
22 export function flushBuffered(destination: Destination) {}
23
24 +export const supportsRequestStorage = false;
25 +export const requestStorage: AsyncLocalStorage<any> = (null: any);
26 +
27 export function beginWriting(destination: Destination) {}
28
29 export function writeChunk(
packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
+5
@@ -191,6 +191,11 @@ export function scheduleWork(callback: () => void) {
191
192 export function flushBuffered(destination: Destination) {}
193
194 +export const supportsRequestStorage = false;
195 +export const requestStorage: AsyncLocalStorage<
196 + Map<Function, mixed>,
197 +> = (null: any);
198 +
199 export function beginWriting(destination: Destination) {}
200
201 export function writeChunk(destination: Destination, chunk: Chunk): void {
packages/react-server-dom-relay/src/ReactServerStreamConfigFB.js
+5
@@ -23,6 +23,11 @@ export function scheduleWork(callback: () => void) {
23
24 export function flushBuffered(destination: Destination) {}
25
26 +export const supportsRequestStorage = false;
27 +export const requestStorage: AsyncLocalStorage<
28 + Map<Function, mixed>,
29 +> = (null: any);
30 +
31 export function beginWriting(destination: Destination) {}
32
33 export function writeChunk(
packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
+5
@@ -186,6 +186,11 @@ export function scheduleWork(callback: () => void) {
186
187 export function flushBuffered(destination: Destination) {}
188
189 +export const supportsRequestStorage = false;
190 +export const requestStorage: AsyncLocalStorage<
191 + Map<Function, mixed>,
192 +> = (null: any);
193 +
194 export function beginWriting(destination: Destination) {}
195
196 export function writeChunk(destination: Destination, chunk: Chunk): void {
packages/react-server/src/ReactFlightCache.js
+19 -9
@@ -9,16 +9,30 @@
9
10 import type {CacheDispatcher} from 'react-reconciler/src/ReactInternalTypes';
11
12 +import {
13 + supportsRequestStorage,
14 + requestStorage,
15 +} from './ReactFlightServerConfig';
16 +
17 function createSignal(): AbortSignal {
18 return new AbortController().signal;
19 }
20
21 +function resolveCache(): Map<Function, mixed> {
22 + if (currentCache) return currentCache;
23 + if (supportsRequestStorage) {
24 + const cache = requestStorage.getStore();
25 + if (cache) return cache;
26 + }
27 + // Since we override the dispatcher all the time, we're effectively always
28 + // active and so to support cache() and fetch() outside of render, we yield
29 + // an empty Map.
30 + return new Map();
31 +}
32 +
33 export const DefaultCacheDispatcher: CacheDispatcher = {
34 getCacheSignal(): AbortSignal {
18 - if (!currentCache) {
19 - throw new Error('Reading the cache is only supported while rendering.');
20 - }
21 - let entry: AbortSignal | void = (currentCache.get(createSignal): any);
35 + let entry: AbortSignal | void = (resolveCache().get(createSignal): any);
36 if (entry === undefined) {
37 entry = createSignal();
38 // $FlowFixMe[incompatible-use] found when upgrading Flow
@@ -27,11 +41,7 @@ export const DefaultCacheDispatcher: CacheDispatcher = {
41 return entry;
42 },
43 getCacheForType<T>(resourceType: () => T): T {
30 - if (!currentCache) {
31 - throw new Error('Reading the cache is only supported while rendering.');
32 - }
33 -
34 - let entry: T | void = (currentCache.get(resourceType): any);
44 + let entry: T | void = (resolveCache().get(resourceType): any);
45 if (entry === undefined) {
46 entry = resourceType();
47 // TODO: Warn if undefined?
packages/react-server/src/ReactFlightServer.js
+17 -4
@@ -43,6 +43,8 @@ import {
43 resolveModuleMetaData,
44 getModuleKey,
45 isModuleReference,
46 + supportsRequestStorage,
47 + requestStorage,
48 } from './ReactFlightServerConfig';
49
50 import {
@@ -157,6 +159,16 @@ export function createRequest(
159 context?: Array<[string, ServerContextJSONValue]>,
160 identifierPrefix?: string,
161 ): Request {
162 + if (
163 + ReactCurrentCache.current !== null &&
164 + ReactCurrentCache.current !== DefaultCacheDispatcher
165 + ) {
166 + throw new Error(
167 + 'Currently React only supports one RSC renderer at a time.',
168 + );
169 + }
170 + ReactCurrentCache.current = DefaultCacheDispatcher;
171 +
172 const abortSet: Set<Task> = new Set();
173 const pingedTasks = [];
174 const request = {
@@ -1155,10 +1167,8 @@ function retryTask(request: Request, task: Task): void {
1167
1168 function performWork(request: Request): void {
1169 const prevDispatcher = ReactCurrentDispatcher.current;
1158 - const prevCacheDispatcher = ReactCurrentCache.current;
1170 const prevCache = getCurrentCache();
1171 ReactCurrentDispatcher.current = HooksDispatcher;
1161 - ReactCurrentCache.current = DefaultCacheDispatcher;
1172 setCurrentCache(request.cache);
1173 prepareToUseHooksForRequest(request);
1174
@@ -1177,7 +1187,6 @@ function performWork(request: Request): void {
1187 fatalError(request, error);
1188 } finally {
1189 ReactCurrentDispatcher.current = prevDispatcher;
1180 - ReactCurrentCache.current = prevCacheDispatcher;
1190 setCurrentCache(prevCache);
1191 resetHooksForRequest();
1192 }
@@ -1254,7 +1263,11 @@ function flushCompletedChunks(
1263 }
1264
1265 export function startWork(request: Request): void {
1257 - scheduleWork(() => performWork(request));
1266 + if (supportsRequestStorage) {
1267 + scheduleWork(() => requestStorage.run(request.cache, performWork, request));
1268 + } else {
1269 + scheduleWork(() => performWork(request));
1270 + }
1271 }
1272
1273 export function startFlowing(request: Request, destination: Destination): void {
packages/react-server/src/ReactFlightServerConfigStream.js
+5
@@ -72,6 +72,11 @@ import type {Chunk} from './ReactServerStreamConfig';
72
73 export type {Destination, Chunk} from './ReactServerStreamConfig';
74
75 +export {
76 + supportsRequestStorage,
77 + requestStorage,
78 +} from './ReactServerStreamConfig';
79 +
80 const stringify = JSON.stringify;
81
82 function serializeRowHeader(tag: string, id: number) {
packages/react-server/src/ReactServerStreamConfigBrowser.js
+7
@@ -21,6 +21,13 @@ export function flushBuffered(destination: Destination) {
21 // transform streams. https://github.com/whatwg/streams/issues/960
22 }
23
24 +// For now we support AsyncLocalStorage as a global for the "browser" builds
25 +// TODO: Move this to some special WinterCG build.
26 +export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
27 +export const requestStorage: AsyncLocalStorage<
28 + Map<Function, mixed>,
29 +> = supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
30 +
31 const VIEW_SIZE = 512;
32 let currentView = null;
33 let writtenBytes = 0;
packages/react-server/src/ReactServerStreamConfigNode.js
+6
@@ -9,6 +9,7 @@
9
10 import type {Writable} from 'stream';
11 import {TextEncoder} from 'util';
12 +import {AsyncLocalStorage} from 'async_hooks';
13
14 interface MightBeFlushable {
15 flush?: () => void;
@@ -33,6 +34,11 @@ export function flushBuffered(destination: Destination) {
34 }
35 }
36
37 +export const supportsRequestStorage = true;
38 +export const requestStorage: AsyncLocalStorage<
39 + Map<Function, mixed>,
40 +> = new AsyncLocalStorage();
41 +
42 const VIEW_SIZE = 2048;
43 let currentView = null;
44 let writtenBytes = 0;
packages/react-server/src/forks/ReactServerStreamConfig.custom.js
+2
@@ -35,6 +35,8 @@ export const writeChunk = $$$hostConfig.writeChunk;
35 export const writeChunkAndReturn = $$$hostConfig.writeChunkAndReturn;
36 export const completeWriting = $$$hostConfig.completeWriting;
37 export const flushBuffered = $$$hostConfig.flushBuffered;
38 +export const supportsRequestStorage = $$$hostConfig.supportsRequestStorage;
39 +export const requestStorage = $$$hostConfig.requestStorage;
40 export const close = $$$hostConfig.close;
41 export const closeWithError = $$$hostConfig.closeWithError;
42 export const stringToChunk = $$$hostConfig.stringToChunk;
packages/react/src/__tests__/ReactFetch-test.js
+20
@@ -16,6 +16,8 @@ global.TextDecoder = require('util').TextDecoder;
16 global.Headers = require('node-fetch').Headers;
17 global.Request = require('node-fetch').Request;
18 global.Response = require('node-fetch').Response;
19 +// Patch for Browser environments to be able to polyfill AsyncLocalStorage
20 +global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
21
22 let fetchCount = 0;
23 async function fetchMock(resource, options) {
@@ -76,6 +78,24 @@ describe('ReactFetch', () => {
78 expect(fetchCount).toBe(1);
79 });
80
81 + // @gate enableFetchInstrumentation && enableCache
82 + it('can dedupe fetches in micro tasks', async () => {
83 + async function getData() {
84 + const r1 = await fetch('hello');
85 + const t1 = await r1.text();
86 + const r2 = await fetch('world');
87 + const t2 = await r2.text();
88 + return t1 + ' ' + t2;
89 + }
90 + function Component() {
91 + return use(getData());
92 + }
93 + expect(await render(Component)).toMatchInlineSnapshot(
94 + `"GET hello [] GET world []"`,
95 + );
96 + expect(fetchCount).toBe(2);
97 + });
98 +
99 // @gate enableFetchInstrumentation && enableCache
100 it('can dedupe fetches using Request and not', async () => {
101 function Component() {
scripts/error-codes/codes.json
+2 -1
@@ -442,5 +442,6 @@
442 "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.",
443 "455": "This CacheSignal was requested outside React which means that it is immediately aborted.",
444 "456": "Calling Offscreen.detach before instance handle has been set.",
445 - "457": "acquireHeadResource encountered a resource type it did not expect: \"%s\". This is a bug in React."
445 + "457": "acquireHeadResource encountered a resource type it did not expect: \"%s\". This is a bug in React.",
446 + "458": "Currently React only supports one RSC renderer at a time."
447 }
scripts/flow/environment.js
+16
@@ -157,3 +157,19 @@ declare module 'pg/lib/utils' {
157 prepareValue(val: any): mixed,
158 };
159 }
160 +
161 +declare class AsyncLocalStorage<T> {
162 + disable(): void;
163 + getStore(): T | void;
164 + run(store: T, callback: (...args: any[]) => void, ...args: any[]): void;
165 + enterWith(store: T): void;
166 +}
167 +
168 +declare module 'async_hooks' {
169 + declare class AsyncLocalStorage<T> {
170 + disable(): void;
171 + getStore(): T | void;
172 + run(store: T, callback: (...args: any[]) => void, ...args: any[]): void;
173 + enterWith(store: T): void;
174 + }
175 +}
scripts/rollup/bundles.js
+6 -5
@@ -320,7 +320,7 @@ const bundles = [
320 global: 'ReactDOMServer',
321 minifyWithProdErrorCodes: false,
322 wrapWithModuleBoundaries: false,
323 - externals: ['react', 'util', 'react-dom'],
323 + externals: ['react', 'util', 'async_hooks', 'react-dom'],
324 },
325 {
326 bundleTypes: __EXPERIMENTAL__ ? [FB_WWW_DEV, FB_WWW_PROD] : [],
@@ -350,7 +350,7 @@ const bundles = [
350 global: 'ReactDOMStatic',
351 minifyWithProdErrorCodes: false,
352 wrapWithModuleBoundaries: false,
353 - externals: ['react', 'util', 'stream', 'react-dom'],
353 + externals: ['react', 'util', 'async_hooks', 'stream', 'react-dom'],
354 },
355
356 /******* React DOM Fizz Server External Runtime *******/
@@ -394,7 +394,7 @@ const bundles = [
394 global: 'ReactServerDOMServer',
395 minifyWithProdErrorCodes: false,
396 wrapWithModuleBoundaries: false,
397 - externals: ['react', 'util', 'react-dom'],
397 + externals: ['react', 'util', 'async_hooks', 'react-dom'],
398 },
399
400 /******* React Server DOM Webpack Client *******/
@@ -462,7 +462,7 @@ const bundles = [
462 bundleTypes: [FB_WWW_DEV, FB_WWW_PROD],
463 moduleType: RENDERER,
464 entry: 'react-server-dom-relay',
465 - global: 'ReactFlightDOMRelayClient', // TODO: Rename to Reader
465 + global: 'ReactFlightDOMRelayClient',
466 minifyWithProdErrorCodes: true,
467 wrapWithModuleBoundaries: false,
468 externals: [
@@ -477,7 +477,7 @@ const bundles = [
477 bundleTypes: [RN_FB_DEV, RN_FB_PROD],
478 moduleType: RENDERER,
479 entry: 'react-server-native-relay/server',
480 - global: 'ReactFlightNativeRelayServer', // TODO: Rename to Writer
480 + global: 'ReactFlightNativeRelayServer',
481 minifyWithProdErrorCodes: false,
482 wrapWithModuleBoundaries: false,
483 externals: [
@@ -486,6 +486,7 @@ const bundles = [
486 'JSResourceReferenceImpl',
487 'ReactNativeInternalFeatureFlags',
488 'util',
489 + 'async_hooks',
490 ],
491 },
492
scripts/rollup/validate/eslintrc.cjs.js
+3
@@ -38,6 +38,9 @@ module.exports = {
38 Uint8Array: 'readonly',
39 Promise: 'readonly',
40
41 + // Temp
42 + AsyncLocalStorage: 'readonly',
43 +
44 // Flight Webpack
45 __webpack_chunk_load__: 'readonly',
46 __webpack_require__: 'readonly',
scripts/rollup/validate/eslintrc.cjs2015.js
+3
@@ -37,6 +37,9 @@ module.exports = {
37 Uint8Array: 'readonly',
38 Promise: 'readonly',
39
40 + // Temp
41 + AsyncLocalStorage: 'readonly',
42 +
43 // Flight Webpack
44 __webpack_chunk_load__: 'readonly',
45 __webpack_require__: 'readonly',
scripts/rollup/validate/eslintrc.esm.js
+3
@@ -36,6 +36,9 @@ module.exports = {
36 Uint8Array: 'readonly',
37 Promise: 'readonly',
38
39 + // Temp
40 + AsyncLocalStorage: 'readonly',
41 +
42 // Flight Webpack
43 __webpack_chunk_load__: 'readonly',
44 __webpack_require__: 'readonly',
scripts/rollup/validate/eslintrc.fb.js
+3
@@ -37,6 +37,9 @@ module.exports = {
37 Uint8Array: 'readonly',
38 Promise: 'readonly',
39
40 + // Temp
41 + AsyncLocalStorage: 'readonly',
42 +
43 // jest
44 jest: 'readonly',
45
scripts/rollup/validate/eslintrc.rn.js
+3
@@ -33,6 +33,9 @@ module.exports = {
33 TaskController: 'readonly',
34 reportError: 'readonly',
35
36 + // Temp
37 + AsyncLocalStorage: 'readonly',
38 +
39 // jest
40 jest: 'readonly',
41
scripts/rollup/validate/eslintrc.umd.js
+3
@@ -42,6 +42,9 @@ module.exports = {
42 Uint8Array: 'readonly',
43 Promise: 'readonly',
44
45 + // Temp
46 + AsyncLocalStorage: 'readonly',
47 +
48 // Flight Webpack
49 __webpack_chunk_load__: 'readonly',
50 __webpack_require__: 'readonly',