@samitouri / QOS-React-1 / commits / 2d7f0c4259

[Flight] Insert an extra await node for awaiting on the promise returned by then callback (#33713)

When a `.then()` callback returns another Promise, there's effectively another "await" on that Promise that happens in the internals but that was not modeled. In effect the Promise returned by `.then()` is blocked on both the original Promise AND the promise returned by the callback. This models that by cloning the original node and treat that as the await on the original Promise. Then we use the existing Node to await the new Promise but its "previous" points to the clone. That way we have a forked node that awaits both. --------- Co-authored-by: Sebastian Sebbie Silbermann <sebastian.silbermann@vercel.com>

Sebastian Markbåge committed Jul 6, 2025 at 15:34 UTC 2d7f0c425902d13617f5c02c1b7c0f2cd95cd6d9

3 files changed +194 -3

packages/react-server/src/ReactFlightServer.js

+10 -2

@@ -2163,13 +2163,21 @@ function visitAsyncNode(

2163 } else {

2164 let isAwaitInUserspace = false;

2165 const fullStack = node.stack;

2166 - if (fullStack.length > 0) {

2166 + let firstFrame = 0;

2167 + while (

2168 + fullStack.length > firstFrame &&

2169 + fullStack[firstFrame][0] === 'Promise.then'

2170 + ) {

2171 + // Skip Promise.then frame itself.

2172 + firstFrame++;

2173 + }

2174 + if (fullStack.length > firstFrame) {

2175 // Check if the very first stack frame that awaited this Promise was in user space.

2176 // TODO: This doesn't take into account wrapper functions such as our fake .then()

2177 // in FlightClient which will always be considered third party awaits if you call

2178 // .then directly.

2179 const filterStackFrame = request.filterStackFrame;

2172 - const callsite = fullStack[0];

2180 + const callsite = fullStack[firstFrame];

2181 const functionName = callsite[0];

2182 const url = devirtualizeURL(callsite[1]);

2183 isAwaitInUserspace = filterStackFrame(url, functionName);

packages/react-server/src/ReactFlightServerConfigDebugNode.js

+30 -1

@@ -257,7 +257,36 @@ export function initAsyncDebugInfo(): void {

257 // the trigger that we originally stored wasn't actually the dependency.

258 // Instead, the current execution context is what ultimately unblocked it.

259 const awaited = pendingOperations.get(currentAsyncId);

260 - resolvedNode.awaited = awaited === undefined ? null : awaited;

260 + if (resolvedNode.tag === PROMISE_NODE) {

261 + // For a Promise we just override the await. We're not interested in

262 + // what created the Promise itself.

263 + resolvedNode.awaited = awaited === undefined ? null : awaited;

264 + } else {

265 + // For an await, there's really two things awaited here. It's the trigger

266 + // that .then() was called on but there seems to also be something else

267 + // in the .then() callback that blocked the returned Promise from resolving

268 + // immediately. We create a fork node which essentially represents an await

269 + // of the Promise returned from the .then() callback. That Promise was blocked

270 + // on the original awaited thing which we stored as "previous".

271 + if (awaited !== undefined) {

272 + const clonedNode: AwaitNode = {

273 + tag: AWAIT_NODE,

274 + owner: resolvedNode.owner,

275 + stack: resolvedNode.stack,

276 + start: resolvedNode.start,

277 + end: resolvedNode.end,

278 + promise: resolvedNode.promise,

279 + awaited: resolvedNode.awaited,

280 + previous: resolvedNode.previous,

281 + };

282 + // We started awaiting on the callback when the original .then() resolved.

283 + resolvedNode.start = resolvedNode.end;

284 + // It resolved now. We could use the end time of "awaited" maybe.

285 + resolvedNode.end = performance.now();

286 + resolvedNode.previous = clonedNode;

287 + resolvedNode.awaited = awaited;

288 + }

289 + }

290 }

291 }

292 },

packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js

+154

@@ -2226,4 +2226,158 @@ describe('ReactFlightAsyncDebugInfo', () => {

2226 `);

2227 }

2228 });

2229 +

2230 + it('can track IO that is chained via then(async ...)', async () => {

2231 + function getData(text) {

2232 + return delay(1).then(async () => {

2233 + return text.toUpperCase();

2234 + });

2235 + }

2236 +

2237 + async function Component({text, promise}) {

2238 + return await getData('hi, sebbie');

2239 + }

2240 +

2241 + const stream = ReactServerDOMServer.renderToPipeableStream(<Component />);

2242 +

2243 + const readable = new Stream.PassThrough(streamOptions);

2244 +

2245 + const result = ReactServerDOMClient.createFromNodeStream(readable, {

2246 + moduleMap: {},

2247 + moduleLoading: {},

2248 + });

2249 + stream.pipe(readable);

2250 +

2251 + expect(await result).toBe('HI, SEBBIE');

2252 +

2253 + await finishLoadingStream(readable);

2254 + if (

2255 + __DEV__ &&

2256 + gate(

2257 + flags =>

2258 + flags.enableComponentPerformanceTrack && flags.enableAsyncDebugInfo,

2259 + )

2260 + ) {

2261 + expect(getDebugInfo(result)).toMatchInlineSnapshot(`

2262 + [

2263 + {

2264 + "time": 0,

2265 + },

2266 + {

2267 + "env": "Server",

2268 + "key": null,

2269 + "name": "Component",

2270 + "props": {},

2271 + "stack": [

2272 + [

2273 + "Object.<anonymous>",

2274 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2275 + 2241,

2276 + 109,

2277 + 2230,

2278 + 58,

2279 + ],

2280 + ],

2281 + },

2282 + {

2283 + "time": 0,

2284 + },

2285 + {

2286 + "awaited": {

2287 + "end": 0,

2288 + "env": "Server",

2289 + "name": "delay",

2290 + "owner": {

2291 + "env": "Server",

2292 + "key": null,

2293 + "name": "Component",

2294 + "props": {},

2295 + "stack": [

2296 + [

2297 + "Object.<anonymous>",

2298 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2299 + 2241,

2300 + 109,

2301 + 2230,

2302 + 58,

2303 + ],

2304 + ],

2305 + },

2306 + "stack": [

2307 + [

2308 + "delay",

2309 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2310 + 160,

2311 + 12,

2312 + 159,

2313 + 3,

2314 + ],

2315 + [

2316 + "getData",

2317 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2318 + 2232,

2319 + 14,

2320 + 2231,

2321 + 5,

2322 + ],

2323 + [

2324 + "Component",

2325 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2326 + 2238,

2327 + 20,

2328 + 2237,

2329 + 5,

2330 + ],

2331 + ],

2332 + "start": 0,

2333 + "value": {

2334 + "value": undefined,

2335 + },

2336 + },

2337 + "env": "Server",

2338 + "owner": {

2339 + "env": "Server",

2340 + "key": null,

2341 + "name": "Component",

2342 + "props": {},

2343 + "stack": [

2344 + [

2345 + "Object.<anonymous>",

2346 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2347 + 2241,

2348 + 109,

2349 + 2230,

2350 + 58,

2351 + ],

2352 + ],

2353 + },

2354 + "stack": [

2355 + [

2356 + "getData",

2357 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2358 + 2232,

2359 + 23,

2360 + 2231,

2361 + 5,

2362 + ],

2363 + [

2364 + "Component",

2365 + "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",

2366 + 2238,

2367 + 20,

2368 + 2237,

2369 + 5,

2370 + ],

2371 + ],

2372 + },

2373 + {

2374 + "time": 0,

2375 + },

2376 + {

2377 + "time": 0,

2378 + },

2379 + ]

2380 + `);

2381 + }

2382 + });

2383 });