Re-use rAF, don't re-schedule it

Author: rickhanlonii

packages/react-dom-bindings/src/client/ReactDOMHostConfig.js

@@ -414,37 +414,59 @@ export const scheduleMicrotask: any =
           .catch(handleErrorInNextTick)
     : scheduleTimeout; // TODO: Determine the best fallback here.
 
-// -------------------
-//     requestAnimationFrame
-// -------------------
-type FrameAlignedTask = {
-  frameNode: any,
-  callbackNode: any,
-};
-
 // TODO: Fix these types
 export const supportsFrameAlignedTask = true;
-export function scheduleFrameAlignedTask(task: any): FrameAlignedTask {
-  // Schedule both tasks, we'll race them and use the first to fire.
-  const raf: any = localRequestAnimationFrame;
 
-  return {
-    frameNode: raf(task),
-    callbackNode: Scheduler.unstable_scheduleCallback(
+type FrameAlignedTask = {|
+  rafNode: number,
+  schedulerNode: number,
+  task: function,
+|};
+
+let currentTask: FrameAlignedTask | null = null;
+function performFrameAlignedWork() {
+  if (currentTask != null) {
+    const task = currentTask.task;
+    localCancelAnimationFrame(currentTask.id);
+    Scheduler.unstable_cancelCallback(currentTask.schedulerNode);
+    currentTask = null;
+    if (task != null) {
+      task();
+    }
+  }
+}
+
+export function scheduleFrameAlignedTask(task: any): any {
+  if (currentTask === null) {
+    const rafNode = localRequestAnimationFrame(performFrameAlignedWork);
+
+    const schedulerNode = Scheduler.unstable_scheduleCallback(
       Scheduler.unstable_NormalPriority,
+      performFrameAlignedWork,
+    );
+
+    currentTask = {
+      rafNode,
+      schedulerNode,
       task,
-    ),
-  };
-}
-export function cancelFrameAlignedTask(task: any) {
-  const caf: any = localCancelAnimationFrame;
-  if (task.frameNode != null) {
-    caf(task.frameNode);
+    };
+  } else {
+    currentTask.task = task;
+    currentTask.schedulerNode = Scheduler.unstable_scheduleCallback(
+      Scheduler.unstable_NormalPriority,
+      performFrameAlignedWork,
+    );
   }
 
-  if (task.callbackNode != null) {
-    Scheduler.unstable_cancelCallback(task.callbackNode);
-  }
+  return currentTask;
+}
+
+export function cancelFrameAlignedTask(task: FrameAlignedTask) {
+  Scheduler.unstable_cancelCallback(task.schedulerNode);
+  task.schedulerNode = null;
+  // We don't cancel the rAF in case it gets re-used later.
+  // But clear the task so if it fires and shouldn't run, it won't.
+  task.task = null;
 }
 
 function handleErrorInNextTick(error) {

packages/react-dom/src/__tests__/ReactDOMFiberAsync-test.js

@@ -618,6 +618,44 @@ describe('ReactDOMFiberAsync', () => {
       expect(counterRef.current.textContent).toBe('Count: 2');
     });
 
+    // @gate enableFrameEndScheduling
+    it('Should re-use scheduled rAF, not cancel and schedule anew', () => {
+      let setState = null;
+      let counterRef = null;
+      function Counter() {
+        const [count, setCount] = React.useState(0);
+        const ref = React.useRef();
+        setState = setCount;
+        counterRef = ref;
+        Scheduler.unstable_yieldValue('Count: ' + count);
+        return <p ref={ref}>Count: {count}</p>;
+      }
+
+      const root = ReactDOMClient.createRoot(container);
+      act(() => {
+        root.render(<Counter />);
+      });
+      expect(Scheduler).toHaveYielded(['Count: 0']);
+
+      window.event = undefined;
+      setState(1);
+      // Unknown updates should schedule a rAF.
+      expect(global.requestAnimationFrameQueue.length).toBe(1);
+      const firstRaf = global.requestAnimationFrameQueue[0];
+
+      setState(2);
+      // Default updates after unknown should re-use the scheduled rAF.
+      expect(global.requestAnimationFrameQueue.length).toBe(1);
+      const secondRaf = global.requestAnimationFrameQueue[0];
+      expect(firstRaf).toBe(secondRaf);
+
+      expect(Scheduler).toHaveYielded([]);
+      expect(counterRef.current.textContent).toBe('Count: 0');
+      global.flushRequestAnimationFrameQueue();
+      expect(Scheduler).toHaveYielded(['Count: 2']);
+      expect(counterRef.current.textContent).toBe('Count: 2');
+    });
+
     // @gate enableFrameEndScheduling
     it('Default update followed by an unknown update is batched, scheduled in a rAF', () => {
       let setState = null;

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -925,9 +925,13 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
     if (
       enableFrameEndScheduling &&
       newCallbackPriority === DefaultLane &&
+      existingCallbackNode !== null &&
+      // TODO: We can't expose the rafNode here,
+      // but how do we know the rAF is not scheduled?
+      existingCallbackNode.rafNode == null &&
       root.hasUnknownUpdates
     ) {
-      // Do nothing, we need to cancel the existing default task and schedule a rAF.
+      // Do nothing, we need to schedule a new rAF.
     } else {
       // The priority hasn't changed. We can reuse the existing task. Exit.
       return;
@@ -940,8 +944,9 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
       enableFrameEndScheduling &&
       supportsFrameAlignedTask &&
       existingCallbackNode != null &&
-      // TODO: is there a better check for callbackNode type?
-      existingCallbackNode.frameNode != null
+      // TODO: we can't expose the scheduler node here,
+      // but how do we know we need to cancel with the host config method?
+      existingCallbackNode.schedulerNode != null
     ) {
       cancelFrameAlignedTask(existingCallbackNode);
     } else {

packages/react-reconciler/src/ReactFiberWorkLoop.old.js

@@ -925,9 +925,13 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
     if (
       enableFrameEndScheduling &&
       newCallbackPriority === DefaultLane &&
+      existingCallbackNode !== null &&
+      // TODO: We can't expose the rafNode here,
+      // but how do we know the rAF is not scheduled?
+      existingCallbackNode.rafNode == null &&
       root.hasUnknownUpdates
     ) {
-      // Do nothing, we need to cancel the existing default task and schedule a rAF.
+      // Do nothing, we need to schedule a new rAF.
     } else {
       // The priority hasn't changed. We can reuse the existing task. Exit.
       return;
@@ -940,8 +944,9 @@ function ensureRootIsScheduled(root: FiberRoot, currentTime: number) {
       enableFrameEndScheduling &&
       supportsFrameAlignedTask &&
       existingCallbackNode != null &&
-      // TODO: is there a better check for callbackNode type?
-      existingCallbackNode.frameNode != null
+      // TODO: we can't expose the scheduler node here,
+      // but how do we know we need to cancel with the host config method?
+      existingCallbackNode.schedulerNode != null
     ) {
       cancelFrameAlignedTask(existingCallbackNode);
     } else {