Adding onSnapshotsInSync (not public) (#2100)

Author: Brian Chen

packages/firestore/src/api/database.ts

@@ -452,6 +452,58 @@ export class Firestore implements firestore.FirebaseFirestore, FirebaseService {
     return this._firestoreClient!.waitForPendingWrites();
   }
 
+  /**
+   * Attaches a listener for a snapshots-in-sync event. The snapshots-in-sync
+   * event indicates that all listeners affected by a given change have fired,
+   * even if a single server-generated change affects multiple listeners.
+   *
+   * NOTE: The snapshots-in-sync event only indicates that listeners are in sync
+   * with each other, but does not relate to whether those snapshots are in sync
+   * with the server. Use SnapshotMetadata in the individual listeners to
+   * determine if a snapshot is from the cache or the server.
+   *
+   * @param onSync A callback to be called every time all snapshot listeners are
+   * in sync with each other.
+   * @return An unsubscribe function that can be called to cancel the snapshot
+   * listener.
+   */
+  _onSnapshotsInSync(observer: PartialObserver<void>): Unsubscribe;
+  _onSnapshotsInSync(onSync: () => void): Unsubscribe;
+  _onSnapshotsInSync(arg: unknown): Unsubscribe {
+    this.ensureClientConfigured();
+
+    if (isPartialObserver(arg)) {
+      return this.onSnapshotsInSyncInternal(arg as PartialObserver<void>);
+    } else {
+      validateArgType('Firestore.onSnapshotsInSync', 'function', 1, arg);
+      const observer: PartialObserver<void> = {
+        next: arg as () => void
+      };
+      return this.onSnapshotsInSyncInternal(observer);
+    }
+  }
+
+  private onSnapshotsInSyncInternal(
+    observer: PartialObserver<void>
+  ): Unsubscribe {
+    const errHandler = (err: Error): void => {
+      throw fail('Uncaught Error in onSnapshotsInSync');
+    };
+    const asyncObserver = new AsyncObserver<void>({
+      next: () => {
+        if (observer.next) {
+          observer.next();
+        }
+      },
+      error: errHandler
+    });
+    this._firestoreClient!.addSnapshotsInSyncListener(asyncObserver);
+    return () => {
+      asyncObserver.mute();
+      this._firestoreClient!.removeSnapshotsInSyncListener(asyncObserver);
+    };
+  }
+
   ensureClientConfigured(): FirestoreClient {
     if (!this._firestoreClient) {
       // Kick off starting the client but don't actually wait for it.

packages/firestore/src/core/event_manager.ts

@@ -53,6 +53,8 @@ export class EventManager implements SyncEngineListener {
 
   private onlineState: OnlineState = OnlineState.Unknown;
 
+  private snapshotsInSyncListeners: Set<Observer<void>> = new Set();
+
   constructor(private syncEngine: SyncEngine) {
     this.syncEngine.subscribe(this);
   }
@@ -69,10 +71,18 @@ export class EventManager implements SyncEngineListener {
     }
     queryInfo.listeners.push(listener);
 
-    listener.applyOnlineStateChange(this.onlineState);
+    // Run global snapshot listeners if a consistent snapshot has been emitted.
+    const raisedEvent = listener.applyOnlineStateChange(this.onlineState);
+    assert(
+      !raisedEvent,
+      "applyOnlineStateChange() shouldn't raise an event for brand-new listeners."
+    );
 
     if (queryInfo.viewSnap) {
-      listener.onViewSnapshot(queryInfo.viewSnap);
+      const raisedEvent = listener.onViewSnapshot(queryInfo.viewSnap);
+      if (raisedEvent) {
+        this.raiseSnapshotsInSyncEvent();
+      }
     }
 
     if (firstListen) {
@@ -105,16 +115,22 @@ export class EventManager implements SyncEngineListener {
   }
 
   onWatchChange(viewSnaps: ViewSnapshot[]): void {
+    let raisedEvent = false;
     for (const viewSnap of viewSnaps) {
       const query = viewSnap.query;
       const queryInfo = this.queries.get(query);
       if (queryInfo) {
         for (const listener of queryInfo.listeners) {
-          listener.onViewSnapshot(viewSnap);
+          if (listener.onViewSnapshot(viewSnap)) {
+            raisedEvent = true;
+          }
         }
         queryInfo.viewSnap = viewSnap;
       }
     }
+    if (raisedEvent) {
+      this.raiseSnapshotsInSyncEvent();
+    }
   }
 
   onWatchError(query: Query, error: Error): void {
@@ -132,11 +148,36 @@ export class EventManager implements SyncEngineListener {
 
   onOnlineStateChange(onlineState: OnlineState): void {
     this.onlineState = onlineState;
+    let raisedEvent = false;
     this.queries.forEach((_, queryInfo) => {
       for (const listener of queryInfo.listeners) {
-        listener.applyOnlineStateChange(onlineState);
+        // Run global snapshot listeners if a consistent snapshot has been emitted.
+        if (listener.applyOnlineStateChange(onlineState)) {
+          raisedEvent = true;
+        }
       }
     });
+    if (raisedEvent) {
+      this.raiseSnapshotsInSyncEvent();
+    }
+  }
+
+  addSnapshotsInSyncListener(observer: Observer<void>): void {
+    this.snapshotsInSyncListeners.add(observer);
+    // Immediately fire an initial event, indicating all existing listeners
+    // are in-sync.
+    observer.next();
+  }
+
+  removeSnapshotsInSyncListener(observer: Observer<void>): void {
+    this.snapshotsInSyncListeners.delete(observer);
+  }
+
+  // Call all global snapshot listeners that have been set.
+  private raiseSnapshotsInSyncEvent(): void {
+    this.snapshotsInSyncListeners.forEach(observer => {
+      observer.next();
+    });
   }
 }
 
@@ -178,7 +219,13 @@ export class QueryListener {
     this.options = options || {};
   }
 
-  onViewSnapshot(snap: ViewSnapshot): void {
+  /**
+   * Applies the new ViewSnapshot to this listener, raising a user-facing event
+   * if applicable (depending on what changed, whether the user has opted into
+   * metadata-only changes, etc.). Returns true if a user-facing event was
+   * indeed raised.
+   */
+  onViewSnapshot(snap: ViewSnapshot): boolean {
     assert(
       snap.docChanges.length > 0 || snap.syncStateChanged,
       'We got a new snapshot with no changes?'
@@ -203,31 +250,38 @@ export class QueryListener {
         /* excludesMetadataChanges= */ true
       );
     }
-
+    let raisedEvent = false;
     if (!this.raisedInitialEvent) {
       if (this.shouldRaiseInitialEvent(snap, this.onlineState)) {
         this.raiseInitialEvent(snap);
+        raisedEvent = true;
       }
     } else if (this.shouldRaiseEvent(snap)) {
       this.queryObserver.next(snap);
+      raisedEvent = true;
     }
 
     this.snap = snap;
+    return raisedEvent;
   }
 
   onError(error: Error): void {
     this.queryObserver.error(error);
   }
 
-  applyOnlineStateChange(onlineState: OnlineState): void {
+  /** Returns whether a snapshot was raised. */
+  applyOnlineStateChange(onlineState: OnlineState): boolean {
     this.onlineState = onlineState;
+    let raisedEvent = false;
     if (
       this.snap &&
       !this.raisedInitialEvent &&
       this.shouldRaiseInitialEvent(this.snap, onlineState)
     ) {
       this.raiseInitialEvent(this.snap);
+      raisedEvent = true;
     }
+    return raisedEvent;
   }
 
   private shouldRaiseInitialEvent(

packages/firestore/src/core/firestore_client.ts

@@ -612,6 +612,23 @@ export class FirestoreClient {
     return this.databaseInfo.databaseId;
   }
 
+  addSnapshotsInSyncListener(observer: Observer<void>): void {
+    this.verifyNotTerminated();
+    this.asyncQueue.enqueueAndForget(() => {
+      this.eventMgr.addSnapshotsInSyncListener(observer);
+      return Promise.resolve();
+    });
+  }
+
+  removeSnapshotsInSyncListener(observer: Observer<void>): void {
+    // Checks for shutdown but does not raise error, allowing remove after
+    // shutdown to be a no-op.
+    if (this.clientTerminated) {
+      return;
+    }
+    this.eventMgr.removeSnapshotsInSyncListener(observer);
+  }
+
   get clientTerminated(): boolean {
     // Technically, the asyncQueue is still running, but only accepting operations
     // related to termination or supposed to be run after termination. It is effectively

packages/firestore/test/integration/api/database.test.ts

@@ -37,6 +37,7 @@ import {
   withTestDoc,
   withTestDocAndInitialData,
   DEFAULT_SETTINGS,
+  onSnapshotsInSync,
   withMockCredentialProviderTestDb
 } from '../util/helpers';
 import { User } from '../../../src/auth/user';
@@ -541,6 +542,43 @@ apiDescribe('Database', (persistence: boolean) => {
     });
   });
 
+  it('onSnapshotsInSync fires after listeners are in sync', () => {
+    const testDocs = {
+      a: { foo: 1 }
+    };
+    return withTestCollection(persistence, testDocs, async coll => {
+      let events: string[] = [];
+      const gotInitialSnapshot = new Deferred<void>();
+      const doc = coll.doc('a');
+
+      doc.onSnapshot(snap => {
+        events.push('doc');
+        gotInitialSnapshot.resolve();
+      });
+      await gotInitialSnapshot.promise;
+      events = [];
+
+      const done = new Deferred<void>();
+      onSnapshotsInSync(doc.firestore, () => {
+        events.push('snapshots-in-sync');
+        if (events.length === 3) {
+          // We should have an initial snapshots-in-sync event, then a snapshot
+          // event for set(), then another event to indicate we're in sync
+          // again.
+          expect(events).to.deep.equal([
+            'snapshots-in-sync',
+            'doc',
+            'snapshots-in-sync'
+          ]);
+          done.resolve();
+        }
+      });
+
+      await doc.set({ foo: 3 });
+      await done.promise;
+    });
+  });
+
   apiDescribe('Queries are validated client-side', (persistence: boolean) => {
     // NOTE: Failure cases are validated in validation_test.ts
 

packages/firestore/test/integration/api/validation.test.ts

@@ -29,7 +29,8 @@ import {
   inOp,
   withAlternateTestDb,
   withTestCollection,
-  withTestDb
+  withTestDb,
+  onSnapshotsInSync
 } from '../util/helpers';
 
 // tslint:disable:no-floating-promises
@@ -335,6 +336,11 @@ apiDescribe('Validation:', (persistence: boolean) => {
       `Unknown option 'bad' passed to function ` +
         `Query.onSnapshot(). Available options: includeMetadataChanges`
     );
+
+    expect(() => onSnapshotsInSync(db, 'bad')).to.throw(
+      `Function Firestore.onSnapshotsInSync() requires its first ` +
+        `argument to be of type function, but it was: "bad"`
+    );
   });
 
   validationIt(persistence, 'get options are validated', db => {

packages/firestore/test/integration/util/helpers.ts

@@ -358,6 +358,15 @@ function wipeDb(db: firestore.FirebaseFirestore): Promise<void> {
   return Promise.resolve(undefined);
 }
 
+// TODO(b/139890752): Remove helper and use public API once this is launched.
+export function onSnapshotsInSync(
+  db: firestore.FirebaseFirestore,
+  onSync: unknown
+): () => void {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  return (db as any)._onSnapshotsInSync(onSync);
+}
+
 // TODO(in-queries): This exists just so we don't have to do the cast
 // repeatedly. Once we expose 'array-contains-any' publicly we can remove it and
 // just use 'array-contains-any' in all the tests.