Add LocalStore implementation for Bundles (#2342)

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java

@@ -24,6 +24,8 @@
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.auth.User;
+import com.google.firebase.firestore.bundle.BundleMetadata;
+import com.google.firebase.firestore.bundle.NamedQuery;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.core.TargetIdGenerator;
@@ -32,6 +34,7 @@
 import com.google.firebase.firestore.model.MaybeDocument;
 import com.google.firebase.firestore.model.NoDocument;
 import com.google.firebase.firestore.model.ObjectValue;
+import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
@@ -121,6 +124,9 @@ public final class LocalStore {
   /** Maps a query to the data about that query. */
   private final TargetCache targetCache;
 
+  /** Holds information about the bundles loaded into the SDK. */
+  private final BundleCache bundleCache;
+
   /** Maps a targetId to data about its query. */
   private final SparseArray<TargetData> queryDataByTarget;
 
@@ -135,6 +141,7 @@ public LocalStore(Persistence persistence, QueryEngine queryEngine, User initial
         persistence.isStarted(), "LocalStore was passed an unstarted persistence implementation");
     this.persistence = persistence;
     targetCache = persistence.getTargetCache();
+    bundleCache = persistence.getBundleCache();
     targetIdGenerator = TargetIdGenerator.forTargetCache(targetCache.getHighestTargetId());
     mutationQueue = persistence.getMutationQueue(initialUser);
     remoteDocuments = persistence.getRemoteDocumentCache();
@@ -372,51 +379,18 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> applyRemoteEvent(RemoteEve
             }
           }
 
-          Map<DocumentKey, MaybeDocument> changedDocs = new HashMap<>();
           Map<DocumentKey, MaybeDocument> documentUpdates = remoteEvent.getDocumentUpdates();
           Set<DocumentKey> limboDocuments = remoteEvent.getResolvedLimboDocuments();
-          // Each loop iteration only affects its "own" doc, so it's safe to get all the remote
-          // documents in advance in a single call.
-          Map<DocumentKey, MaybeDocument> existingDocs =
-              remoteDocuments.getAll(documentUpdates.keySet());
-
-          for (Entry<DocumentKey, MaybeDocument> entry : documentUpdates.entrySet()) {
-            DocumentKey key = entry.getKey();
-            MaybeDocument doc = entry.getValue();
-            MaybeDocument existingDoc = existingDocs.get(key);
-
-            // Note: The order of the steps below is important, since we want to ensure that
-            // rejected limbo resolutions (which fabricate NoDocuments with SnapshotVersion.NONE)
-            // never add documents to cache.
-            if (doc instanceof NoDocument && doc.getVersion().equals(SnapshotVersion.NONE)) {
-              // NoDocuments with SnapshotVersion.NONE are used in manufactured events. We remove
-              // these documents from cache since we lost access.
-              remoteDocuments.remove(doc.getKey());
-              changedDocs.put(key, doc);
-            } else if (existingDoc == null
-                || doc.getVersion().compareTo(existingDoc.getVersion()) > 0
-                || (doc.getVersion().compareTo(existingDoc.getVersion()) == 0
-                    && existingDoc.hasPendingWrites())) {
-              hardAssert(
-                  !SnapshotVersion.NONE.equals(remoteEvent.getSnapshotVersion()),
-                  "Cannot add a document when the remote version is zero");
-              remoteDocuments.add(doc, remoteEvent.getSnapshotVersion());
-              changedDocs.put(key, doc);
-            } else {
-              Logger.debug(
-                  "LocalStore",
-                  "Ignoring outdated watch update for %s."
-                      + "Current version: %s  Watch version: %s",
-                  key,
-                  existingDoc.getVersion(),
-                  doc.getVersion());
-            }
 
+          for (DocumentKey key : documentUpdates.keySet()) {
             if (limboDocuments.contains(key)) {
               persistence.getReferenceDelegate().updateLimboDocument(key);
             }
           }
 
+          Map<DocumentKey, MaybeDocument> changedDocs =
+              populateDocumentChanges(documentUpdates, null, remoteEvent.getSnapshotVersion());
+
           // HACK: The only reason we allow snapshot version NONE is so that we can synthesize
           // remote events when we get permission denied errors while trying to resolve the
           // state of a locally cached document that is in limbo.
@@ -434,6 +408,65 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> applyRemoteEvent(RemoteEve
         });
   }
 
+  /**
+   * Populates the remote document cache with documents from backend or a bundle. Returns the
+   * document changes resulting from applying those documents.
+   *
+   * <p>Note: this function will use `documentVersions` if it is defined. When it is not defined, it
+   * resorts to `globalVersion`.
+   *
+   * @param documents Documents to be applied.
+   * @param documentVersions A DocumentKey-to-SnapshotVersion map if documents have their own read
+   *     time.
+   * @param globalVersion A SnapshotVersion representing the read time if all documents have the
+   *     same read time.
+   */
+  private Map<DocumentKey, MaybeDocument> populateDocumentChanges(
+      Map<DocumentKey, MaybeDocument> documents,
+      @Nullable Map<DocumentKey, SnapshotVersion> documentVersions,
+      SnapshotVersion globalVersion) {
+    Map<DocumentKey, MaybeDocument> changedDocs = new HashMap<>();
+
+    // Each loop iteration only affects its "own" doc, so it's safe to get all the remote
+    // documents in advance in a single call.
+    Map<DocumentKey, MaybeDocument> existingDocs = remoteDocuments.getAll(documents.keySet());
+
+    for (Entry<DocumentKey, MaybeDocument> entry : documents.entrySet()) {
+      DocumentKey key = entry.getKey();
+      MaybeDocument doc = entry.getValue();
+      MaybeDocument existingDoc = existingDocs.get(key);
+      SnapshotVersion readTime =
+          documentVersions != null ? documentVersions.get(key) : globalVersion;
+
+      // Note: The order of the steps below is important, since we want to ensure that
+      // rejected limbo resolutions (which fabricate NoDocuments with SnapshotVersion.NONE)
+      // never add documents to cache.
+      if (doc instanceof NoDocument && doc.getVersion().equals(SnapshotVersion.NONE)) {
+        // NoDocuments with SnapshotVersion.NONE are used in manufactured events. We remove
+        // these documents from cache since we lost access.
+        remoteDocuments.remove(doc.getKey());
+        changedDocs.put(key, doc);
+      } else if (existingDoc == null
+          || doc.getVersion().compareTo(existingDoc.getVersion()) > 0
+          || (doc.getVersion().compareTo(existingDoc.getVersion()) == 0
+              && existingDoc.hasPendingWrites())) {
+        hardAssert(
+            !SnapshotVersion.NONE.equals(readTime),
+            "Cannot add a document when the remote version is zero");
+        remoteDocuments.add(doc, readTime);
+        changedDocs.put(key, doc);
+      } else {
+        Logger.debug(
+            "LocalStore",
+            "Ignoring outdated watch update for %s." + "Current version: %s  Watch version: %s",
+            key,
+            existingDoc.getVersion(),
+            doc.getVersion());
+      }
+    }
+    return changedDocs;
+  }
+
   /**
    * Returns true if the newTargetData should be persisted during an update of an active target.
    * TargetData should always be persisted when a target is being released and should not call this
@@ -575,6 +608,102 @@ TargetData getTargetData(Target target) {
     return targetCache.getTargetData(target);
   }
 
+  /**
+   * Returns a boolean indicating if the given bundle has already been loaded and its create time is
+   * newer than the currently loading bundle.
+   */
+  public boolean hasNewerBundle(BundleMetadata bundleMetadata) {
+    return persistence.runTransaction(
+        "Has newer bundle",
+        () -> {
+          BundleMetadata existingMetadata =
+              bundleCache.getBundleMetadata(bundleMetadata.getBundleId());
+          return existingMetadata != null
+              && existingMetadata.getCreateTime().compareTo(bundleMetadata.getCreateTime()) > 0;
+        });
+  }
+
+  /** Saves the given BundleMetadata to local persistence. */
+  public void saveBundle(BundleMetadata bundleMetadata) {
+    persistence.runTransaction(
+        "Save bundle",
+        () -> {
+          bundleCache.saveBundleMetadata(bundleMetadata);
+        });
+  }
+
+  /**
+   * Applies the documents from a bundle to the "ground-state" (remote) documents.
+   *
+   * <p>LocalDocuments are re-calculated if there are remaining mutations in the queue.
+   */
+  public ImmutableSortedMap<DocumentKey, MaybeDocument> applyBundledDocuments(
+      ImmutableSortedMap<DocumentKey, MaybeDocument> documents, String bundleId) {
+    // Allocates a target to hold all document keys from the bundle, such that
+    // they will not get garbage collected right away.
+    TargetData umbrellaTargetData = allocateTarget(newUmbrellaTarget(bundleId));
+
+    return persistence.runTransaction(
+        "Apply bundle documents",
+        () -> {
+          ImmutableSortedSet<DocumentKey> documentKeys = DocumentKey.emptyKeySet();
+          Map<DocumentKey, MaybeDocument> documentMap = new HashMap<>();
+          Map<DocumentKey, SnapshotVersion> versionMap = new HashMap<>();
+
+          for (Entry<DocumentKey, MaybeDocument> entry : documents) {
+            DocumentKey documentKey = entry.getKey();
+            MaybeDocument document = entry.getValue();
+
+            if (document instanceof Document) {
+              documentKeys = documentKeys.insert(documentKey);
+            }
+            documentMap.put(documentKey, document);
+            versionMap.put(documentKey, document.getVersion());
+          }
+
+          targetCache.removeMatchingKeysForTargetId(umbrellaTargetData.getTargetId());
+          targetCache.addMatchingKeys(documentKeys, umbrellaTargetData.getTargetId());
+
+          Map<DocumentKey, MaybeDocument> changedDocs =
+              populateDocumentChanges(documentMap, versionMap, SnapshotVersion.NONE);
+          return localDocuments.getLocalViewOfDocuments(changedDocs);
+        });
+  }
+
+  /** Saves the given NamedQuery to local persistence. */
+  public void saveNamedQuery(NamedQuery namedQuery, ImmutableSortedSet<DocumentKey> documentKeys) {
+    // Allocate a target for the named query such that it can be resumed from associated read time
+    // if users use it to listen.
+    // NOTE: this also means if no corresponding target exists, the new target will remain active
+    // and will not get collected, unless users happen to unlisten the query somehow.
+    TargetData existingTargetData = allocateTarget(namedQuery.getBundledQuery().getTarget());
+    int targetId = existingTargetData.getTargetId();
+
+    persistence.runTransaction(
+        "Saved named query",
+        () -> {
+          // Only update the matching documents if it is newer than what the SDK already has
+          if (namedQuery.getReadTime().compareTo(existingTargetData.getSnapshotVersion()) > 0) {
+            // Update existing target data because the query from the bundle is newer.
+            TargetData newTargetData =
+                existingTargetData.withResumeToken(ByteString.EMPTY, namedQuery.getReadTime());
+            queryDataByTarget.append(targetId, newTargetData);
+
+            targetCache.updateTargetData(newTargetData);
+            targetCache.removeMatchingKeysForTargetId(targetId);
+            targetCache.addMatchingKeys(documentKeys, targetId);
+          }
+
+          bundleCache.saveNamedQuery(namedQuery);
+        });
+  }
+
+  /** Returns the NameQuery associated with queryName or null if not found. */
+  public @Nullable NamedQuery getNamedQuery(String queryName) {
+    return persistence.runTransaction(
+        "Get named query", () -> bundleCache.getNamedQuery(queryName));
+  }
+
   /** Mutable state for the transaction in allocateQuery. */
   private static class AllocateQueryHolder {
     TargetData cached;
@@ -673,4 +802,15 @@ public LruGarbageCollector.Results collectGarbage(LruGarbageCollector garbageCol
     return persistence.runTransaction(
         "Collect garbage", () -> garbageCollector.collect(queryDataByTarget));
   }
+
+  /**
+   * Creates a new target using the given bundle name, which will be used to hold the keys of all
+   * documents from the bundle in query-document mappings. This ensures that the loaded documents do
+   * not get garbage collected right away.
+   */
+  private static Target newUmbrellaTarget(String bundleName) {
+    // It is OK that the path used for the query is not valid, because this will not be read and
+    // queried.
+    return Query.atPath(ResourcePath.fromString("__bundle__/docs/" + bundleName)).toTarget();
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryTargetCache.java

@@ -157,7 +157,8 @@ public void removeMatchingKeys(ImmutableSortedSet<DocumentKey> keys, int targetI
     }
   }
 
-  private void removeMatchingKeysForTargetId(int targetId) {
+  @Override
+  public void removeMatchingKeysForTargetId(int targetId) {
     references.removeReferencesForId(targetId);
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteTargetCache.java

@@ -278,7 +278,8 @@ public void removeMatchingKeys(ImmutableSortedSet<DocumentKey> keys, int targetI
     }
   }
 
-  private void removeMatchingKeysForTargetId(int targetId) {
+  @Override
+  public void removeMatchingKeysForTargetId(int targetId) {
     db.execute("DELETE FROM target_documents WHERE target_id = ?", targetId);
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/TargetCache.java

@@ -108,6 +108,9 @@ interface TargetCache {
   /** Removes the given document keys from the cached query results of the given target ID. */
   void removeMatchingKeys(ImmutableSortedSet<DocumentKey> keys, int targetId);
 
+  /** Removes all the keys in the query results of the given target ID. */
+  void removeMatchingKeysForTargetId(int targetId);
+
   ImmutableSortedSet<DocumentKey> getMatchingKeysForTargetId(int targetId);
 
   /** @return True if the document is part of any target */

firebase-firestore/src/main/java/com/google/firebase/firestore/model/ObjectValue.java

@@ -134,6 +134,11 @@ public ObjectValue.Builder toBuilder() {
     return new ObjectValue.Builder(this);
   }
 
+  @Override
+  public String toString() {
+    return "ObjectValue{" + internalValue + "}";
+  }
+
   /**
    * An ObjectValue.Builder provides APIs to set and delete fields from an ObjectValue. All
    * operations mutate the existing instance.