Bugfix: Expiring a partially completed tree

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.

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 (
+        <>
+          <Text text="A" />
+          <Text text="B" />
+          <Oops />
+          <Text text="C" />
+          <Text text="D" />
+        </>
+      );
+    }
+
+    ReactNoop.render(<App />);
+
+    // 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 {

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 (
+        <>
+          <Text text={`A${step}`} />
+          <Text text={`B${step}`} />
+          <Text text={`C${step}`} />
+        </>
+      );
+    }
+
+    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(<App step={1} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={2} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={3} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={4} />);
+      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(<App step={1} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={2} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={3} />);
+      Scheduler.unstable_advanceTime(1000);
+      expect(Scheduler).toFlushAndYieldThrough(['A1']);
+
+      interrupt();
+
+      ReactNoop.render(<App step={4} />);
+      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;