From 760abec70b90c2aff8a0dd986e2aa8b87d2e865c Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Mon, 29 Mar 2021 22:37:55 -0500
Subject: [PATCH] Support nesting of startTransition and flushSync

---
 .../src/ReactFiberWorkLoop.new.js             |  4 ++
 .../src/ReactFiberWorkLoop.old.js             |  4 ++
 .../src/__tests__/ReactFlushSync-test.js      | 43 +++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
index 1cff8db1fe654..b0e54d7d74e8e 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -240,6 +240,7 @@ const ceil = Math.ceil;
 const {
   ReactCurrentDispatcher,
   ReactCurrentOwner,
+  ReactCurrentBatchConfig,
   IsSomeRendererActing,
 } = ReactSharedInternals;
 
@@ -1151,8 +1152,10 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
   const prevExecutionContext = executionContext;
   executionContext |= BatchedContext;
 
+  const prevTransition = ReactCurrentBatchConfig.transition;
   const previousPriority = getCurrentUpdatePriority();
   try {
+    ReactCurrentBatchConfig.transition = 0;
     setCurrentUpdatePriority(DiscreteEventPriority);
     if (fn) {
       return fn(a);
@@ -1161,6 +1164,7 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
     }
   } finally {
     setCurrentUpdatePriority(previousPriority);
+    ReactCurrentBatchConfig.transition = prevTransition;
     executionContext = prevExecutionContext;
     // Flush the immediate callbacks that were scheduled during this batch.
     // Note that this will happen even if batchedUpdates is higher up
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
index 98f4d119e4088..ef77be304abc9 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
@@ -240,6 +240,7 @@ const ceil = Math.ceil;
 const {
   ReactCurrentDispatcher,
   ReactCurrentOwner,
+  ReactCurrentBatchConfig,
   IsSomeRendererActing,
 } = ReactSharedInternals;
 
@@ -1151,8 +1152,10 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
   const prevExecutionContext = executionContext;
   executionContext |= BatchedContext;
 
+  const prevTransition = ReactCurrentBatchConfig.transition;
   const previousPriority = getCurrentUpdatePriority();
   try {
+    ReactCurrentBatchConfig.transition = 0;
     setCurrentUpdatePriority(DiscreteEventPriority);
     if (fn) {
       return fn(a);
@@ -1161,6 +1164,7 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
     }
   } finally {
     setCurrentUpdatePriority(previousPriority);
+    ReactCurrentBatchConfig.transition = prevTransition;
     executionContext = prevExecutionContext;
     // Flush the immediate callbacks that were scheduled during this batch.
     // Note that this will happen even if batchedUpdates is higher up
diff --git a/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js b/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
index 3c96f5ee65914..bc7499a3d43a7 100644
--- a/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
@@ -3,6 +3,7 @@ let ReactNoop;
 let Scheduler;
 let useState;
 let useEffect;
+let startTransition;
 
 describe('ReactFlushSync', () => {
   beforeEach(() => {
@@ -13,6 +14,7 @@ describe('ReactFlushSync', () => {
     Scheduler = require('scheduler');
     useState = React.useState;
     useEffect = React.useEffect;
+    startTransition = React.unstable_startTransition;
   });
 
   function Text({text}) {
@@ -54,4 +56,45 @@ describe('ReactFlushSync', () => {
     });
     expect(root).toMatchRenderedOutput('1, 1');
   });
+
+  // @gate experimental
+  test('nested with startTransition', async () => {
+    let setSyncState;
+    let setState;
+    function App() {
+      const [syncState, _setSyncState] = useState(0);
+      const [state, _setState] = useState(0);
+      setSyncState = _setSyncState;
+      setState = _setState;
+      return <Text text={`${syncState}, ${state}`} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await ReactNoop.act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['0, 0']);
+    expect(root).toMatchRenderedOutput('0, 0');
+
+    await ReactNoop.act(async () => {
+      ReactNoop.flushSync(() => {
+        startTransition(() => {
+          // This should be async even though flushSync is on the stack, because
+          // startTransition is closer.
+          setState(1);
+          ReactNoop.flushSync(() => {
+            // This should be async even though startTransition is on the stack,
+            // because flushSync is closer.
+            setSyncState(1);
+          });
+        });
+      });
+      // Only the sync update should have flushed
+      expect(Scheduler).toHaveYielded(['1, 0']);
+      expect(root).toMatchRenderedOutput('1, 0');
+    });
+    // Now the async update has flushed, too.
+    expect(Scheduler).toHaveYielded(['1, 1']);
+    expect(root).toMatchRenderedOutput('1, 1');
+  });
 });