When a root expires, flush all expired work in a single batch

Instead of flushing each level one at a time.

packages/react-noop-renderer/src/createReactNoop.js

@@ -40,6 +40,7 @@ if (__DEV__) {
 
 function createReactNoop(reconciler: Function, useMutation: boolean) {
   let scheduledCallback = null;
+  let scheduledCallbackTimeout = -1;
   let instanceCounter = 0;
   let hostDiffCounter = 0;
   let hostUpdateCounter = 0;
@@ -251,14 +252,27 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       return inst;
     },
 
-    scheduleDeferredCallback(callback) {
+    scheduleDeferredCallback(callback, options) {
       if (scheduledCallback) {
         throw new Error(
           'Scheduling a callback twice is excessive. Instead, keep track of ' +
             'whether the callback has already been scheduled.',
         );
       }
       scheduledCallback = callback;
+      if (
+        typeof options === 'object' &&
+        options !== null &&
+        typeof options.timeout === 'number'
+      ) {
+        const newTimeout = options.timeout;
+        if (
+          scheduledCallbackTimeout === -1 ||
+          scheduledCallbackTimeout > newTimeout
+        ) {
+          scheduledCallbackTimeout = newTimeout;
+        }
+      }
       return 0;
     },
 
@@ -267,6 +281,7 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
         throw new Error('No callback is scheduled.');
       }
       scheduledCallback = null;
+      scheduledCallbackTimeout = -1;
     },
 
     scheduleTimeout: setTimeout,
@@ -409,10 +424,9 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
           didStop = true;
           return 0;
         },
-        // React's scheduler has its own way of keeping track of expired
-        // work and doesn't read this, so don't bother setting it to the
-        // correct value.
-        didTimeout: false,
+        didTimeout:
+          scheduledCallbackTimeout !== -1 &&
+          elapsedTimeInMs > scheduledCallbackTimeout,
       });
 
       if (yieldedValues !== null) {

packages/react-reconciler/src/ReactFiberPendingPriority.js

@@ -242,6 +242,17 @@ export function findEarliestOutstandingPriorityLevel(
   return earliestExpirationTime;
 }
 
+export function didExpireAtExpirationTime(
+  root: FiberRoot,
+  currentTime: ExpirationTime,
+): void {
+  const expirationTime = root.expirationTime;
+  if (expirationTime !== NoWork && currentTime >= expirationTime) {
+    // The root has expired. Flush all work up to the current time.
+    root.nextExpirationTimeToWorkOn = currentTime;
+  }
+}
+
 function findNextExpirationTimeToWorkOn(completedExpirationTime, root) {
   const earliestSuspendedTime = root.earliestSuspendedTime;
   const latestSuspendedTime = root.latestSuspendedTime;

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -80,6 +80,7 @@ import {
   hasLowerPriorityWork,
   isPriorityLevelSuspended,
   findEarliestOutstandingPriorityLevel,
+  didExpireAtExpirationTime,
 } from './ReactFiberPendingPriority';
 import {
   recordEffect,
@@ -2109,6 +2110,22 @@ function findHighestPriorityRoot() {
 }
 
 function performAsyncWork(dl) {
+  if (dl.didTimeout) {
+    // The callback timed out. That means at least one update has expired.
+    // Iterate through the root schedule. If they contain expired work, set
+    // the next render expiration time to the current time. This has the effect
+    // of flushing all expired work in a single batch, instead of flushing each
+    // level one at a time.
+    if (firstScheduledRoot !== null) {
+      recomputeCurrentRendererTime();
+      let root: FiberRoot = firstScheduledRoot;
+      do {
+        didExpireAtExpirationTime(root, currentRendererTime);
+        // The root schedule is circular, so this is never null.
+        root = (root.nextScheduledRoot: any);
+      } while (root !== firstScheduledRoot);
+    }
+  }
   performWork(NoWork, dl);
 }
 

packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js

@@ -453,4 +453,153 @@ describe('ReactIncrementalUpdates', () => {
     });
     expect(ReactNoop.getChildren()).toEqual([span('derived state')]);
   });
+
+  it('flushes all expired updates in a single batch', () => {
+    class Foo extends React.Component {
+      componentDidUpdate() {
+        ReactNoop.yield('Commit: ' + this.props.prop);
+      }
+      componentDidMount() {
+        ReactNoop.yield('Commit: ' + this.props.prop);
+      }
+      render() {
+        ReactNoop.yield('Render: ' + this.props.prop);
+        return <span prop={this.props.prop} />;
+      }
+    }
+
+    // First, as a sanity check, assert what happens when four low pri
+    // updates in separate batches are all flushed in the same callback
+    ReactNoop.render(<Foo prop="" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="he" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="hell" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="hello" />);
+
+    // There should be a separate render and commit for each update
+    expect(ReactNoop.flush()).toEqual([
+      'Render: ',
+      'Commit: ',
+      'Render: he',
+      'Commit: he',
+      'Render: hell',
+      'Commit: hell',
+      'Render: hello',
+      'Commit: hello',
+    ]);
+    expect(ReactNoop.getChildren()).toEqual([span('hello')]);
+
+    // Now do the same thing, except this time expire all the updates
+    // before flushing them.
+    ReactNoop.render(<Foo prop="" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="go" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="good" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo prop="goodbye" />);
+
+    ReactNoop.advanceTime(10000);
+    jest.advanceTimersByTime(10000);
+
+    // All the updates should render and commit in a single batch.
+    expect(ReactNoop.flush()).toEqual(['Render: goodbye', 'Commit: goodbye']);
+    expect(ReactNoop.getChildren()).toEqual([span('goodbye')]);
+  });
+
+  it('flushes all expired updates in a single batch across multiple roots', () => {
+    // Same as previous test, but with two roots.
+    class Foo extends React.Component {
+      componentDidUpdate() {
+        ReactNoop.yield('Commit: ' + this.props.prop);
+      }
+      componentDidMount() {
+        ReactNoop.yield('Commit: ' + this.props.prop);
+      }
+      render() {
+        ReactNoop.yield('Render: ' + this.props.prop);
+        return <span prop={this.props.prop} />;
+      }
+    }
+
+    // First, as a sanity check, assert what happens when four low pri
+    // updates in separate batches are all flushed in the same callback
+    ReactNoop.renderToRootWithID(<Foo prop="" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="" />, 'b');
+
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="he" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="he" />, 'b');
+
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="hell" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="hell" />, 'b');
+
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="hello" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="hello" />, 'b');
+
+    // There should be a separate render and commit for each update
+    expect(ReactNoop.flush()).toEqual([
+      'Render: ',
+      'Commit: ',
+      'Render: ',
+      'Commit: ',
+      'Render: he',
+      'Commit: he',
+      'Render: he',
+      'Commit: he',
+      'Render: hell',
+      'Commit: hell',
+      'Render: hell',
+      'Commit: hell',
+      'Render: hello',
+      'Commit: hello',
+      'Render: hello',
+      'Commit: hello',
+    ]);
+    expect(ReactNoop.getChildren('a')).toEqual([span('hello')]);
+    expect(ReactNoop.getChildren('b')).toEqual([span('hello')]);
+
+    // Now do the same thing, except this time expire all the updates
+    // before flushing them.
+    ReactNoop.renderToRootWithID(<Foo prop="" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="" />, 'b');
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="go" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="go" />, 'b');
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="good" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="good" />, 'b');
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.renderToRootWithID(<Foo prop="goodbye" />, 'a');
+    ReactNoop.renderToRootWithID(<Foo prop="goodbye" />, 'b');
+
+    ReactNoop.advanceTime(10000);
+    jest.advanceTimersByTime(10000);
+
+    // All the updates should render and commit in a single batch.
+    expect(ReactNoop.flush()).toEqual([
+      'Render: goodbye',
+      'Commit: goodbye',
+      'Render: goodbye',
+      'Commit: goodbye',
+    ]);
+    expect(ReactNoop.getChildren('a')).toEqual([span('goodbye')]);
+    expect(ReactNoop.getChildren('b')).toEqual([span('goodbye')]);
+  });
 });

packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js

@@ -782,6 +782,50 @@ describe('ReactSuspense', () => {
     expect(ReactNoop.getChildren()).toEqual([div(span('Async'))]);
   });
 
+  it('flushes all expired updates in a single batch', async () => {
+    class Foo extends React.Component {
+      componentDidUpdate() {
+        ReactNoop.yield('Commit: ' + this.props.text);
+      }
+      componentDidMount() {
+        ReactNoop.yield('Commit: ' + this.props.text);
+      }
+      render() {
+        return (
+          <Placeholder fallback={<Text text="Loading..." />}>
+            <AsyncText ms={20000} text={this.props.text} />
+          </Placeholder>
+        );
+      }
+    }
+
+    ReactNoop.render(<Foo text="" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo text="go" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo text="good" />);
+    ReactNoop.expire(1000);
+    jest.advanceTimersByTime(1000);
+    ReactNoop.render(<Foo text="goodbye" />);
+
+    ReactNoop.advanceTime(10000);
+    jest.advanceTimersByTime(10000);
+
+    expect(ReactNoop.flush()).toEqual([
+      'Suspend! [goodbye]',
+      'Loading...',
+      'Commit: goodbye',
+    ]);
+    expect(ReactNoop.getChildren()).toEqual([span('Loading...')]);
+
+    ReactNoop.advanceTime(20000);
+    await advanceTimers(20000);
+    expect(ReactNoop.clearYields()).toEqual(['Promise resolved [goodbye]']);
+    expect(ReactNoop.getChildren()).toEqual([span('Loading...')]);
+  });
+
   describe('a Delay component', () => {
     function Never() {
       // Throws a promise that resolves after some arbitrarily large