From 077822e9d0f8a26bd505ee38e6353449ac8a845b Mon Sep 17 00:00:00 2001
From: Ankeet Maini <ankeet.maini@gmail.com>
Date: Sun, 13 Nov 2016 02:21:36 +0530
Subject: [PATCH] Handles risky callbacks on setState. Fixes #8238 (#8242)

* Handles risky callbacks on setState. Fixes #8238

* Updates try-catch to cover `callback` when context is not present.

* Updates code to trapErrors instead of swallowing them.

* Fixes flow errors

* Incorporates review comments

* Traps only the first error.

Removes `callbackWasCalled` and updates fiber tests.
---
 scripts/fiber/tests-passing.txt               |  1 +
 .../shared/fiber/ReactFiberCommitWork.js      | 13 +++++-
 .../shared/fiber/ReactFiberUpdateQueue.js     | 18 +++++---
 .../fiber/__tests__/ReactIncremental-test.js  | 44 +++++++++++++++++++
 4 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/scripts/fiber/tests-passing.txt b/scripts/fiber/tests-passing.txt
index 438bec74ba..b296e23566 100644
--- a/scripts/fiber/tests-passing.txt
+++ b/scripts/fiber/tests-passing.txt
@@ -806,6 +806,7 @@ src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
 * skips will/DidUpdate when bailing unless an update was already in progress
 * performs batched updates at the end of the batch
 * can nest batchedUpdates
+* can handle if setState callback throws
 
 src/renderers/shared/fiber/__tests__/ReactIncrementalErrorHandling-test.js
 * catches render error in a boundary during mounting
diff --git a/src/renderers/shared/fiber/ReactFiberCommitWork.js b/src/renderers/shared/fiber/ReactFiberCommitWork.js
index 4676d772ed..06bcb2a75c 100644
--- a/src/renderers/shared/fiber/ReactFiberCommitWork.js
+++ b/src/renderers/shared/fiber/ReactFiberCommitWork.js
@@ -321,7 +321,11 @@ module.exports = function<T, P, I, TI, C>(
         }
         if (finishedWork.effectTag & Callback) {
           if (finishedWork.callbackList) {
-            callCallbacks(finishedWork.callbackList, instance);
+            const callbackError = callCallbacks(finishedWork.callbackList, instance);
+            // since we only want to keep the first error
+            if (!error) {
+              error = callbackError;
+            }
             finishedWork.callbackList = null;
           }
         }
@@ -332,10 +336,15 @@ module.exports = function<T, P, I, TI, C>(
       }
       case HostContainer: {
         const rootFiber = finishedWork.stateNode;
+        let error = null;
         if (rootFiber.callbackList) {
           const { callbackList } = rootFiber;
           rootFiber.callbackList = null;
-          callCallbacks(callbackList, rootFiber.current.child.stateNode);
+          error = callCallbacks(callbackList, rootFiber.current.child.stateNode);
+        }
+
+        if (error) {
+          trapError(rootFiber, error, false);
         }
         return;
       }
diff --git a/src/renderers/shared/fiber/ReactFiberUpdateQueue.js b/src/renderers/shared/fiber/ReactFiberUpdateQueue.js
index c3c8cf7904..77c1290014 100644
--- a/src/renderers/shared/fiber/ReactFiberUpdateQueue.js
+++ b/src/renderers/shared/fiber/ReactFiberUpdateQueue.js
@@ -69,20 +69,26 @@ exports.addCallbackToQueue = function(queue : UpdateQueue, callback: Function) :
   return queue;
 };
 
-exports.callCallbacks = function(queue : UpdateQueue, context : any) {
+exports.callCallbacks = function(queue : UpdateQueue, context : any) : Error | null {
   let node : ?UpdateQueueNode = queue;
+  let error = null;
   while (node) {
     const callback = node.callback;
     if (callback && !node.callbackWasCalled) {
-      node.callbackWasCalled = true;
-      if (typeof context !== 'undefined') {
-        callback.call(context);
-      } else {
-        callback();
+      try {
+        node.callbackWasCalled = true;
+        if (typeof context !== 'undefined') {
+          callback.call(context);
+        } else {
+          callback();
+        }
+      } catch (e) {
+        error = e;
       }
     }
     node = node.next;
   }
+  return error;
 };
 
 exports.mergeUpdateQueue = function(queue : UpdateQueue, instance : any, prevState : any, props : any) : any {
diff --git a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
index 5df5c10caf..928af62d59 100644
--- a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
+++ b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
@@ -1465,4 +1465,48 @@ describe('ReactIncremental', () => {
     ]);
     expect(instance.state.n).toEqual(4);
   });
+
+  it('can handle if setState callback throws', () => {
+    var ops = [];
+    var instance;
+
+    class Foo extends React.Component {
+      state = { n: 0 };
+      render() {
+        instance = this;
+        return <div />;
+      }
+    }
+
+    ReactNoop.render(<Foo />);
+    ReactNoop.flush();
+    ops = [];
+
+    expect(instance.state.n).toEqual(0);
+
+    // first good callback
+    instance.setState({ n: 1 }, () => ops.push('first good callback'));
+    ReactNoop.flush();
+
+    // callback throws
+    instance.setState({ n: 2 }, () => {
+      throw new Error('Bail');
+    });
+    expect(() => {
+      ReactNoop.flush();
+    }).toThrow('Bail');
+    
+    // should set state to 2 even if callback throws up
+    expect(instance.state.n).toEqual(2);
+
+    // another good callback
+    instance.setState({ n: 3 }, () => ops.push('second good callback'));
+    ReactNoop.flush();
+
+    expect(ops).toEqual([
+      'first good callback',
+      'second good callback',
+    ]);
+    expect(instance.state.n).toEqual(3);
+  });
 });