Merge pull request #8585 from acdlite/batchstateandcallback

Schedule state and callback at the same time

src/isomorphic/classic/class/ReactClass.js

@@ -728,10 +728,7 @@ var ReactClassMixin = {
    * type signature and the only use case for this, is to avoid that.
    */
   replaceState: function(newState, callback) {
-    this.updater.enqueueReplaceState(this, newState);
-    if (callback) {
-      this.updater.enqueueCallback(this, callback, 'replaceState');
-    }
+    this.updater.enqueueReplaceState(this, newState, callback, 'replaceState');
   },
 
   /**

src/isomorphic/modern/class/ReactComponent.js

@@ -65,10 +65,7 @@ ReactComponent.prototype.setState = function(partialState, callback) {
     'setState(...): takes an object of state variables to update or a ' +
     'function which returns an object of state variables.'
   );
-  this.updater.enqueueSetState(this, partialState);
-  if (callback) {
-    this.updater.enqueueCallback(this, callback, 'setState');
-  }
+  this.updater.enqueueSetState(this, partialState, callback, 'setState');
 };
 
 /**
@@ -86,10 +83,7 @@ ReactComponent.prototype.setState = function(partialState, callback) {
  * @protected
  */
 ReactComponent.prototype.forceUpdate = function(callback) {
-  this.updater.enqueueForceUpdate(this);
-  if (callback) {
-    this.updater.enqueueCallback(this, callback, 'forceUpdate');
-  }
+  this.updater.enqueueForceUpdate(this, callback, 'forceUpdate');
 };
 
 /**

src/isomorphic/modern/class/ReactNoopUpdateQueue.js

@@ -44,16 +44,6 @@ var ReactNoopUpdateQueue = {
     return false;
   },
 
-  /**
-   * Enqueue a callback that will be executed after all the pending updates
-   * have processed.
-   *
-   * @param {ReactClass} publicInstance The instance to use as `this` context.
-   * @param {?function} callback Called after state is updated.
-   * @internal
-   */
-  enqueueCallback: function(publicInstance, callback) { },
-
   /**
    * Forces an update. This should only be invoked when it is known with
    * certainty that we are **not** in a DOM transaction.
@@ -65,9 +55,11 @@ var ReactNoopUpdateQueue = {
    * `componentWillUpdate` and `componentDidUpdate`.
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueForceUpdate: function(publicInstance) {
+  enqueueForceUpdate: function(publicInstance, callback, callerName) {
     warnNoop(publicInstance, 'forceUpdate');
   },
 
@@ -80,9 +72,11 @@ var ReactNoopUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} completeState Next state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueReplaceState: function(publicInstance, completeState) {
+  enqueueReplaceState: function(publicInstance, completeState, callback, callerName) {
     warnNoop(publicInstance, 'replaceState');
   },
 
@@ -94,9 +88,11 @@ var ReactNoopUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} partialState Next partial state to be merged with state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueSetState: function(publicInstance, partialState) {
+  enqueueSetState: function(publicInstance, partialState, callback, callerName) {
     warnNoop(publicInstance, 'setState');
   },
 };

src/renderers/dom/stack/server/ReactServerUpdateQueue.js

@@ -58,20 +58,6 @@ class ReactServerUpdateQueue {
     return false;
   }
 
-  /**
-   * Enqueue a callback that will be executed after all the pending updates
-   * have processed.
-   *
-   * @param {ReactClass} publicInstance The instance to use as `this` context.
-   * @param {?function} callback Called after state is updated.
-   * @internal
-   */
-  enqueueCallback(publicInstance: ReactComponent<any, any, any>, callback?: Function, callerName?: string) {
-    if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueCallback(publicInstance, callback, callerName);
-    }
-  }
-
   /**
    * Forces an update. This should only be invoked when it is known with
    * certainty that we are **not** in a DOM transaction.
@@ -83,11 +69,13 @@ class ReactServerUpdateQueue {
    * `componentWillUpdate` and `componentDidUpdate`.
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueForceUpdate(publicInstance: ReactComponent<any, any, any>) {
+  enqueueForceUpdate(publicInstance: ReactComponent<any, any, any>, callback?: Function, callerName?: string) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueForceUpdate(publicInstance);
+      ReactUpdateQueue.enqueueForceUpdate(publicInstance, callback, callerName);
     } else {
       warnNoop(publicInstance, 'forceUpdate');
     }
@@ -102,11 +90,18 @@ class ReactServerUpdateQueue {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object|function} completeState Next state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueReplaceState(publicInstance: ReactComponent<any, any, any>, completeState: Object|Function) {
+  enqueueReplaceState(
+    publicInstance: ReactComponent<any, any, any>,
+    completeState: Object|Function,
+    callback?: Function,
+    callerName?: string
+  ) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueReplaceState(publicInstance, completeState);
+      ReactUpdateQueue.enqueueReplaceState(publicInstance, completeState, callback, callerName);
     } else {
       warnNoop(publicInstance, 'replaceState');
     }
@@ -122,9 +117,14 @@ class ReactServerUpdateQueue {
    * @param {object|function} partialState Next partial state to be merged with state.
    * @internal
    */
-  enqueueSetState(publicInstance: ReactComponent<any, any, any>, partialState: Object|Function) {
+  enqueueSetState(
+    publicInstance: ReactComponent<any, any, any>,
+    partialState: Object|Function,
+    callback?: Function,
+    callerName?: string
+  ) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueSetState(publicInstance, partialState);
+      ReactUpdateQueue.enqueueSetState(publicInstance, partialState, callback, callerName);
     } else {
       warnNoop(publicInstance, 'setState');
     }

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -71,10 +71,8 @@ if (__DEV__) {
 module.exports = function<T, P, I, TI, C, CX>(
   config : HostConfig<T, P, I, TI, C, CX>,
   hostContext : HostContext<C, CX>,
-  scheduleSetState: (fiber : Fiber, partialState : any) => void,
-  scheduleReplaceState: (fiber : Fiber, state : any) => void,
-  scheduleForceUpdate: (fiber : Fiber) => void,
-  scheduleUpdateCallback: (fiber : Fiber, callback : Function) => void,
+  scheduleUpdate : (fiber : Fiber, priorityLevel : PriorityLevel) => void,
+  getPriorityContext : () => PriorityLevel,
 ) {
 
   const { shouldSetTextContent } = config;
@@ -91,12 +89,7 @@ module.exports = function<T, P, I, TI, C, CX>(
     mountClassInstance,
     resumeMountClassInstance,
     updateClassInstance,
-  } = ReactFiberClassComponent(
-    scheduleSetState,
-    scheduleReplaceState,
-    scheduleForceUpdate,
-    scheduleUpdateCallback
-  );
+  } = ReactFiberClassComponent(scheduleUpdate, getPriorityContext);
 
   function markChildAsProgressed(current, workInProgress, priorityLevel) {
     // We now have clones. Let's store them as the currently progressed work.

src/renderers/shared/fiber/ReactFiberClassComponent.js

@@ -19,6 +19,9 @@ var {
   getMaskedContext,
 } = require('ReactFiberContext');
 var {
+  addUpdate,
+  addReplaceUpdate,
+  addForceUpdate,
   beginUpdateQueue,
 } = require('ReactFiberUpdateQueue');
 var { hasContextChanged } = require('ReactFiberContext');
@@ -31,30 +34,30 @@ var invariant = require('invariant');
 const isArray = Array.isArray;
 
 module.exports = function(
-  scheduleSetState: (fiber : Fiber, partialState : any) => void,
-  scheduleReplaceState: (fiber : Fiber, state : any) => void,
-  scheduleForceUpdate: (fiber : Fiber) => void,
-  scheduleUpdateCallback: (fiber : Fiber, callback : Function) => void,
+  scheduleUpdate : (fiber : Fiber, priorityLevel : PriorityLevel) => void,
+  getPriorityContext : () => PriorityLevel,
 ) {
 
   // Class component state updater
   const updater = {
     isMounted,
-    enqueueSetState(instance, partialState) {
+    enqueueSetState(instance, partialState, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleSetState(fiber, partialState);
+      const priorityLevel = getPriorityContext();
+      addUpdate(fiber, partialState, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
-    enqueueReplaceState(instance, state) {
+    enqueueReplaceState(instance, state, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleReplaceState(fiber, state);
+      const priorityLevel = getPriorityContext();
+      addReplaceUpdate(fiber, state, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
-    enqueueForceUpdate(instance) {
+    enqueueForceUpdate(instance, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleForceUpdate(fiber);
-    },
-    enqueueCallback(instance, callback) {
-      const fiber = ReactInstanceMap.get(instance);
-      scheduleUpdateCallback(fiber, callback);
+      const priorityLevel = getPriorityContext();
+      addForceUpdate(fiber, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
   };
 

src/renderers/shared/fiber/ReactFiberReconciler.js

@@ -17,6 +17,10 @@ import type { FiberRoot } from 'ReactFiberRoot';
 import type { PriorityLevel } from 'ReactPriorityLevel';
 import type { ReactNodeList } from 'ReactTypes';
 
+var {
+  addTopLevelUpdate,
+} = require('ReactFiberUpdateQueue');
+
 var {
   findCurrentUnmaskedContext,
   isContextProvider,
@@ -98,14 +102,21 @@ getContextForSubtree._injectFiber(function(fiber : Fiber) {
 module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C, CX>) : Reconciler<C, I, TI> {
 
   var {
-    scheduleTopLevelSetState,
-    scheduleUpdateCallback,
+    scheduleUpdate,
+    getPriorityContext,
     performWithPriority,
     batchedUpdates,
     syncUpdates,
     deferredUpdates,
   } = ReactFiberScheduler(config);
 
+  function scheduleTopLevelUpdate(current : Fiber, element : ReactNodeList, callback : ?Function) {
+    const priorityLevel = getPriorityContext();
+    const nextState = { element };
+    addTopLevelUpdate(current, nextState, callback || null, priorityLevel);
+    scheduleUpdate(current, priorityLevel);
+  }
+
   return {
 
     mountContainer(element : ReactNodeList, containerInfo : C, parentComponent : ?ReactComponent<any, any, any>, callback: ?Function) : OpaqueNode {
@@ -113,10 +124,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
       const root = createFiberRoot(containerInfo, context);
       const current = root.current;
 
-      scheduleTopLevelSetState(current, { element });
-      if (callback) {
-        scheduleUpdateCallback(current, callback);
-      }
+      scheduleTopLevelUpdate(current, element, callback);
 
       if (__DEV__ && ReactFiberInstrumentation.debugTool) {
         ReactFiberInstrumentation.debugTool.onMountContainer(root);
@@ -135,10 +143,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
 
       root.pendingContext = getContextForSubtree(parentComponent);
 
-      scheduleTopLevelSetState(current, { element });
-      if (callback) {
-        scheduleUpdateCallback(current, callback);
-      }
+      scheduleTopLevelUpdate(current, element, callback);
 
       if (__DEV__) {
         if (ReactFiberInstrumentation.debugTool) {

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -55,11 +55,6 @@ var {
 
 var {
   getPendingPriority,
-  addUpdate,
-  addReplaceUpdate,
-  addForceUpdate,
-  addCallback,
-  addTopLevelUpdate,
 } = require('ReactFiberUpdateQueue');
 
 var {
@@ -79,10 +74,8 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
   const { beginWork, beginFailedWork } = ReactFiberBeginWork(
     config,
     hostContext,
-    scheduleSetState,
-    scheduleReplaceState,
-    scheduleForceUpdate,
-    scheduleUpdateCallback,
+    scheduleUpdate,
+    getPriorityContext,
   );
   const { completeWork } = ReactFiberCompleteWork(config, hostContext);
   const {
@@ -986,7 +979,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     }
   }
 
-  function scheduleUpdateAtPriority(fiber : Fiber, priorityLevel : PriorityLevel) {
+  function scheduleUpdate(fiber : Fiber, priorityLevel : PriorityLevel) {
     // If we're in a batch, downgrade sync priority to task priority
     if (priorityLevel === SynchronousPriority && isPerformingWork) {
       priorityLevel = TaskPriority;
@@ -1049,34 +1042,15 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     }
   }
 
+  function getPriorityContext() : PriorityLevel {
+    if (priorityContext === SynchronousPriority && isPerformingWork) {
+      return TaskPriority;
+    }
+    return priorityContext;
+  }
+
   function scheduleErrorRecovery(fiber : Fiber) {
-    scheduleUpdateAtPriority(fiber, TaskPriority);
-  }
-
-  function scheduleSetState(fiber : Fiber, partialState : any) {
-    addUpdate(fiber, partialState, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleReplaceState(fiber : Fiber, state : any) {
-    addReplaceUpdate(fiber, state, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleForceUpdate(fiber : Fiber) {
-    addForceUpdate(fiber, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleUpdateCallback(fiber : Fiber, callback : Function) {
-    addCallback(fiber, callback, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  // TODO: This indirection will be removed as part of #8585
-  function scheduleTopLevelSetState(fiber : Fiber, partialState : any) {
-    addTopLevelUpdate(fiber, partialState, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
+    scheduleUpdate(fiber, TaskPriority);
   }
 
   function performWithPriority(priorityLevel : PriorityLevel, fn : Function) {
@@ -1126,8 +1100,8 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
   }
 
   return {
-    scheduleTopLevelSetState: scheduleTopLevelSetState,
-    scheduleUpdateCallback: scheduleUpdateCallback,
+    scheduleUpdate: scheduleUpdate,
+    getPriorityContext: getPriorityContext,
     performWithPriority: performWithPriority,
     batchedUpdates: batchedUpdates,
     syncUpdates: syncUpdates,

src/renderers/shared/fiber/ReactFiberUpdateQueue.js

@@ -284,12 +284,13 @@ function insertUpdate(fiber : Fiber, update : Update, methodName : ?string) : Up
 function addUpdate(
   fiber : Fiber,
   partialState : PartialState<any, any> | null,
+  callback : Callback | null,
   priorityLevel : PriorityLevel
 ) : void {
   const update = {
     priorityLevel,
     partialState,
-    callback: null,
+    callback,
     isReplace: false,
     isForced: false,
     isTopLevelUnmount: false,
@@ -306,12 +307,13 @@ exports.addUpdate = addUpdate;
 function addReplaceUpdate(
   fiber : Fiber,
   state : any | null,
+  callback : Callback | null,
   priorityLevel : PriorityLevel
 ) : void {
   const update = {
     priorityLevel,
     partialState: state,
-    callback: null,
+    callback,
     isReplace: true,
     isForced: false,
     isTopLevelUnmount: false,
@@ -326,11 +328,15 @@ function addReplaceUpdate(
 }
 exports.addReplaceUpdate = addReplaceUpdate;
 
-function addForceUpdate(fiber : Fiber, priorityLevel : PriorityLevel) : void {
+function addForceUpdate(
+  fiber : Fiber,
+  callback : Callback | null,
+  priorityLevel : PriorityLevel
+) : void {
   const update = {
     priorityLevel,
     partialState: null,
-    callback: null,
+    callback,
     isReplace: false,
     isForced: true,
     isTopLevelUnmount: false,
@@ -344,21 +350,6 @@ function addForceUpdate(fiber : Fiber, priorityLevel : PriorityLevel) : void {
 }
 exports.addForceUpdate = addForceUpdate;
 
-
-function addCallback(fiber : Fiber, callback: Callback, priorityLevel : PriorityLevel) : void {
-  const update : Update = {
-    priorityLevel,
-    partialState: null,
-    callback,
-    isReplace: false,
-    isForced: false,
-    isTopLevelUnmount: false,
-    next: null,
-  };
-  insertUpdate(fiber, update);
-}
-exports.addCallback = addCallback;
-
 function getPendingPriority(queue : UpdateQueue) : PriorityLevel {
   return queue.first ? queue.first.priorityLevel : NoWork;
 }
@@ -367,14 +358,18 @@ exports.getPendingPriority = getPendingPriority;
 function addTopLevelUpdate(
   fiber : Fiber,
   partialState : PartialState<any, any>,
+  callback : Callback | null,
   priorityLevel : PriorityLevel
 ) : void {
-  const isTopLevelUnmount = partialState === null;
+  const isTopLevelUnmount = Boolean(
+    partialState &&
+    partialState.element === null
+  );
 
   const update = {
     priorityLevel,
     partialState,
-    callback: null,
+    callback,
     isReplace: false,
     isForced: false,
     isTopLevelUnmount,

src/renderers/shared/shared/__tests__/ReactCompositeComponentState-test.js

@@ -171,20 +171,8 @@ describe('ReactCompositeComponent-state', () => {
       ['componentDidMount-end', 'orange'],
       ['setState-sunrise', 'orange'],
       ['setState-orange', 'orange'],
-    ];
-
-    // In Fiber, the initial callback is not enqueued until after any work
-    // scheduled by lifecycles has flushed (same semantics as a regular setState
-    // callback outside of a batch). In Stack, the initial render is scheduled
-    // inside of batchedUpdates, so the callback gets flushed right after
-    // componentDidMount.
-    // TODO: We should fix this, in both Stack and Fiber, so that the behavior
-    // is consistent regardless of whether you're in a batch.
-    if (!ReactDOMFeatureFlags.useFiber) {
-      expected.push(['initial-callback', 'orange']);
-    }
-
-    expected.push(['shouldComponentUpdate-currentState', 'orange'],
+      ['initial-callback', 'orange'],
+      ['shouldComponentUpdate-currentState', 'orange'],
       ['shouldComponentUpdate-nextState', 'yellow'],
       ['componentWillUpdate-currentState', 'orange'],
       ['componentWillUpdate-nextState', 'yellow'],
@@ -192,18 +180,11 @@ describe('ReactCompositeComponent-state', () => {
       ['componentDidUpdate-currentState', 'yellow'],
       ['componentDidUpdate-prevState', 'orange'],
       ['setState-yellow', 'yellow'],
-    );
-
-    if (ReactDOMFeatureFlags.useFiber) {
-      expected.push(['initial-callback', 'yellow']);
-    }
-
-    expected.push(
       ['componentWillReceiveProps-start', 'yellow'],
       // setState({color:'green'}) only enqueues a pending state.
       ['componentWillReceiveProps-end', 'yellow'],
       // pending state queue is processed
-    );
+    ];
 
     if (ReactDOMFeatureFlags.useFiber) {
       // In Stack, this is never called because replaceState drops all updates

src/renderers/shared/shared/__tests__/ReactUpdates-test.js

@@ -648,16 +648,18 @@ describe('ReactUpdates', () => {
           'Inner-render-1-0',
           'Inner-didUpdate-1-0',
         'Outer-didUpdate-1',
+           // Happens in a batch, so don't re-render yet
           'Inner-setState-1',
-            'Inner-render-1-1',
-            'Inner-didUpdate-1-1',
-            'Inner-callback-1',
         'Outer-callback-1',
 
-      'Outer-setState-2',
+        // Happens in a batch
+        'Outer-setState-2',
+
+        // Flush batched updates all at once
         'Outer-render-2',
           'Inner-render-2-1',
           'Inner-didUpdate-2-1',
+          'Inner-callback-1',
         'Outer-didUpdate-2',
           'Inner-setState-2',
         'Outer-callback-2',

src/renderers/shared/stack/reconciler/ReactUpdateQueue.js

@@ -112,40 +112,6 @@ var ReactUpdateQueue = {
     }
   },
 
-  /**
-   * Enqueue a callback that will be executed after all the pending updates
-   * have processed.
-   *
-   * @param {ReactClass} publicInstance The instance to use as `this` context.
-   * @param {?function} callback Called after state is updated.
-   * @param {string} callerName Name of the calling function in the public API.
-   * @internal
-   */
-  enqueueCallback: function(publicInstance, callback, callerName) {
-    ReactUpdateQueue.validateCallback(callback, callerName);
-    var internalInstance = getInternalInstanceReadyForUpdate(publicInstance);
-
-    // Previously we would throw an error if we didn't have an internal
-    // instance. Since we want to make it a no-op instead, we mirror the same
-    // behavior we have in other enqueue* methods.
-    // We also need to ignore callbacks in componentWillMount. See
-    // enqueueUpdates.
-    if (!internalInstance) {
-      return null;
-    }
-
-    if (internalInstance._pendingCallbacks) {
-      internalInstance._pendingCallbacks.push(callback);
-    } else {
-      internalInstance._pendingCallbacks = [callback];
-    }
-    // TODO: The callback here is ignored when setState is called from
-    // componentWillMount. Either fix it or disallow doing so completely in
-    // favor of getInitialState. Alternatively, we can disallow
-    // componentWillMount during server-side rendering.
-    enqueueUpdate(internalInstance);
-  },
-
   enqueueCallbackInternal: function(internalInstance, callback) {
     if (internalInstance._pendingCallbacks) {
       internalInstance._pendingCallbacks.push(callback);
@@ -166,9 +132,11 @@ var ReactUpdateQueue = {
    * `componentWillUpdate` and `componentDidUpdate`.
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueForceUpdate: function(publicInstance) {
+  enqueueForceUpdate: function(publicInstance, callback, callerName) {
     var internalInstance = getInternalInstanceReadyForUpdate(
       publicInstance,
       'forceUpdate'
@@ -178,6 +146,15 @@ var ReactUpdateQueue = {
       return;
     }
 
+    if (callback) {
+      ReactUpdateQueue.validateCallback(callback, callerName);
+      if (internalInstance._pendingCallbacks) {
+        internalInstance._pendingCallbacks.push(callback);
+      } else {
+        internalInstance._pendingCallbacks = [callback];
+      }
+    }
+
     internalInstance._pendingForceUpdate = true;
 
     enqueueUpdate(internalInstance);
@@ -192,9 +169,11 @@ var ReactUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} completeState Next state.
+   * @param {?function} callback Called after state is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueReplaceState: function(publicInstance, completeState) {
+  enqueueReplaceState: function(publicInstance, completeState, callback, callerName) {
     var internalInstance = getInternalInstanceReadyForUpdate(
       publicInstance,
       'replaceState'
@@ -207,6 +186,15 @@ var ReactUpdateQueue = {
     internalInstance._pendingStateQueue = [completeState];
     internalInstance._pendingReplaceState = true;
 
+    if (callback) {
+      ReactUpdateQueue.validateCallback(callback, callerName);
+      if (internalInstance._pendingCallbacks) {
+        internalInstance._pendingCallbacks.push(callback);
+      } else {
+        internalInstance._pendingCallbacks = [callback];
+      }
+    }
+
     enqueueUpdate(internalInstance);
   },
 
@@ -218,9 +206,11 @@ var ReactUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} partialState Next partial state to be merged with state.
+   * @param {?function} callback Called after state is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueSetState: function(publicInstance, partialState) {
+  enqueueSetState: function(publicInstance, partialState, callback, callerName) {
     if (__DEV__) {
       ReactInstrumentation.debugTool.onSetState();
       warning(
@@ -244,6 +234,15 @@ var ReactUpdateQueue = {
       (internalInstance._pendingStateQueue = []);
     queue.push(partialState);
 
+    if (callback) {
+      ReactUpdateQueue.validateCallback(callback, callerName);
+      if (internalInstance._pendingCallbacks) {
+        internalInstance._pendingCallbacks.push(callback);
+      } else {
+        internalInstance._pendingCallbacks = [callback];
+      }
+    }
+
     enqueueUpdate(internalInstance);
   },