Move suspended render logic to ensureRootIsScheduled (#26328)
When the work loop is suspended, we shouldn't schedule a new render task until the promise has resolved. When I originally implemented this, I wasn't sure where to put this logic — `ensureRootIsScheduled` is the more natural place for it, but that's also a really hot path, so I chose to do it elsewhere, and left a TODO to reconsider later. Now it's later. I'm working on a refactor to move the `ensureRootIsScheduled` call to always happen in a microtask, so that if there are multiple updates/pings in a single event, they get batched into a single operation. Which means I can put the logic in that function where it belongs.
Andrew Clark committed
Mar 6, 2023 at 12:27 UTC
6e1756a5a9bb0d95ec17983f7dc38cd2c2663b39
2 files changed
+68
-18
packages/react-reconciler/src/ReactFiberWorkLoop.js
+27
-18
@@ -322,7 +322,7 @@ const SuspendedOnError: SuspendedReason = 1;
322
const SuspendedOnData: SuspendedReason = 2;
323
const SuspendedOnImmediate: SuspendedReason = 3;
324
const SuspendedOnDeprecatedThrowPromise: SuspendedReason = 4;
325
-const SuspendedAndReadyToUnwind: SuspendedReason = 5;
325
+const SuspendedAndReadyToContinue: SuspendedReason = 5;
326
const SuspendedOnHydration: SuspendedReason = 6;
327
328
// When this is true, the work-in-progress fiber just suspended (or errored) and
@@ -892,6 +892,18 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
892
return;
893
}
894
895
+ // If this root is currently suspended and waiting for data to resolve, don't
896
+ // schedule a task to render it. We'll either wait for a ping, or wait to
897
+ // receive an update.
898
+ if (
899
+ workInProgressSuspendedReason === SuspendedOnData &&
900
+ workInProgressRoot === root
901
+ ) {
902
+ root.callbackPriority = NoLane;
903
+ root.callbackNode = null;
904
+ return;
905
+ }
906
+
907
// We use the highest priority lane to represent the priority of the callback.
908
const newCallbackPriority = getHighestPriorityLane(nextLanes);
909
@@ -1153,20 +1165,6 @@ function performConcurrentWorkOnRoot(
1165
if (root.callbackNode === originalCallbackNode) {
1166
// The task node scheduled for this root is the same one that's
1167
// currently executed. Need to return a continuation.
1156
- if (
1157
- workInProgressSuspendedReason === SuspendedOnData &&
1158
- workInProgressRoot === root
1159
- ) {
1160
- // Special case: The work loop is currently suspended and waiting for
1161
- // data to resolve. Unschedule the current task.
1162
- //
1163
- // TODO: The factoring is a little weird. Arguably this should be checked
1164
- // in ensureRootIsScheduled instead. I went back and forth, not totally
1165
- // sure yet.
1166
- root.callbackPriority = NoLane;
1167
- root.callbackNode = null;
1168
- return null;
1169
- }
1168
return performConcurrentWorkOnRoot.bind(null, root);
1169
}
1170
return null;
@@ -1858,7 +1856,7 @@ function handleThrow(root: FiberRoot, thrownValue: any): void {
1856
case SuspendedOnData:
1857
case SuspendedOnImmediate:
1858
case SuspendedOnDeprecatedThrowPromise:
1861
- case SuspendedAndReadyToUnwind: {
1859
+ case SuspendedAndReadyToContinue: {
1860
const wakeable: Wakeable = (thrownValue: any);
1861
markComponentSuspended(
1862
erroredWork,
@@ -2216,6 +2214,17 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
2214
// `status` field, but if the promise already has a status, we won't
2215
// have added a listener until right here.
2216
const onResolution = () => {
2217
+ // Check if the root is still suspended on this promise.
2218
+ if (
2219
+ workInProgressSuspendedReason === SuspendedOnData &&
2220
+ workInProgressRoot === root
2221
+ ) {
2222
+ // Mark the root as ready to continue rendering.
2223
+ workInProgressSuspendedReason = SuspendedAndReadyToContinue;
2224
+ }
2225
+ // Ensure the root is scheduled. We should do this even if we're
2226
+ // currently working on a different root, so that we resume
2227
+ // rendering later.
2228
ensureRootIsScheduled(root, now());
2229
};
2230
thenable.then(onResolution, onResolution);
@@ -2225,10 +2234,10 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
2234
// If this fiber just suspended, it's possible the data is already
2235
// cached. Yield to the main thread to give it a chance to ping. If
2236
// it does, we can retry immediately without unwinding the stack.
2228
- workInProgressSuspendedReason = SuspendedAndReadyToUnwind;
2237
+ workInProgressSuspendedReason = SuspendedAndReadyToContinue;
2238
break outer;
2239
}
2231
- case SuspendedAndReadyToUnwind: {
2240
+ case SuspendedAndReadyToContinue: {
2241
const thenable: Thenable<mixed> = (thrownValue: any);
2242
if (isThenableResolved(thenable)) {
2243
// The data resolved. Try rendering the component again.
packages/react-reconciler/src/__tests__/ReactThenable-test.js
+41
@@ -650,6 +650,47 @@ describe('ReactThenable', () => {
650
assertLog(['Something different']);
651
});
652
653
+ // @gate enableUseHook
654
+ test('when waiting for data to resolve, an update on a different root does not cause work to be dropped', async () => {
655
+ const getCachedAsyncText = cache(getAsyncText);
656
+
657
+ function App() {
658
+ return <Text text={use(getCachedAsyncText('Hi'))} />;
659
+ }
660
+
661
+ const root1 = ReactNoop.createRoot();
662
+ await act(async () => {
663
+ root1.render(<Suspense fallback={<Text text="Loading..." />} />);
664
+ });
665
+
666
+ // Start a transition on one root. It will suspend.
667
+ await act(async () => {
668
+ startTransition(() => {
669
+ root1.render(
670
+ <Suspense fallback={<Text text="Loading..." />}>
671
+ <App />
672
+ </Suspense>,
673
+ );
674
+ });
675
+ });
676
+ assertLog(['Async text requested [Hi]']);
677
+
678
+ // While we're waiting for the first root's data to resolve, a second
679
+ // root renders.
680
+ const root2 = ReactNoop.createRoot();
681
+ await act(async () => {
682
+ root2.render('Do re mi');
683
+ });
684
+ expect(root2).toMatchRenderedOutput('Do re mi');
685
+
686
+ // Once the first root's data is ready, we should finish its transition.
687
+ await act(async () => {
688
+ await resolveTextRequests('Hi');
689
+ });
690
+ assertLog(['Hi']);
691
+ expect(root1).toMatchRenderedOutput('Hi');
692
+ });
693
+
694
// @gate enableUseHook
695
test('while suspended, hooks cannot be called (i.e. current dispatcher is unset correctly)', async () => {
696
function App() {