Keep Query state consistent when primary lease is lost (#1162)

Author: schmidt-sebastian

packages/firestore/src/core/sync_engine.ts

@@ -315,11 +315,13 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       );
 
       if (!targetRemainsActive) {
-        this.sharedClientState.clearQueryState(queryView.targetId);
-        this.remoteStore.unlisten(queryView.targetId);
         await this.localStore
           .releaseQuery(query, /*keepPersistedQueryData=*/ false)
-          .then(() => this.removeAndCleanupQuery(queryView))
+          .then(() => {
+            this.sharedClientState.clearQueryState(queryView.targetId);
+            this.remoteStore.unlisten(queryView.targetId);
+            return this.removeAndCleanupQuery(queryView);
+          })
           .then(() => this.localStore.collectGarbage())
           .catch(err => this.ignoreIfPrimaryLeaseLoss(err));
       }
@@ -882,14 +884,22 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       }
     } else if (isPrimary === false && this.isPrimary !== false) {
       this.isPrimary = false;
-      const activeQueries = await this.synchronizeQueryViewsAndRaiseSnapshots(
-        objUtils.indices(this.queryViewsByTarget)
-      );
+
+      const activeTargets: TargetId[] = [];
+
+      let p = Promise.resolve();
+      objUtils.forEachNumber(this.queryViewsByTarget, (targetId, queryView) => {
+        if (this.sharedClientState.isLocalQueryTarget(targetId)) {
+          activeTargets.push(targetId);
+        } else {
+          p = p.then(() => this.unlisten(queryView.query));
+        }
+        this.remoteStore.unlisten(queryView.targetId);
+      });
+      await p;
+
+      await this.synchronizeQueryViewsAndRaiseSnapshots(activeTargets);
       this.resetLimboDocuments();
-      for (const queryData of activeQueries) {
-        // TODO(multitab): Remove query views for non-local queries.
-        this.remoteStore.unlisten(queryData.targetId);
-      }
       await this.remoteStore.applyPrimaryState(false);
     }
   }
@@ -1038,10 +1048,12 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       // Check that the query is still active since the query might have been
       // removed if it has been rejected by the backend.
       if (queryView) {
-        this.remoteStore.unlisten(targetId);
         await this.localStore
           .releaseQuery(queryView.query, /*keepPersistedQueryData=*/ false)
-          .then(() => this.removeAndCleanupQuery(queryView))
+          .then(() => {
+            this.remoteStore.unlisten(targetId);
+            return this.removeAndCleanupQuery(queryView);
+          })
           .catch(err => this.ignoreIfPrimaryLeaseLoss(err));
       }
     }

packages/firestore/src/local/shared_client_state.ts

@@ -105,6 +105,9 @@ export interface SharedClientState {
   /** Removes the Query Target ID association from the local client. */
   removeLocalQueryTarget(targetId: TargetId): void;
 
+  /** Checks whether the target is associated with the local client. */
+  isLocalQueryTarget(targetId: TargetId): boolean;
+
   /**
    * Processes an update to a query target.
    *
@@ -700,6 +703,10 @@ export class WebStorageSharedClientState implements SharedClientState {
     this.persistClientState();
   }
 
+  isLocalQueryTarget(targetId: TargetId): boolean {
+    return this.localClientState.activeTargetIds.has(targetId);
+  }
+
   clearQueryState(targetId: TargetId): void {
     this.removeItem(this.toLocalStorageQueryTargetMetadataKey(targetId));
   }
@@ -1086,6 +1093,10 @@ export class MemorySharedClientState implements SharedClientState {
     this.localState.removeQueryTarget(targetId);
   }
 
+  isLocalQueryTarget(targetId: TargetId): boolean {
+    return this.localState.activeTargetIds.has(targetId);
+  }
+
   clearQueryState(targetId: TargetId): void {
     delete this.queryState[targetId];
   }

packages/firestore/src/util/obj.ts

@@ -39,13 +39,6 @@ export function size<V>(obj: Dict<V>): number {
   return count;
 }
 
-/** Extracts the numeric indices from a dictionary. */
-export function indices<V>(obj: { [numberKey: number]: V }): number[] {
-  return Object.keys(obj).map(key => {
-    return Number(key);
-  });
-}
-
 /** Returns the given value if it's defined or the defaultValue otherwise. */
 export function defaulted<V>(value: V | undefined, defaultValue: V): V {
   return value !== undefined ? value : defaultValue;

packages/firestore/test/unit/specs/listen_spec.test.ts

@@ -1054,11 +1054,49 @@ describeSpec('Listens:', [], () => {
         .expectEvents(query, { added: [docC] })
         .client(0)
         .runTimer(TimerId.ClientMetadataRefresh)
+        // Client 0 recovers from its lease loss and applies the updates from
+        // client 1
         .expectPrimaryState(false)
         .expectEvents(query, { added: [docB, docC] })
     );
   });
 
+  specTest('Query bounces between primaries', ['multi-client'], () => {
+    const query = Query.atPath(path('collection'));
+    const docA = doc('collection/a', 1000, { key: 'a' });
+    const docB = doc('collection/b', 2000, { key: 'b' });
+    const docC = doc('collection/c', 3000, { key: 'c' });
+
+    // Client 0 listens to a query. Client 1 is the primary when the query is
+    // first listened to, then the query switches to client 0 and back to client
+    // 1.
+    return client(1)
+      .expectPrimaryState(true)
+      .client(0)
+      .userListens(query)
+      .client(1)
+      .expectListen(query)
+      .watchAcksFull(query, 1000, docA)
+      .client(0)
+      .expectEvents(query, { added: [docA] })
+      .client(2)
+      .stealPrimaryLease()
+      .expectListen(query, 'resume-token-1000')
+      .client(1)
+      .runTimer(TimerId.ClientMetadataRefresh)
+      .expectPrimaryState(false)
+      .client(2)
+      .watchAcksFull(query, 2000, docB)
+      .client(0)
+      .expectEvents(query, { added: [docB] })
+      .client(1)
+      .stealPrimaryLease()
+      .expectListen(query, 'resume-token-2000')
+      .watchAcksFull(query, 2000, docC)
+      .client(0)
+      .expectEvents(query, { added: [docC] });
+  });
+
   specTest(
     'Unresponsive primary ignores watch update',
     ['multi-client'],