Add ReactUpdates.setImmediate for async callbacks

Callbacks passed to this setImmediate function are called at the end of the current update cycle, which is guaranteed to be asynchronous but in the same event loop (with the default batching strategy).

This is useful for new-style refs (facebook#1373, facebook#1554) and for fixing facebook#1698.

Test Plan: jest

Author: sophiebits

src/core/ReactUpdates.js

@@ -29,6 +29,8 @@ var mixInto = require('mixInto');
 var warning = require('warning');
 
 var dirtyComponents = [];
+var setImmediateCallbackQueue = CallbackQueue.getPooled();
+var setImmediateEnqueued = false;
 
 var batchingStrategy = null;
 
@@ -73,7 +75,7 @@ var TRANSACTION_WRAPPERS = [NESTED_UPDATES, UPDATE_QUEUEING];
 function ReactUpdatesFlushTransaction() {
   this.reinitializeTransaction();
   this.dirtyComponentsLength = null;
-  this.callbackQueue = CallbackQueue.getPooled(null);
+  this.callbackQueue = CallbackQueue.getPooled();
   this.reconcileTransaction =
     ReactUpdates.ReactReconcileTransaction.getPooled();
 }
@@ -170,11 +172,21 @@ var flushBatchedUpdates = ReactPerf.measure(
     // ReactUpdatesFlushTransaction's wrappers will clear the dirtyComponents
     // array and perform any updates enqueued by mount-ready handlers (i.e.,
     // componentDidUpdate) but we need to check here too in order to catch
-    // updates enqueued by setState callbacks.
-    while (dirtyComponents.length) {
-      var transaction = ReactUpdatesFlushTransaction.getPooled();
-      transaction.perform(runBatchedUpdates, null, transaction);
-      ReactUpdatesFlushTransaction.release(transaction);
+    // updates enqueued by setState callbacks and setImmediate calls.
+    while (dirtyComponents.length || setImmediateEnqueued) {
+      if (dirtyComponents.length) {
+        var transaction = ReactUpdatesFlushTransaction.getPooled();
+        transaction.perform(runBatchedUpdates, null, transaction);
+        ReactUpdatesFlushTransaction.release(transaction);
+      }
+
+      if (setImmediateEnqueued) {
+        setImmediateEnqueued = false;
+        var queue = setImmediateCallbackQueue;
+        setImmediateCallbackQueue = CallbackQueue.getPooled();
+        queue.notifyAll();
+        CallbackQueue.release(queue);
+      }
     }
   }
 );
@@ -221,6 +233,20 @@ function enqueueUpdate(component, callback) {
   }
 }
 
+/**
+ * Enqueue a callback to be run at the end of the current batching cycle. Throws
+ * if no updates are currently being performed.
+ */
+function setImmediate(callback, context) {
+  invariant(
+    batchingStrategy.isBatchingUpdates,
+    'ReactUpdates.setImmediate: Can\'t enqueue an immediate callback in a ' +
+    'context where updates are not being batched.'
+  );
+  setImmediateCallbackQueue.enqueue(callback, context);
+  setImmediateEnqueued = true;
+}
+
 var ReactUpdatesInjection = {
   injectReconcileTransaction: function(ReconcileTransaction) {
     invariant(
@@ -259,7 +285,8 @@ var ReactUpdates = {
   batchedUpdates: batchedUpdates,
   enqueueUpdate: enqueueUpdate,
   flushBatchedUpdates: flushBatchedUpdates,
-  injection: ReactUpdatesInjection
+  injection: ReactUpdatesInjection,
+  setImmediate: setImmediate
 };
 
 module.exports = ReactUpdates;

src/core/__tests__/ReactUpdates-test.js

@@ -748,4 +748,79 @@ describe('ReactUpdates', function() {
     React.renderComponent(<A x={2} />, container);
     expect(callbackCount).toBe(1);
   });
+
+  it('calls setImmediate callbacks properly', function() {
+    var callbackCount = 0;
+    var A = React.createClass({
+      render: function() {
+        return <div />;
+      },
+      componentDidUpdate: function() {
+        var component = this;
+        ReactUpdates.setImmediate(function() {
+          expect(this).toBe(component);
+          callbackCount++;
+          ReactUpdates.setImmediate(function() {
+            callbackCount++;
+          });
+          expect(callbackCount).toBe(1);
+        }, this);
+        expect(callbackCount).toBe(0);
+      }
+    });
+
+    var container = document.createElement('div');
+    var component = React.renderComponent(<A />, container);
+    component.forceUpdate();
+    expect(callbackCount).toBe(2);
+  });
+
+  it('calls setImmediate callbacks with queued updates', function() {
+    var log = [];
+    var A = React.createClass({
+      getInitialState: () => ({updates: 0}),
+      render: function() {
+        log.push('render-' + this.state.updates);
+        return <div />;
+      },
+      componentDidUpdate: function() {
+        if (this.state.updates === 1) {
+          ReactUpdates.setImmediate(function() {
+            this.setState({updates: 2}, function() {
+              ReactUpdates.setImmediate(function() {
+                log.push('setImmediate-1.2');
+              });
+              log.push('setState-cb');
+            });
+            log.push('setImmediate-1.1');
+          }, this);
+        } else if (this.state.updates === 2) {
+          ReactUpdates.setImmediate(function() {
+            log.push('setImmediate-2');
+          });
+        }
+        log.push('didUpdate-' + this.state.updates);
+      }
+    });
+
+    var container = document.createElement('div');
+    var component = React.renderComponent(<A />, container);
+    component.setState({updates: 1});
+    expect(log).toEqual([
+      'render-0',
+      // We do the first update...
+      'render-1',
+      'didUpdate-1',
+      // ...which calls a setImmediate and enqueues a second update...
+      'setImmediate-1.1',
+      // ...which runs and enqueues the setImmediate-2 log in its didUpdate...
+      'render-2',
+      'didUpdate-2',
+      // ...and runs the setState callback, which enqueues the log for
+      // setImmediate-1.2.
+      'setState-cb',
+      'setImmediate-2',
+      'setImmediate-1.2'
+    ]);
+  });
 });