Fix Dispatcher reset edge cases (#38)

* WIP Ensure that Dispatcher is always reset correctly * Refactor and fix test execution * Add test for exceptions handling inside yield Co-authored-by: Maxim <khmm12@gmail.com>

Author: kitten

src/__tests__/suspense.test.js

@@ -37,6 +37,27 @@ describe('renderPrepass', () => {
  })
  })
 
+ it('rejects promise after getting of exception inside', () => {
+ const exception = new TypeError('Something went wrong')
+
+ const Inner = jest.fn(() => {
+ throw exception
+ })
+
+ const Outer = () => {
+ const start = Date.now()
+ while (Date.now() - start < 40) {}
+ return <Inner />
+ }
+
+ const render$ = renderPrepass(<Outer />)
+
+ return render$.catch(error => {
+ expect(Inner).toHaveBeenCalledTimes(1)
+ expect(error).toBe(exception)
+ })
+ })
+
  it('preserves the correct legacy context values across yields', () => {
  let called = false
  const Inner = (_, context) => {

src/index.js

@@ -2,15 +2,9 @@
 
 import React, { type Node, type Element } from 'react'
 import type { Visitor, YieldFrame, Frame, AbstractElement } from './types'
-import { visitChildren, resumeVisitChildren } from './visitor'
+import { visitChildren, resumeVisitChildren, update } from './visitor'
 import { getChildrenArray } from './element'
 
-import {
- updateFunctionComponent,
- updateClassComponent,
- updateLazyComponent
-} from './render'
-
 import {
  setCurrentContextStore,
  setCurrentContextMap,
@@ -21,7 +15,29 @@ const {
  ReactCurrentDispatcher
 } = (React: any).__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
 
-let prevDispatcher = ReactCurrentDispatcher.current
+/** wrapWithDispatcher accepts a function and wraps it
+ in one that sets up our ReactCurrentDispatcher and
+ resets it afterwards */
+function wrapWithDispatcher<T: Function>(exec: T): T {
+ // $FlowFixMe
+ return (...args) => {
+ const prevDispatcher = ReactCurrentDispatcher.current
+
+ try {
+ // The "Dispatcher" is what handles hook calls and
+ // a React internal that needs to be set to our dispatcher
+ ReactCurrentDispatcher.current = Dispatcher
+ return exec(...args)
+ } finally {
+ // We're resetting the dispatcher after we're done
+ ReactCurrentDispatcher.current = prevDispatcher
+ }
+ }
+}
+
+const resumeWithDispatcher = wrapWithDispatcher(resumeVisitChildren)
+const visitWithDispatcher = wrapWithDispatcher(visitChildren)
+const updateWithDispatcher = wrapWithDispatcher(update)
 
 /** visitChildren walks all elements (depth-first) and while it walks the
  element tree some components will suspend and put a `Frame` onto
@@ -34,50 +50,35 @@ const updateWithFrame = (
  visitor: Visitor
 ): Promise<void> => {
  if (frame.kind === 'frame.yield') {
- const yieldFrame: YieldFrame = frame
-
- return new Promise(resolve => {
+ return new Promise((resolve, reject) => {
  setImmediate(() => {
- prevDispatcher = ReactCurrentDispatcher.current
- ReactCurrentDispatcher.current = Dispatcher
- resumeVisitChildren(yieldFrame, queue, visitor)
- ReactCurrentDispatcher.current = prevDispatcher
- resolve()
+ try {
+ resumeWithDispatcher(frame, queue, visitor)
+ resolve()
+ } catch (err) {
+ reject(err)
+ }
  })
  })
  }
 
  return frame.thenable.then(() => {
- prevDispatcher = ReactCurrentDispatcher.current
- ReactCurrentDispatcher.current = Dispatcher
-
- let children = []
-
  // Update the component after we've suspended to rerender it,
  // at which point we'll actually get its children
- if (frame.kind === 'frame.class') {
- children = updateClassComponent(queue, frame)
- } else if (frame.kind === 'frame.hooks') {
- children = updateFunctionComponent(queue, frame)
- } else if (frame.kind === 'frame.lazy') {
- children = updateLazyComponent(queue, frame)
- }
-
+ const children = updateWithDispatcher(frame, queue)
  // Now continue walking the previously suspended component's
  // children (which might also suspend)
- visitChildren(getChildrenArray(children), queue, visitor)
- ReactCurrentDispatcher.current = prevDispatcher
+ visitWithDispatcher(getChildrenArray(children), queue, visitor)
  })
 }
 
 const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
- if (queue.length === 0) {
- return Promise.resolve()
- }
-
- return updateWithFrame(queue.shift(), queue, visitor).then(() =>
- flushFrames(queue, visitor)
- )
+ const frame = queue.shift()
+ return frame
+ ? updateWithFrame(frame, queue, visitor).then(() =>
+ flushFrames(queue, visitor)
+ )
+ : Promise.resolve()
 }
 
 const defaultVisitor = () => undefined
@@ -93,17 +94,9 @@ const renderPrepass = (element: Node, visitor?: Visitor): Promise<void> => {
  setCurrentContextStore(new Map())
 
  try {
- // The "Dispatcher" is what handles hook calls and
- // a React internal that needs to be set to our
- // dispatcher and reset after we're done
- prevDispatcher = ReactCurrentDispatcher.current
- ReactCurrentDispatcher.current = Dispatcher
-
- visitChildren(getChildrenArray(element), queue, fn)
+ visitWithDispatcher(getChildrenArray(element), queue, fn)
  } catch (error) {
  return Promise.reject(error)
- } finally {
- ReactCurrentDispatcher.current = prevDispatcher
  }
 
  return flushFrames(queue, fn)

src/visitor.js

@@ -11,8 +11,11 @@ import {
 
 import {
  mountFunctionComponent,
+ updateFunctionComponent,
  mountClassComponent,
+ updateClassComponent,
  mountLazyComponent,
+ updateLazyComponent,
  mountStyledComponent,
  isStyledElement
 } from './render'
@@ -263,3 +266,16 @@ export const resumeVisitChildren = (
  queue.unshift(makeYieldFrame(frame.children, frame.map, frame.store))
  }
 }
+
+export const update = (frame: Frame, queue: Frame[]): Node => {
+ switch (frame.kind) {
+ case 'frame.class':
+ return updateClassComponent(queue, frame)
+ case 'frame.hooks':
+ return updateFunctionComponent(queue, frame)
+ case 'frame.lazy':
+ return updateLazyComponent(queue, frame)
+ default:
+ return []
+ }
+}