[2/3] Query/Target split: make SyncEngine able to handle the mapping between Target and Query. (#2281)

wu-hui · web-flow · commit ca58a2939bf8 · 2019-10-21T15:06:27.000-04:00
diff --git a/packages/firestore/src/core/sync_engine.ts b/packages/firestore/src/core/sync_engine.ts
@@ -147,7 +147,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
   private queryViewsByQuery = new ObjectMap<Query, QueryView>(q =>
     q.canonicalId()
   );
-  private queryViewsByTarget: { [targetId: number]: QueryView } = {};
+  private queriesByTarget: { [targetId: number]: Query[] } = {};
   private limboTargetsByKey = new SortedMap<DocumentKey, TargetId>(
     DocumentKey.comparator
   );
@@ -223,7 +223,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       targetId = queryData.targetId;
       viewSnapshot = await this.initializeViewAndComputeSnapshot(
         query,
-        queryData.targetId,
+        targetId,
         status === 'current'
       );
       if (this.isPrimary) {
@@ -270,7 +270,10 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
 
     const data = new QueryView(query, targetId, view);
     this.queryViewsByQuery.set(query, data);
-    this.queryViewsByTarget[targetId] = data;
+    if (!this.queriesByTarget[targetId]) {
+      this.queriesByTarget[targetId] = [];
+    }
+    this.queriesByTarget[targetId].push(query);
     return viewChange.snapshot!;
   }
 
@@ -302,6 +305,8 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
     const queryView = this.queryViewsByQuery.get(query)!;
     assert(!!queryView, 'Trying to unlisten on query not found:' + query);
 
+    // TODO(wuandy): Note this does not handle the case where multiple queries
+    // map to one target, and user request to unlisten on of the queries.
     if (this.isPrimary) {
       // We need to remove the local query target first to allow us to verify
       // whether any other client is still interested in this target.
@@ -312,18 +317,18 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
 
       if (!targetRemainsActive) {
         await this.localStore
-          .releaseQuery(query, /*keepPersistedQueryData=*/ false)
+          .releaseTarget(queryView.targetId, /*keepPersistedQueryData=*/ false)
           .then(() => {
             this.sharedClientState.clearQueryState(queryView.targetId);
             this.remoteStore.unlisten(queryView.targetId);
-            this.removeAndCleanupQuery(queryView);
+            this.removeAndCleanupTarget(queryView.targetId);
           })
           .catch(ignoreIfPrimaryLeaseLoss);
       }
     } else {
-      this.removeAndCleanupQuery(queryView);
-      await this.localStore.releaseQuery(
-        query,
+      this.removeAndCleanupTarget(queryView.targetId);
+      await this.localStore.releaseTarget(
+        queryView.targetId,
         /*keepPersistedQueryData=*/ true
       );
     }
@@ -495,13 +500,10 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       );
       return this.applyRemoteEvent(event);
     } else {
-      const queryView = this.queryViewsByTarget[targetId];
-      assert(!!queryView, 'Unknown targetId: ' + targetId);
       await this.localStore
-        .releaseQuery(queryView.query, /* keepPersistedQueryData */ false)
-        .then(() => this.removeAndCleanupQuery(queryView))
+        .releaseTarget(targetId, /* keepPersistedQueryData */ false)
+        .then(() => this.removeAndCleanupTarget(targetId, err))
         .catch(ignoreIfPrimaryLeaseLoss);
-      this.syncEngineListener!.onWatchError(queryView.query, err);
     }
   }
 
@@ -681,17 +683,30 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
     }
   }
 
-  private removeAndCleanupQuery(queryView: QueryView): void {
-    this.sharedClientState.removeLocalQueryTarget(queryView.targetId);
+  private removeAndCleanupTarget(
+    targetId: number,
+    error: Error | null = null
+  ): void {
+    this.sharedClientState.removeLocalQueryTarget(targetId);
+
+    assert(
+      this.queriesByTarget[targetId] &&
+        this.queriesByTarget[targetId].length !== 0,
+      `There are no queries mapped to target id ${targetId}`
+    );
+
+    for (const query of this.queriesByTarget[targetId]) {
+      this.queryViewsByQuery.delete(query);
+      if (error) {
+        this.syncEngineListener!.onWatchError(query, error);
+      }
+    }
 
-    this.queryViewsByQuery.delete(queryView.query);
-    delete this.queryViewsByTarget[queryView.targetId];
+    delete this.queriesByTarget[targetId];
 
     if (this.isPrimary) {
-      const limboKeys = this.limboDocumentRefs.referencesForId(
-        queryView.targetId
-      );
-      this.limboDocumentRefs.removeReferencesForId(queryView.targetId);
+      const limboKeys = this.limboDocumentRefs.referencesForId(targetId);
+      this.limboDocumentRefs.removeReferencesForId(targetId);
       limboKeys.forEach(limboKey => {
         const isReferenced = this.limboDocumentRefs.containsKey(limboKey);
         if (!isReferenced) {
@@ -710,6 +725,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       // This target already got removed, because the query failed.
       return;
     }
+
     this.remoteStore.unlisten(limboTargetId);
     this.limboTargetsByKey = this.limboTargetsByKey.remove(key);
     delete this.limboResolutionsByTarget[limboTargetId];
@@ -885,13 +901,19 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       const activeTargets: TargetId[] = [];
 
       let p = Promise.resolve();
-      objUtils.forEachNumber(this.queryViewsByTarget, (targetId, queryView) => {
+      objUtils.forEachNumber(this.queriesByTarget, (targetId, _) => {
         if (this.sharedClientState.isLocalQueryTarget(targetId)) {
           activeTargets.push(targetId);
         } else {
-          p = p.then(() => this.unlisten(queryView.query));
+          p = p.then(() => {
+            this.removeAndCleanupTarget(targetId);
+            return this.localStore.releaseTarget(
+              targetId,
+              /*keepPersistedQueryData=*/ true
+            );
+          });
         }
-        this.remoteStore.unlisten(queryView.targetId);
+        this.remoteStore.unlisten(targetId);
       });
       await p;
 
@@ -926,27 +948,34 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
     const newViewSnapshots: ViewSnapshot[] = [];
     for (const targetId of targets) {
       let queryData: QueryData;
-      const queryView = this.queryViewsByTarget[targetId];
-      if (queryView) {
+      const queries = this.queriesByTarget[targetId];
+
+      if (queries && queries.length !== 0) {
         // For queries that have a local View, we need to update their state
         // in LocalStore (as the resume token and the snapshot version
         // might have changed) and reconcile their views with the persisted
         // state (the list of syncedDocuments may have gotten out of sync).
-        await this.localStore.releaseQuery(
-          queryView.query,
+        await this.localStore.releaseTarget(
+          targetId,
           /*keepPersistedQueryData=*/ true
         );
-        queryData = await this.localStore.allocateQuery(queryView.query);
-        const viewChange = await this.synchronizeViewAndComputeSnapshot(
-          queryView
-        );
-        if (viewChange.snapshot) {
-          newViewSnapshots.push(viewChange.snapshot);
+        queryData = await this.localStore.allocateTarget(queries[0].toTarget());
+
+        for (const query of queries) {
+          const queryView = this.queryViewsByQuery.get(query);
+          assert(!!queryView, `No query view found for ${query}`);
+
+          const viewChange = await this.synchronizeViewAndComputeSnapshot(
+            queryView!
+          );
+          if (viewChange.snapshot) {
+            newViewSnapshots.push(viewChange.snapshot);
+          }
         }
       } else {
         assert(
           this.isPrimary === true,
-          'A secondary tab should never have an active query without an active view.'
+          'A secondary tab should never have an active target without an active query.'
         );
         // For queries that never executed on this client, we need to
         // allocate the target in LocalStore and initialize a new View.
@@ -959,8 +988,10 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
           /*current=*/ false
         );
       }
-      activeQueries.push(queryData);
+
+      activeQueries.push(queryData!);
     }
+
     this.syncEngineListener!.onWatchChange(newViewSnapshots);
     return activeQueries;
   }
@@ -983,7 +1014,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
       return;
     }
 
-    if (this.queryViewsByTarget[targetId]) {
+    if (this.queriesByTarget[targetId]) {
       switch (state) {
         case 'current':
         case 'not-current': {
@@ -999,13 +1030,11 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
           break;
         }
         case 'rejected': {
-          const queryView = this.queryViewsByTarget[targetId];
-          this.removeAndCleanupQuery(queryView);
-          await this.localStore.releaseQuery(
-            queryView.query,
-            /*keepPersistedQueryData=*/ true
+          await this.localStore.releaseTarget(
+            targetId,
+            /* keepPersistedQueryData */ true
           );
-          this.syncEngineListener!.onWatchError(queryView.query, error!);
+          this.removeAndCleanupTarget(targetId, error);
           break;
         }
         default:
@@ -1025,7 +1054,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
 
     for (const targetId of added) {
       assert(
-        !this.queryViewsByTarget[targetId],
+        !this.queriesByTarget[targetId],
         'Trying to add an already active target'
       );
       const target = await this.localStore.getTarget(targetId);
@@ -1040,18 +1069,20 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
     }
 
     for (const targetId of removed) {
-      const queryView = this.queryViewsByTarget[targetId];
-      // Check that the query is still active since the query might have been
+      // Check that the target is still active since the target might have been
       // removed if it has been rejected by the backend.
-      if (queryView) {
-        await this.localStore
-          .releaseQuery(queryView.query, /*keepPersistedQueryData=*/ false)
-          .then(() => {
-            this.remoteStore.unlisten(targetId);
-            this.removeAndCleanupQuery(queryView);
-          })
-          .catch(ignoreIfPrimaryLeaseLoss);
+      if (!this.queriesByTarget[targetId]) {
+        continue;
       }
+
+      // Release queries that are still active.
+      await this.localStore
+        .releaseTarget(targetId, /* keepPersistedQueryData */ false)
+        .then(() => {
+          this.remoteStore.unlisten(targetId);
+          this.removeAndCleanupTarget(targetId);
+        })
+        .catch(ignoreIfPrimaryLeaseLoss);
     }
   }
 
@@ -1074,9 +1105,17 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
     if (limboResolution && limboResolution.receivedDocument) {
       return documentKeySet().add(limboResolution.key);
     } else {
-      return this.queryViewsByTarget[targetId]
-        ? this.queryViewsByTarget[targetId].view.syncedDocuments
-        : documentKeySet();
+      let keySet = documentKeySet();
+      const queries = this.queriesByTarget[targetId];
+      if (!queries) {
+        return keySet;
+      }
+      for (const query of queries) {
+        const queryView = this.queryViewsByQuery.get(query);
+        assert(!!queryView, `No query view found for ${query}`);
+        keySet = keySet.unionWith(queryView!.view.syncedDocuments);
+      }
+      return keySet;
     }
   }
 }
diff --git a/packages/firestore/src/local/local_store.ts b/packages/firestore/src/local/local_store.ts
@@ -885,11 +885,6 @@ export class LocalStore {
     }
   }
 
-  //TODO(wuandy): Delete this temp mothod, it's for change isolation only.
-  releaseQuery(query: Query, keepPersistedQueryData: boolean): Promise<void> {
-    return this.releaseTarget(query.toTarget(), keepPersistedQueryData);
-  }
-
   /**
    * Unpin all the documents associated with the given query. If
    * `keepPersistedQueryData` is set to false and Eager GC enabled, the method
@@ -899,24 +894,17 @@ export class LocalStore {
    */
   // PORTING NOTE: `keepPersistedQueryData` is multi-tab only.
   releaseTarget(
-    target: Target,
+    targetId: number,
     keepPersistedQueryData: boolean
   ): Promise<void> {
-    let targetId: number;
+    const queryData = this.queryDataByTarget.get(targetId)!;
+    assert(!!queryData, 'Tried to release nonexistent target: ' + targetId);
 
     const mode = keepPersistedQueryData
       ? 'readwrite-idempotent'
       : 'readwrite-primary-idempotent';
     return this.persistence
       .runTransaction('Release query', mode, txn => {
-        const cachedTargetId = this.targetIdByTarget.get(target);
-        assert(
-          cachedTargetId !== undefined,
-          'Tried to release nonexistent target: ' + target
-        );
-        targetId = cachedTargetId!;
-        const queryData = this.queryDataByTarget.get(targetId)!;
-
         // References for documents sent via Watch are automatically removed
         // when we delete a query's target data from the reference delegate.
         // Since this does not remove references for locally mutated documents,
@@ -943,7 +931,7 @@ export class LocalStore {
       })
       .then(() => {
         this.queryDataByTarget = this.queryDataByTarget.remove(targetId);
-        this.targetIdByTarget.delete(target);
+        this.targetIdByTarget.delete(queryData.target);
       });
   }
 
diff --git a/packages/firestore/test/unit/local/local_store.test.ts b/packages/firestore/test/unit/local/local_store.test.ts
