Fix UMD builds by re-exporting the scheduler priorities (#14914)

* Fix react-dom/server context memory retention

* Test for pollution of later renders

* Inline loop

* More tests

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

@@ -482,5 +482,98 @@ describe('ReactDOMServerIntegration', () => {
         );
       }
     });
+
+    // Regression test for https://github.com/facebook/react/issues/14705
+    it('does not pollute later renders when stream destroyed', () => {
+      const LoggedInUser = React.createContext('default');
+
+      const AppWithUser = user => (
+        <LoggedInUser.Provider value={user}>
+          <header>
+            <LoggedInUser.Consumer>{whoAmI => whoAmI}</LoggedInUser.Consumer>
+          </header>
+        </LoggedInUser.Provider>
+      );
+
+      const stream = ReactDOMServer.renderToNodeStream(
+        AppWithUser('Amy'),
+      ).setEncoding('utf8');
+
+      // This is an implementation detail because we test a memory leak
+      const {threadID} = stream.partialRenderer;
+
+      // Read enough to render Provider but not enough for it to be exited
+      stream._read(10);
+      expect(LoggedInUser[threadID]).toBe('Amy');
+
+      stream.destroy();
+
+      const AppWithUserNoProvider = () => (
+        <LoggedInUser.Consumer>{whoAmI => whoAmI}</LoggedInUser.Consumer>
+      );
+
+      const stream2 = ReactDOMServer.renderToNodeStream(
+        AppWithUserNoProvider(),
+      ).setEncoding('utf8');
+
+      // Sanity check to ensure 2nd render has same threadID as 1st render,
+      // otherwise this test is not testing what it's meant to
+      expect(stream2.partialRenderer.threadID).toBe(threadID);
+
+      const markup = stream2.read(Infinity);
+
+      expect(markup).toBe('default');
+    });
+
+    // Regression test for https://github.com/facebook/react/issues/14705
+    it('frees context value reference when stream destroyed', () => {
+      const LoggedInUser = React.createContext('default');
+
+      const AppWithUser = user => (
+        <LoggedInUser.Provider value={user}>
+          <header>
+            <LoggedInUser.Consumer>{whoAmI => whoAmI}</LoggedInUser.Consumer>
+          </header>
+        </LoggedInUser.Provider>
+      );
+
+      const stream = ReactDOMServer.renderToNodeStream(
+        AppWithUser('Amy'),
+      ).setEncoding('utf8');
+
+      // This is an implementation detail because we test a memory leak
+      const {threadID} = stream.partialRenderer;
+
+      // Read enough to render Provider but not enough for it to be exited
+      stream._read(10);
+      expect(LoggedInUser[threadID]).toBe('Amy');
+
+      stream.destroy();
+      expect(LoggedInUser[threadID]).toBe('default');
+    });
+
+    it('does not pollute sync renders after an error', () => {
+      const LoggedInUser = React.createContext('default');
+      const Crash = () => {
+        throw new Error('Boo!');
+      };
+      const AppWithUser = user => (
+        <LoggedInUser.Provider value={user}>
+          <LoggedInUser.Consumer>{whoAmI => whoAmI}</LoggedInUser.Consumer>
+          <Crash />
+        </LoggedInUser.Provider>
+      );
+
+      expect(() => {
+        ReactDOMServer.renderToString(AppWithUser('Casper'));
+      }).toThrow('Boo');
+
+      // Should not report a value from failed render
+      expect(
+        ReactDOMServer.renderToString(
+          <LoggedInUser.Consumer>{whoAmI => whoAmI}</LoggedInUser.Consumer>,
+        ),
+      ).toBe('default');
+    });
   });
 });

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -715,6 +715,7 @@ class ReactDOMServerRenderer {
   destroy() {
     if (!this.exhausted) {
       this.exhausted = true;
+      this.clearProviders();
       freeThreadID(this.threadID);
     }
   }
@@ -776,6 +777,15 @@ class ReactDOMServerRenderer {
     context[this.threadID] = previousValue;
   }
 
+  clearProviders(): void {
+    // Restore any remaining providers on the stack to previous values
+    for (let index = this.contextIndex; index >= 0; index--) {
+      const context: ReactContext<any> = this.contextStack[index];
+      const previousValue = this.contextValueStack[index];
+      context[this.threadID] = previousValue;
+    }
+  }
+
   read(bytes: number): string | null {
     if (this.exhausted) {
       return null;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -607,7 +607,6 @@ function updateReducer<S, I, A>(
         }
 
         hook.memoizedState = newState;
-
         // Don't persist the state accumlated from the render phase updates to
         // the base state unless the queue is empty.
         // TODO: Not sure if this is the desired semantics, but it's what we
@@ -616,6 +615,9 @@ function updateReducer<S, I, A>(
           hook.baseState = newState;
         }
 
+        queue.eagerReducer = reducer;
+        queue.eagerState = newState;
+
         return [newState, dispatch];
       }
     }

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

@@ -669,6 +669,76 @@ describe('ReactHooks', () => {
     }).toThrow('is not a function');
   });
 
+  it('does not forget render phase useState updates inside an effect', () => {
+    const {useState, useEffect} = React;
+
+    function Counter() {
+      const [counter, setCounter] = useState(0);
+      if (counter === 0) {
+        setCounter(x => x + 1);
+        setCounter(x => x + 1);
+      }
+      useEffect(() => {
+        setCounter(x => x + 1);
+        setCounter(x => x + 1);
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
+  it('does not forget render phase useReducer updates inside an effect with hoisted reducer', () => {
+    const {useReducer, useEffect} = React;
+
+    const reducer = x => x + 1;
+    function Counter() {
+      const [counter, increment] = useReducer(reducer, 0);
+      if (counter === 0) {
+        increment();
+        increment();
+      }
+      useEffect(() => {
+        increment();
+        increment();
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
+  it('does not forget render phase useReducer updates inside an effect with inline reducer', () => {
+    const {useReducer, useEffect} = React;
+
+    function Counter() {
+      const [counter, increment] = useReducer(x => x + 1, 0);
+      if (counter === 0) {
+        increment();
+        increment();
+      }
+      useEffect(() => {
+        increment();
+        increment();
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
   it('warns for bad useImperativeHandle first arg', () => {
     const {useImperativeHandle} = React;
     function App() {

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

@@ -454,7 +454,9 @@ describe('ReactHooksWithNoopRenderer', () => {
 
       // Test that it works on update, too. This time the log is a bit different
       // because we started with reducerB instead of reducerA.
-      counter.current.dispatch('reset');
+      ReactNoop.act(() => {
+        counter.current.dispatch('reset');
+      });
       ReactNoop.render(<Counter ref={counter} />);
       expect(ReactNoop.flush()).toEqual([
         'Render: 0',

packages/react/src/ReactSharedInternals.js

@@ -18,6 +18,11 @@ import {
   unstable_continueExecution,
   unstable_wrapCallback,
   unstable_getCurrentPriorityLevel,
+  unstable_IdlePriority,
+  unstable_ImmediatePriority,
+  unstable_LowPriority,
+  unstable_NormalPriority,
+  unstable_UserBlockingPriority,
 } from 'scheduler';
 import {
   __interactionsRef,
@@ -60,6 +65,11 @@ if (__UMD__) {
       unstable_pauseExecution,
       unstable_continueExecution,
       unstable_getCurrentPriorityLevel,
+      unstable_IdlePriority,
+      unstable_ImmediatePriority,
+      unstable_LowPriority,
+      unstable_NormalPriority,
+      unstable_UserBlockingPriority,
     },
     SchedulerTracing: {
       __interactionsRef,

packages/scheduler/npm/umd/scheduler.development.js

@@ -108,5 +108,25 @@
     unstable_continueExecution: unstable_continueExecution,
     unstable_pauseExecution: unstable_pauseExecution,
     unstable_getFirstCallbackNode: unstable_getFirstCallbackNode,
+    get unstable_IdlePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_IdlePriority;
+    },
+    get unstable_ImmediatePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_ImmediatePriority;
+    },
+    get unstable_LowPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_LowPriority;
+    },
+    get unstable_NormalPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_NormalPriority;
+    },
+    get unstable_UserBlockingPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_UserBlockingPriority;
+    },
   });
 });

packages/scheduler/npm/umd/scheduler.production.min.js

@@ -102,5 +102,25 @@
     unstable_continueExecution: unstable_continueExecution,
     unstable_pauseExecution: unstable_pauseExecution,
     unstable_getFirstCallbackNode: unstable_getFirstCallbackNode,
+    get unstable_IdlePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_IdlePriority;
+    },
+    get unstable_ImmediatePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_ImmediatePriority;
+    },
+    get unstable_LowPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_LowPriority;
+    },
+    get unstable_NormalPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_NormalPriority;
+    },
+    get unstable_UserBlockingPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_UserBlockingPriority;
+    },
   });
 });

packages/scheduler/npm/umd/scheduler.profiling.min.js

@@ -102,5 +102,25 @@
     unstable_continueExecution: unstable_continueExecution,
     unstable_pauseExecution: unstable_pauseExecution,
     unstable_getFirstCallbackNode: unstable_getFirstCallbackNode,
+    get unstable_IdlePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_IdlePriority;
+    },
+    get unstable_ImmediatePriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_ImmediatePriority;
+    },
+    get unstable_LowPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_LowPriority;
+    },
+    get unstable_NormalPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_NormalPriority;
+    },
+    get unstable_UserBlockingPriority() {
+      return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .Scheduler.unstable_UserBlockingPriority;
+    },
   });
 });

packages/scheduler/src/__tests__/SchedulerUMDBundle-test.internal.js

@@ -17,8 +17,17 @@ describe('Scheduling UMD bundle', () => {
   });
 
   function filterPrivateKeys(name) {
-    // TODO: Figure out how to forward priority levels.
-    return !name.startsWith('_') && !name.endsWith('Priority');
+    // Be very careful adding things to this whitelist!
+    // It's easy to introduce bugs by doing it:
+    // https://github.com/facebook/react/issues/14904
+    switch (name) {
+      case '__interactionsRef':
+      case '__subscriberRef':
+        // Don't forward these. (TODO: why?)
+        return false;
+      default:
+        return true;
+    }
   }
 
   function validateForwardedAPIs(api, forwardedAPIs) {