From 098600c42addca272bcb833bf55f53a65126e8aa Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Mon, 3 May 2021 13:34:21 -0500
Subject: [PATCH] Re-land "Fix: flushSync changes priority inside effect
 (#21122)"

This re-lands commit 0e3c7e1d62efb6238b69e5295d45b9bd2dcf9181.
---
 .../src/ReactFiberWorkLoop.new.js             | 22 +++----
 .../src/ReactFiberWorkLoop.old.js             | 22 +++----
 .../src/__tests__/ReactFlushSync-test.js      | 64 +++++++++++++++++++
 3 files changed, 86 insertions(+), 22 deletions(-)
 create mode 100644 packages/react-reconciler/src/__tests__/ReactFlushSync-test.js

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
index 68a811e3fe088..7c659b7775a1e 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -1142,16 +1142,6 @@ export function unbatchedUpdates<A, R>(fn: (a: A) => R, a: A): R {
 
 export function flushSync<A, R>(fn: A => R, a: A): R {
   const prevExecutionContext = executionContext;
-  if ((prevExecutionContext & (RenderContext | CommitContext)) !== NoContext) {
-    if (__DEV__) {
-      console.error(
-        'flushSync was called from inside a lifecycle method. React cannot ' +
-          'flush when React is already rendering. Consider moving this call to ' +
-          'a scheduler task or micro task.',
-      );
-    }
-    return fn(a);
-  }
   executionContext |= BatchedContext;
 
   const previousPriority = getCurrentUpdatePriority();
@@ -1168,7 +1158,17 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
     // Flush the immediate callbacks that were scheduled during this batch.
     // Note that this will happen even if batchedUpdates is higher up
     // the stack.
-    flushSyncCallbacks();
+    if ((executionContext & (RenderContext | CommitContext)) === NoContext) {
+      flushSyncCallbacks();
+    } else {
+      if (__DEV__) {
+        console.error(
+          'flushSync was called from inside a lifecycle method. React cannot ' +
+            'flush when React is already rendering. Consider moving this call to ' +
+            'a scheduler task or micro task.',
+        );
+      }
+    }
   }
 }
 
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
index 729e8e6717d05..530646f8a4d29 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
@@ -1142,16 +1142,6 @@ export function unbatchedUpdates<A, R>(fn: (a: A) => R, a: A): R {
 
 export function flushSync<A, R>(fn: A => R, a: A): R {
   const prevExecutionContext = executionContext;
-  if ((prevExecutionContext & (RenderContext | CommitContext)) !== NoContext) {
-    if (__DEV__) {
-      console.error(
-        'flushSync was called from inside a lifecycle method. React cannot ' +
-          'flush when React is already rendering. Consider moving this call to ' +
-          'a scheduler task or micro task.',
-      );
-    }
-    return fn(a);
-  }
   executionContext |= BatchedContext;
 
   const previousPriority = getCurrentUpdatePriority();
@@ -1168,7 +1158,17 @@ export function flushSync<A, R>(fn: A => R, a: A): R {
     // Flush the immediate callbacks that were scheduled during this batch.
     // Note that this will happen even if batchedUpdates is higher up
     // the stack.
-    flushSyncCallbacks();
+    if ((executionContext & (RenderContext | CommitContext)) === NoContext) {
+      flushSyncCallbacks();
+    } else {
+      if (__DEV__) {
+        console.error(
+          'flushSync was called from inside a lifecycle method. React cannot ' +
+            'flush when React is already rendering. Consider moving this call to ' +
+            'a scheduler task or micro task.',
+        );
+      }
+    }
   }
 }
 
diff --git a/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js b/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
new file mode 100644
index 0000000000000..f763e44705aaa
--- /dev/null
+++ b/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
@@ -0,0 +1,64 @@
+let React;
+let ReactNoop;
+let Scheduler;
+let useState;
+let useEffect;
+
+describe('ReactFlushSync', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
+    useState = React.useState;
+    useEffect = React.useEffect;
+  });
+
+  function Text({text}) {
+    Scheduler.unstable_yieldValue(text);
+    return text;
+  }
+
+  // @gate experimental || !enableSyncDefaultUpdates
+  test('changes priority of updates in useEffect', async () => {
+    function App() {
+      const [syncState, setSyncState] = useState(0);
+      const [state, setState] = useState(0);
+      useEffect(() => {
+        if (syncState !== 1) {
+          setState(1);
+          ReactNoop.flushSync(() => setSyncState(1));
+        }
+      }, [syncState, state]);
+      return <Text text={`${syncState}, ${state}`} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await ReactNoop.act(async () => {
+      if (gate(flags => flags.enableSyncDefaultUpdates)) {
+        React.unstable_startTransition(() => {
+          root.render(<App />);
+        });
+      } else {
+        root.render(<App />);
+      }
+      // This will yield right before the passive effect fires
+      expect(Scheduler).toFlushUntilNextPaint(['0, 0']);
+
+      // The passive effect will schedule a sync update and a normal update.
+      // They should commit in two separate batches. First the sync one.
+      expect(() => {
+        expect(Scheduler).toFlushUntilNextPaint(['1, 0']);
+      }).toErrorDev('flushSync was called from inside a lifecycle method');
+
+      // The remaining update is not sync
+      ReactNoop.flushSync();
+      expect(Scheduler).toHaveYielded([]);
+
+      // Now flush it.
+      expect(Scheduler).toFlushUntilNextPaint(['1, 1']);
+    });
+    expect(root).toMatchRenderedOutput('1, 1');
+  });
+});