Allow live Resolvers to be batched

Differential Revision: D40575460

fbshipit-source-id: b5331afdd115a49d2273d6b571577c50c05f63ad

packages/relay-runtime/store/__tests__/resolvers/LiveResolvers-test.js

@@ -129,3 +129,75 @@ test('unsubscribe happens when record is updated due to missing data', () => {
   expect(data.greeting).toBe('Yo user 2');
   expect(__debug.state.subscribers.size).toBe(1);
 });
+
+test('Updates can be batched', () => {
+  const source = RelayRecordSource.create({
+    'client:root': {
+      __id: 'client:root',
+      __typename: '__Root',
+    },
+  });
+  const operation = createOperationDescriptor(
+    graphql`
+      query LiveResolversTestBatchingQuery {
+        # Together these fields create two subscriptions to the underlying
+        # GLOBAL_STORE.
+        counter_no_fragment
+        counter_no_fragment_with_arg(prefix: "sup")
+      }
+    `,
+    {},
+  );
+  const store = new LiveResolverStore(source, {
+    gcReleaseBufferSize: 0,
+  });
+  const environment = new RelayModernEnvironment({
+    network: RelayNetwork.create(jest.fn()),
+    store,
+  });
+
+  const snapshot = environment.lookup(operation.fragment);
+
+  const handler = jest.fn();
+  environment.subscribe(snapshot, handler);
+
+  expect(handler.mock.calls.length).toBe(0);
+
+  // Update without batching
+  GLOBAL_STORE.dispatch({type: 'INCREMENT'});
+
+  // We get notified once per live resolver. :(
+  expect(handler.mock.calls.length).toBe(2);
+
+  let lastCallCount = handler.mock.calls.length;
+
+  // Update _with_ batching.
+  store.batchLiveStateUpdates(() => {
+    GLOBAL_STORE.dispatch({type: 'INCREMENT'});
+  });
+
+  // We get notified once per batch! :)
+  expect(handler.mock.calls.length - lastCallCount).toBe(1);
+
+  lastCallCount = handler.mock.calls.length;
+
+  // Update with batching, but update throws.
+  // This might happen if some other subscriber to the store throws when they
+  // get notified of an error.
+  expect(() => {
+    store.batchLiveStateUpdates(() => {
+      GLOBAL_STORE.dispatch({type: 'INCREMENT'});
+      throw new Error('An Example Error');
+    });
+  }).toThrowError('An Example Error');
+
+  // We still notify our subscribers
+  expect(handler.mock.calls.length - lastCallCount).toBe(1);
+
+  // Nested calls to batchLiveStateUpdate throw
+  expect(() => {
+    store.batchLiveStateUpdates(() => {
+      store.batchLiveStateUpdates(() => {});
+    });
+  }).toThrow('Unexpected nested call to batchLiveStateUpdates.');
+});

packages/relay-runtime/store/__tests__/resolvers/UserProfilePictureUriSuspendsWhenTheCounterIsOdd.js

@@ -28,7 +28,8 @@ const {readFragment} = require('relay-runtime/store/ResolverFragments');
  * @onType User
  * @live
  *
- * This field is returning the profile picture url, when s
+ * This field returns the profile picture url, when the GLOBAL_STORE number is
+ * even and suspends when the number is odd.
  */
 function user_profile_picture_uri_suspends_when_the_counter_is_odd(
   rootKey: UserProfilePictureUriSuspendsWhenTheCounterIsOdd$key,

packages/relay-runtime/store/experimental-live-resolvers/LiveResolverCache.js

@@ -85,6 +85,8 @@ class LiveResolverCache implements ResolverCache {
   _recordIDToResolverIDs: Map<DataID, Set<ResolverID>>;
   _getRecordSource: () => MutableRecordSource;
   _store: LiveResolverStore;
+  _handlingBatch: boolean; // Flag indicating that Live Resolver updates are being batched.
+  _liveResolverBatchRecordSource: ?MutableRecordSource; // Lazily created record source for batched Live Resolver updates.
 
   constructor(
     getRecordSource: () => MutableRecordSource,
@@ -94,6 +96,8 @@ class LiveResolverCache implements ResolverCache {
     this._recordIDToResolverIDs = new Map();
     this._getRecordSource = getRecordSource;
     this._store = store;
+    this._handlingBatch = false;
+    this._liveResolverBatchRecordSource = null;
   }
 
   readFromCacheOrEvaluate<T>(
@@ -403,7 +407,6 @@ class LiveResolverCache implements ResolverCache {
         return;
       }
 
-      const nextSource = RelayRecordSource.create();
       const nextRecord = RelayModernRecord.clone(currentRecord);
 
       // Mark the field as dirty. The next time it's read, we will call
@@ -414,13 +417,48 @@ class LiveResolverCache implements ResolverCache {
         true,
       );
 
-      nextSource.set(linkedID, nextRecord);
-      this._store.publish(nextSource);
+      this._setLiveResolverUpdate(linkedID, nextRecord);
+    };
+  }
 
-      // In the future, this notify might be deferred if we are within a
-      // transaction.
+  _setLiveResolverUpdate(linkedId: DataID, record: Record) {
+    if (this._handlingBatch) {
+      // Lazily create the batched record source.
+      if (this._liveResolverBatchRecordSource == null) {
+        this._liveResolverBatchRecordSource = RelayRecordSource.create();
+      }
+      this._liveResolverBatchRecordSource.set(linkedId, record);
+      // We will wait for the batch to complete before we publish/notify...
+    } else {
+      const nextSource = RelayRecordSource.create();
+      nextSource.set(linkedId, record);
+
+      // We are not within a batch, so we will immediately publish/notify.
+      this._store.publish(nextSource);
       this._store.notify();
-    };
+    }
+  }
+
+  batchLiveStateUpdates(callback: () => void) {
+    invariant(
+      !this._handlingBatch,
+      'Unexpected nested call to batchLiveStateUpdates.',
+    );
+    this._handlingBatch = true;
+    try {
+      callback();
+    } finally {
+      // We lazily create the record source. If one has not been created, there
+      // is nothing to publish.
+      if (this._liveResolverBatchRecordSource != null) {
+        this._store.publish(this._liveResolverBatchRecordSource);
+        this._store.notify();
+      }
+
+      // Reset batched state.
+      this._liveResolverBatchRecordSource = null;
+      this._handlingBatch = false;
+    }
   }
 
   _setRelayResolverValue(

packages/relay-runtime/store/experimental-live-resolvers/LiveResolverStore.js

@@ -191,6 +191,29 @@ class LiveResolverStore implements Store {
     return this._resolverCache.getLiveResolverPromise(recordID);
   }
 
+  /**
+   * When an external data proider knows it's going to notify us about multiple
+   * Live Resolver state updates in a single tick, it can batch them into a
+   * single Relay update by notifying us within a batch. All updates recieved by
+   * Relay during the evaluation of the provided `callback` will be aggregated
+   * into a single Relay update.
+   *
+   * A typical use with a Flux store might look like this:
+   *
+   * const originalDispatch = fluxStore.dispatch;
+   *
+   * function wrapped(action) {
+   *   relayStore.batchLiveStateUpdates(() => {
+   *     originalDispatch(action);
+   *   })
+   * }
+   *
+   * fluxStore.dispatch = wrapped;
+   */
+  batchLiveStateUpdates(callback: () => void) {
+    this._resolverCache.batchLiveStateUpdates(callback);
+  }
+
   check(
     operation: OperationDescriptor,
     options?: CheckOptions,