From 8570d19af69e5ecfc4f0d4e8b2af937efbc860b7 Mon Sep 17 00:00:00 2001
From: Naoyuki Kanezawa <naoyuki.kanezawa@gmail.com>
Date: Fri, 6 Jan 2017 02:27:39 +0900
Subject: [PATCH] Handle errors of React lifecycle methods (#661)

* handle errors of lifecycle methods

* handle errors of render method
---
 client/next-dev.js    |  9 ++++++-
 client/next.js        | 30 ++++++++++++++++++---
 client/patch-react.js | 62 +++++++++++++++++++++++++++++++++++++++++++
 lib/app.js            | 19 +++++++++++--
 4 files changed, 114 insertions(+), 6 deletions(-)
 create mode 100644 client/patch-react.js

diff --git a/client/next-dev.js b/client/next-dev.js
index 0fc62c70..3d62249c 100644
--- a/client/next-dev.js
+++ b/client/next-dev.js
@@ -1,4 +1,11 @@
 import 'react-hot-loader/patch'
-import * as next from './next'
+import patch from './patch-react'
 
+// apply patch first
+patch((err) => {
+  console.error(err)
+  next.renderError(err)
+})
+
+const next = require('./next')
 window.next = next
diff --git a/client/next.js b/client/next.js
index 70d0b06c..a1897675 100644
--- a/client/next.js
+++ b/client/next.js
@@ -1,5 +1,5 @@
 import { createElement } from 'react'
-import { render } from 'react-dom'
+import ReactDOM from 'react-dom'
 import HeadManager from './head-manager'
 import { rehydrate } from '../lib/css'
 import { createRouter } from '../lib/router'
@@ -29,7 +29,31 @@ export const router = createRouter(pathname, query, {
 
 const headManager = new HeadManager()
 const container = document.getElementById('__next')
-const appProps = { Component, props, router, headManager }
+const defaultProps = { Component, ErrorComponent, props, router, headManager }
 
 if (ids && ids.length) rehydrate(ids)
-render(createElement(App, appProps), container)
+
+render()
+
+export function render (props = {}) {
+  try {
+    doRender(props)
+  } catch (err) {
+    renderError(err)
+  }
+}
+
+export async function renderError (err) {
+  const { pathname, query } = router
+  const props = await ErrorComponent.getInitialProps({ err, pathname, query })
+  try {
+    doRender({ Component: ErrorComponent, props })
+  } catch (err2) {
+    console.error(err2)
+  }
+}
+
+function doRender (props) {
+  const appProps = { ...defaultProps, ...props }
+  ReactDOM.render(createElement(App, appProps), container)
+}
diff --git a/client/patch-react.js b/client/patch-react.js
new file mode 100644
index 00000000..b16b8bb9
--- /dev/null
+++ b/client/patch-react.js
@@ -0,0 +1,62 @@
+// monkeypatch React for fixing https://github.com/facebook/react/issues/2461
+// based on https://gist.github.com/Aldredcz/4d63b0a9049b00f54439f8780be7f0d8
+
+import React from 'react'
+
+let patched = false
+
+export default (handleError = () => {}) => {
+  if (patched) {
+    throw new Error('React is already monkeypatched')
+  }
+
+  patched = true
+
+  const { createElement } = React
+
+  React.createElement = function (Component, ...rest) {
+    if (typeof Component === 'function') {
+      const { prototype } = Component
+      if (prototype && prototype.render) {
+        prototype.render = wrapRender(prototype.render)
+      } else {
+        // stateless component
+        Component = wrapRender(Component)
+      }
+    }
+
+    return createElement.call(this, Component, ...rest)
+  }
+
+  const { Component: { prototype: componentPrototype } } = React
+  const { forceUpdate } = componentPrototype
+
+  componentPrototype.forceUpdate = function (...args) {
+    if (this.render) {
+      this.render = wrapRender(this.render)
+    }
+    return forceUpdate.apply(this, args)
+  }
+
+  function wrapRender (render) {
+    if (render.__wrapped) {
+      return render.__wrapped
+    }
+
+    const _render = function (...args) {
+      try {
+        return render.apply(this, args)
+      } catch (err) {
+        handleError(err)
+        return null
+      }
+    }
+
+    // copy all properties
+    Object.assign(_render, render)
+
+    render.__wrapped = _render.__wrapped = _render
+
+    return _render
+  }
+}
diff --git a/lib/app.js b/lib/app.js
index 126ed594..5367ea05 100644
--- a/lib/app.js
+++ b/lib/app.js
@@ -19,7 +19,7 @@ export default class App extends Component {
     try {
       this.setState(state)
     } catch (err) {
-      console.error(err)
+      this.handleError(err)
     }
   }
 
@@ -37,7 +37,7 @@ export default class App extends Component {
       try {
         this.setState(state)
       } catch (err) {
-        console.error(err)
+        this.handleError(err)
       }
     })
   }
@@ -58,6 +58,21 @@ export default class App extends Component {
       <Component {...props} />
     </AppContainer>
   }
+
+  async handleError (err) {
+    console.error(err)
+
+    const { router, ErrorComponent } = this.props
+    const { pathname, query } = router
+    const props = await ErrorComponent.getInitialProps({ err, pathname, query })
+    const state = propsToState({ Component: ErrorComponent, props, router })
+
+    try {
+      this.setState(state)
+    } catch (err2) {
+      console.error(err2)
+    }
+  }
 }
 
 function propsToState (props) {