diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
index a4664613ac224..ac548ea66b81c 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -644,9 +644,16 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
// event time. The next update will compute a new event time.
currentEventTime = NoWork;
- if (didTimeout) {
- // The render task took too long to complete. Mark the current time as
- // expired to synchronously render all expired work in a single batch.
+ // Check if the render expired. If so, restart at the current time so that we
+ // can finish all the expired work in a single batch. However, we should only
+ // do this if we're starting a new tree. If we're in the middle of an existing
+ // tree, we'll continue working on that (without yielding) so that the work
+ // doesn't get dropped. If there's another expired level after that, we'll hit
+ // this path again, at which point we can batch all the subsequent levels
+ // together.
+ if (didTimeout && workInProgress === null) {
+ // Mark the current time as expired to synchronously render all expired work
+ // in a single batch.
const currentTime = requestCurrentTimeForUpdate();
markRootExpiredAtTime(root, currentTime);
// This will schedule a synchronous callback.
diff --git a/packages/react-reconciler/src/__tests__/ReactIncrementalErrorHandling-test.internal.js b/packages/react-reconciler/src/__tests__/ReactIncrementalErrorHandling-test.internal.js
index bcb74b6ea07f8..27eaf52666a7e 100644
--- a/packages/react-reconciler/src/__tests__/ReactIncrementalErrorHandling-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactIncrementalErrorHandling-test.internal.js
@@ -381,6 +381,55 @@ describe('ReactIncrementalErrorHandling', () => {
expect(ReactNoop.getChildren()).toEqual([]);
});
+ it('retries one more time if an error occurs during a render that expires midway through the tree', () => {
+ function Oops() {
+ Scheduler.unstable_yieldValue('Oops');
+ throw new Error('Oops');
+ }
+
+ function Text({text}) {
+ Scheduler.unstable_yieldValue(text);
+ return text;
+ }
+
+ function App() {
+ return (
+ <>
+
+
+
+
+
+ >
+ );
+ }
+
+ ReactNoop.render();
+
+ // Render part of the tree
+ expect(Scheduler).toFlushAndYieldThrough(['A', 'B']);
+
+ // Expire the render midway through
+ expect(() => ReactNoop.expire(10000)).toThrow('Oops');
+
+ expect(Scheduler).toHaveYielded([
+ // The render expired, but we shouldn't throw out the partial work.
+ // Finish the current level.
+ 'Oops',
+ 'C',
+ 'D',
+
+ // Since the error occured during a partially concurrent render, we should
+ // retry one more time, synchonrously.
+ 'A',
+ 'B',
+ 'Oops',
+ 'C',
+ 'D',
+ ]);
+ expect(ReactNoop.getChildren()).toEqual([]);
+ });
+
it('calls componentDidCatch multiple times for multiple errors', () => {
let id = 0;
class BadMount extends React.Component {
diff --git a/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js b/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js
index bcb9f7189c05a..43e77da4ac549 100644
--- a/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js
@@ -655,6 +655,116 @@ describe('ReactIncrementalUpdates', () => {
});
});
+ it('does not throw out partially completed tree if it expires midway through', () => {
+ function Text({text}) {
+ Scheduler.unstable_yieldValue(text);
+ return text;
+ }
+
+ function App({step}) {
+ return (
+ <>
+
+
+
+ >
+ );
+ }
+
+ function interrupt() {
+ ReactNoop.flushSync(() => {
+ ReactNoop.renderToRootWithID(null, 'other-root');
+ });
+ }
+
+ // First, as a sanity check, assert what happens when four low pri
+ // updates in separate batches are all flushed in the same callback
+ ReactNoop.act(() => {
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ // Each update flushes in a separate commit.
+ // Note: This isn't necessarily the ideal behavior. It might be better to
+ // batch all of these updates together. The fact that they don't is an
+ // implementation detail. The important part of this unit test is what
+ // happens when they expire, in which case they really should be batched to
+ // avoid blocking the main thread for a long time.
+ expect(Scheduler).toFlushAndYield([
+ // A1 already completed. Finish rendering the first level.
+ 'B1',
+ 'C1',
+ // The remaining two levels complete sequentially.
+ 'A2',
+ 'B2',
+ 'C2',
+ 'A3',
+ 'B3',
+ 'C3',
+ 'A4',
+ 'B4',
+ 'C4',
+ ]);
+ });
+
+ ReactNoop.act(() => {
+ // Now do the same thing over again, but this time, expire all the updates
+ // instead of flushing them normally.
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ interrupt();
+
+ ReactNoop.render();
+ Scheduler.unstable_advanceTime(1000);
+ expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+ // Expire all the updates
+ ReactNoop.expire(10000);
+
+ expect(Scheduler).toHaveYielded([
+ // A1 already completed. Finish rendering the first level.
+ 'B1',
+ 'C1',
+ // Then render the remaining two levels in a single batch
+ 'A4',
+ 'B4',
+ 'C4',
+ ]);
+ });
+ });
+
it('when rebasing, does not exclude updates that were already committed, regardless of priority', async () => {
const {useState, useLayoutEffect} = React;