[Flight] Optimize Retention of Weak Promises Abit (#33736)
We don't really need to retain a reference to whatever Promise another Promise was created in. Only awaits need to retain both their trigger and their previous context.
Sebastian Markbåge committed
Jul 9, 2025 at 09:07 UTC
49ded1d12a4c2771f4fa6e4592de2f97184700a9
1 file changed
+43
-23
packages/react-server/src/ReactFlightServerConfigDebugNode.js
+43
-23
@@ -45,14 +45,10 @@ const pendingOperations: Map<number, AsyncSequence> =
45
// and so on. By putting this relationship in a WeakMap this could be done as a single pass in the VM.
46
// We don't actually ever have to read from this map since we have WeakRef reference to these Promises
47
// if they're still alive. It's also optional information so we could just expose only if GC didn't run.
48
-const awaitedPromise: WeakMap<Promise<any>, Promise<any>> = __DEV__ &&
49
-enableAsyncDebugInfo
50
- ? new WeakMap()
51
- : (null: any);
52
-const previousPromise: WeakMap<Promise<any>, Promise<any>> = __DEV__ &&
53
-enableAsyncDebugInfo
54
- ? new WeakMap()
55
- : (null: any);
48
+const awaitedPromise: WeakMap<
49
+ Promise<any>,
50
+ Promise<any> | [Promise<any>, Promise<any>],
51
+> = __DEV__ && enableAsyncDebugInfo ? new WeakMap() : (null: any);
52
53
// Keep the last resolved await as a workaround for async functions missing data.
54
let lastRanAwait: null | AwaitNode = null;
@@ -88,14 +84,6 @@ export function initAsyncDebugInfo(): void {
84
const trigger = pendingOperations.get(triggerAsyncId);
85
let node: AsyncSequence;
86
if (type === 'PROMISE') {
91
- if (trigger !== undefined && trigger.promise !== null) {
92
- const triggerPromise = trigger.promise.deref();
93
- if (triggerPromise !== undefined) {
94
- // Keep the awaited Promise alive as long as the child is alive so we can
95
- // trace its value at the end.
96
- awaitedPromise.set(resource, triggerPromise);
97
- }
98
- }
87
const currentAsyncId = executionAsyncId();
88
if (currentAsyncId !== triggerAsyncId) {
89
// When you call .then() on a native Promise, or await/Promise.all() a thenable,
@@ -104,18 +92,42 @@ export function initAsyncDebugInfo(): void {
92
// We don't track awaits on things that started outside our tracked scope.
93
return;
94
}
107
- const current = pendingOperations.get(currentAsyncId);
108
- if (current !== undefined && current.promise !== null) {
109
- const currentPromise = current.promise.deref();
110
- if (currentPromise !== undefined) {
111
- // Keep the previous Promise alive as long as the child is alive so we can
95
+ let retain: null | Promise<any> | [Promise<any>, Promise<any>] =
96
+ null;
97
+ const triggerPromiseRef = trigger.promise;
98
+ if (triggerPromiseRef !== null) {
99
+ const triggerPromise = triggerPromiseRef.deref();
100
+ if (triggerPromise !== undefined) {
101
+ // Keep the awaited Promise alive as long as the child is alive so we can
102
// trace its value at the end.
113
- previousPromise.set(resource, currentPromise);
103
+ retain = triggerPromise;
104
}
105
}
106
+
107
+ const current = pendingOperations.get(currentAsyncId);
108
+ if (current !== undefined) {
109
+ const currentPromiseRef = current.promise;
110
+ if (currentPromiseRef !== null) {
111
+ const currentPromise = currentPromiseRef.deref();
112
+ if (currentPromise !== undefined) {
113
+ // Keep the previous Promise alive as long as the child is alive so we can
114
+ // trace its value at the end.
115
+ if (retain === null) {
116
+ retain = currentPromise;
117
+ } else {
118
+ retain = [(retain: any), currentPromise];
119
+ }
120
+ }
121
+ }
122
+ }
123
+
124
+ if (retain !== null) {
125
+ awaitedPromise.set(resource, retain);
126
+ }
127
// If the thing we're waiting on is another Await we still track that sequence
128
// so that we can later pick the best stack trace in user space.
129
let stack = null;
130
+ let promiseRef: WeakRef<Promise<any>>;
131
if (
132
trigger.stack !== null &&
133
(trigger.tag === AWAIT_NODE ||
@@ -124,7 +136,15 @@ export function initAsyncDebugInfo(): void {
136
// We already had a stack for an await. In a chain of awaits we'll only need one good stack.
137
// We mark it with an empty stack to signal to any await on this await that we have a stack.
138
stack = emptyStack;
139
+ if (resource._debugInfo !== undefined) {
140
+ // We may need to forward this debug info at the end so we need to retain this promise.
141
+ promiseRef = new WeakRef((resource: Promise<any>));
142
+ } else {
143
+ // Otherwise, we can just refer to the inner one since that's the one we'll log anyway.
144
+ promiseRef = trigger.promise;
145
+ }
146
} else {
147
+ promiseRef = new WeakRef((resource: Promise<any>));
148
const request = resolveRequest();
149
if (request === null) {
150
// We don't collect stacks for awaits that weren't in the scope of a specific render.
@@ -144,7 +164,7 @@ export function initAsyncDebugInfo(): void {
164
stack: stack,
165
start: performance.now(),
166
end: -1.1, // set when resolved.
147
- promise: new WeakRef((resource: Promise<any>)),
167
+ promise: promiseRef,
168
awaited: trigger, // The thing we're awaiting on. Might get overrriden when we resolve.
169
previous: current === undefined ? null : current, // The path that led us here.
170
}: UnresolvedAwaitNode);