Port performance optimizations to speed up reading large collections from Android

Firestore/Example/Tests/Local/FSTRemoteDocumentCacheTests.mm

@@ -81,6 +81,37 @@ - (void)testSetAndReadADocument {
   [self setAndReadADocumentAtPath:kDocPath];
 }
 
+- (void)testSetAndReadSeveralDocuments {
+  if (!self.remoteDocumentCache) return;
+
+  self.persistence.run("testSetAndReadSeveralDocuments", [=]() {
+    NSArray<FSTDocument *> *written =
+        @[ [self setTestDocumentAtPath:kDocPath], [self setTestDocumentAtPath:kLongDocPath] ];
+    MaybeDocumentMap read = [self.remoteDocumentCache
+        entriesForKeys:DocumentKeySet{testutil::Key(kDocPath), testutil::Key(kLongDocPath)}];
+    [self expectMap:read hasDocsInArray:written exactly:YES];
+  });
+}
+
+- (void)testSetAndReadSeveralDocumentsIncludingMissingDocument {
+  if (!self.remoteDocumentCache) return;
+
+  self.persistence.run("testSetAndReadSeveralDocumentsIncludingMissingDocument", [=]() {
+    NSArray<FSTDocument *> *written =
+        @[ [self setTestDocumentAtPath:kDocPath], [self setTestDocumentAtPath:kLongDocPath] ];
+    MaybeDocumentMap read =
+        [self.remoteDocumentCache entriesForKeys:DocumentKeySet{
+                                                     testutil::Key(kDocPath),
+                                                     testutil::Key(kLongDocPath),
+                                                     testutil::Key("foo/nonexistent"),
+                                                 }];
+    [self expectMap:read hasDocsInArray:written exactly:NO];
+    auto found = read.find(DocumentKey::FromPathString("foo/nonexistent"));
+    XCTAssertTrue(found != read.end());
+    XCTAssertNil(found->second);
+  });
+}
+
 - (void)testSetAndReadADocumentAtDeepPath {
   if (!self.remoteDocumentCache) return;
 
@@ -144,7 +175,7 @@ - (void)testDocumentsMatchingQuery {
 
     FSTQuery *query = FSTTestQuery("b");
     DocumentMap results = [self.remoteDocumentCache documentsMatchingQuery:query];
-    [self expectMap:results
+    [self expectMap:results.underlying_map()
         hasDocsInArray:@[
           FSTTestDoc("b/1", kVersion, _kDocData, FSTDocumentStateSynced),
           FSTTestDoc("b/2", kVersion, _kDocData, FSTDocumentStateSynced)
@@ -162,16 +193,16 @@ - (FSTDocument *)setTestDocumentAtPath:(const absl::string_view)path {
   return doc;
 }
 
-- (void)expectMap:(const DocumentMap &)map
+- (void)expectMap:(const MaybeDocumentMap &)map
     hasDocsInArray:(NSArray<FSTDocument *> *)expected
            exactly:(BOOL)exactly {
   if (exactly) {
     XCTAssertEqual(map.size(), [expected count]);
   }
   for (FSTDocument *doc in expected) {
     FSTDocument *actual = nil;
-    auto found = map.underlying_map().find(doc.key);
-    if (found != map.underlying_map().end()) {
+    auto found = map.find(doc.key);
+    if (found != map.end()) {
       actual = static_cast<FSTDocument *>(found->second);
     }
     XCTAssertEqualObjects(actual, doc);

Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm

@@ -85,6 +85,24 @@ - (nullable FSTMaybeDocument *)entryForKey:(const DocumentKey &)documentKey {
   }
 }
 
+- (MaybeDocumentMap)entriesForKeys:(const DocumentKeySet &)keys {
+  MaybeDocumentMap results;
+
+  LevelDbRemoteDocumentKey currentKey;
+  auto it = _db.currentTransaction->NewIterator();
+
+  for (const DocumentKey &key : keys) {
+    it->Seek([self remoteDocumentKey:key]);
+    if (!it->Valid() || !currentKey.Decode(it->key()) || currentKey.document_key() != key) {
+      results = results.insert(key, nil);
+    } else {
+      results = results.insert(key, [self decodeMaybeDocument:it->value() withKey:key]);
+    }
+  }
+
+  return results;
+}
+
 - (DocumentMap)documentsMatchingQuery:(FSTQuery *)query {
   DocumentMap results;
 

Firestore/Source/Local/FSTLocalDocumentsView.h

@@ -55,6 +55,13 @@ NS_ASSUME_NONNULL_BEGIN
 - (firebase::firestore::model::MaybeDocumentMap)documentsForKeys:
     (const firebase::firestore::model::DocumentKeySet &)keys;
 
+/**
+ * Similar to `documentsForKeys`, but creates the local view from the given
+ * `baseDocs` without retrieving documents from the local store.
+ */
+- (firebase::firestore::model::MaybeDocumentMap)localViewsForDocuments:
+    (const firebase::firestore::model::MaybeDocumentMap &)baseDocs;
+
 /** Performs a query against the local view of all documents. */
 - (firebase::firestore::model::DocumentMap)documentsMatchingQuery:(FSTQuery *)query;
 

Firestore/Source/Local/FSTLocalDocumentsView.mm

@@ -80,13 +80,48 @@ - (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key
   return document;
 }
 
+// Returns the view of the given `docs` as they would appear after applying all
+// mutations in the given `batches`.
+- (MaybeDocumentMap)applyLocalMutationsToDocuments:(const MaybeDocumentMap &)docs
+                                       fromBatches:(NSArray<FSTMutationBatch *> *)batches {
+  MaybeDocumentMap results;
+
+  for (const auto &kv : docs) {
+    const DocumentKey &key = kv.first;
+    FSTMaybeDocument *localView = kv.second;
+    for (FSTMutationBatch *batch in batches) {
+      localView = [batch applyToLocalDocument:localView documentKey:key];
+    }
+    results = results.insert(key, localView);
+  }
+  return results;
+}
+
 - (MaybeDocumentMap)documentsForKeys:(const DocumentKeySet &)keys {
+  MaybeDocumentMap docs = [self.remoteDocumentCache entriesForKeys:keys];
+  return [self localViewsForDocuments:docs];
+}
+
+/**
+ * Similar to `documentsForKeys`, but creates the local view from the given
+ * `baseDocs` without retrieving documents from the local store.
+ */
+- (MaybeDocumentMap)localViewsForDocuments:(const MaybeDocumentMap &)baseDocs {
   MaybeDocumentMap results;
+
+  DocumentKeySet allKeys;
+  for (const auto &kv : baseDocs) {
+    allKeys = allKeys.insert(kv.first);
+  }
   NSArray<FSTMutationBatch *> *batches =
-      [self.mutationQueue allMutationBatchesAffectingDocumentKeys:keys];
-  for (const DocumentKey &key : keys) {
-    // TODO(mikelehen): PERF: Consider fetching all remote documents at once rather than one-by-one.
-    FSTMaybeDocument *maybeDoc = [self documentForKey:key inBatches:batches];
+      [self.mutationQueue allMutationBatchesAffectingDocumentKeys:allKeys];
+
+  MaybeDocumentMap docs = [self applyLocalMutationsToDocuments:baseDocs fromBatches:batches];
+
+  for (const auto &kv : docs) {
+    const DocumentKey &key = kv.first;
+    FSTMaybeDocument *maybeDoc = kv.second;
+
     // TODO(http://b/32275378): Don't conflate missing / deleted.
     if (!maybeDoc) {
       maybeDoc = [FSTDeletedDocument documentWithKey:key

Firestore/Source/Local/FSTLocalSerializer.mm

@@ -67,7 +67,11 @@ - (FSTPBMaybeDocument *)encodedMaybeDocument:(FSTMaybeDocument *)document {
     proto.hasCommittedMutations = deletedDocument.hasCommittedMutations;
   } else if ([document isKindOfClass:[FSTDocument class]]) {
     FSTDocument *existingDocument = (FSTDocument *)document;
-    proto.document = [self encodedDocument:existingDocument];
+    if (existingDocument.proto != nil) {
+      proto.document = existingDocument.proto;
+    } else {
+      proto.document = [self encodedDocument:existingDocument];
+    }
     proto.hasCommittedMutations = existingDocument.hasCommittedMutations;
   } else if ([document isKindOfClass:[FSTUnknownDocument class]]) {
     FSTUnknownDocument *unknownDocument = (FSTUnknownDocument *)document;

Firestore/Source/Local/FSTLocalStore.mm

@@ -264,14 +264,24 @@ - (MaybeDocumentMap)applyRemoteEvent:(FSTRemoteEvent *)remoteEvent {
       }
     }
 
-    // TODO(klimt): This could probably be an NSMutableDictionary.
-    DocumentKeySet changedDocKeys;
+    MaybeDocumentMap changedDocs;
     const DocumentKeySet &limboDocuments = remoteEvent.limboDocumentChanges;
+    DocumentKeySet updatedKeys;
+    for (const auto &kv : remoteEvent.documentUpdates) {
+      updatedKeys = updatedKeys.insert(kv.first);
+    }
+    // Each loop iteration only affects its "own" doc, so it's safe to get all the remote
+    // documents in advance in a single call.
+    MaybeDocumentMap existingDocs = [self.remoteDocumentCache entriesForKeys:updatedKeys];
+
     for (const auto &kv : remoteEvent.documentUpdates) {
       const DocumentKey &key = kv.first;
       FSTMaybeDocument *doc = kv.second;
-      changedDocKeys = changedDocKeys.insert(key);
-      FSTMaybeDocument *existingDoc = [self.remoteDocumentCache entryForKey:key];
+      FSTMaybeDocument *existingDoc = nil;
+      auto foundExisting = existingDocs.find(key);
+      if (foundExisting != existingDocs.end()) {
+        existingDoc = foundExisting->second;
+      }
 
       // If a document update isn't authoritative, make sure we don't apply an old document version
       // to the remote cache. We make an exception for SnapshotVersion.MIN which can happen for
@@ -280,6 +290,7 @@ - (MaybeDocumentMap)applyRemoteEvent:(FSTRemoteEvent *)remoteEvent {
           (authoritativeUpdates.contains(doc.key) && !existingDoc.hasPendingWrites) ||
           doc.version >= existingDoc.version) {
         [self.remoteDocumentCache addEntry:doc];
+        changedDocs = changedDocs.insert(key, doc);
       } else {
         LOG_DEBUG(
             "FSTLocalStore Ignoring outdated watch update for %s. "
@@ -306,7 +317,7 @@ - (MaybeDocumentMap)applyRemoteEvent:(FSTRemoteEvent *)remoteEvent {
       [self.queryCache setLastRemoteSnapshotVersion:remoteVersion];
     }
 
-    return [self.localDocuments documentsForKeys:changedDocKeys];
+    return [self.localDocuments localViewsForDocuments:changedDocs];
   });
 }
 

Firestore/Source/Local/FSTMemoryRemoteDocumentCache.mm

@@ -68,6 +68,16 @@ - (nullable FSTMaybeDocument *)entryForKey:(const DocumentKey &)key {
   return found != _docs.end() ? found->second : nil;
 }
 
+- (MaybeDocumentMap)entriesForKeys:(const DocumentKeySet &)keys {
+  MaybeDocumentMap results;
+  for (const DocumentKey &key : keys) {
+    // Make sure each key has a corresponding entry, which is null in case the document is not
+    // found.
+    results = results.insert(key, [self entryForKey:key]);
+  }
+  return results;
+}
+
 - (DocumentMap)documentsMatchingQuery:(FSTQuery *)query {
   DocumentMap result;
 

Firestore/Source/Local/FSTRemoteDocumentCache.h

@@ -56,6 +56,16 @@ NS_ASSUME_NONNULL_BEGIN
 - (nullable FSTMaybeDocument *)entryForKey:
     (const firebase::firestore::model::DocumentKey &)documentKey;
 
+/**
+ * Looks up a set of entries in the cache.
+ *
+ * @param documentKeys The keys of the entries to look up.
+ * @return The cached Document or NoDocument entries indexed by key. If an entry is not cached,
+ *     the corresponding key will be mapped to a null value.
+ */
+- (firebase::firestore::model::MaybeDocumentMap)entriesForKeys:
+    (const firebase::firestore::model::DocumentKeySet &)documentKeys;
+
 /**
  * Executes a query against the cached FSTDocument entries
  *

Firestore/Source/Model/FSTDocument.h

@@ -21,6 +21,7 @@
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTFieldValue;
+@class GCFSDocument;
 @class FSTObjectValue;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -57,12 +58,24 @@ typedef NS_ENUM(NSInteger, FSTDocumentState) {
                          version:(firebase::firestore::model::SnapshotVersion)version
                            state:(FSTDocumentState)state;
 
++ (instancetype)documentWithData:(FSTObjectValue *)data
+                             key:(firebase::firestore::model::DocumentKey)key
+                         version:(firebase::firestore::model::SnapshotVersion)version
+                           state:(FSTDocumentState)state
+                           proto:(GCFSDocument *)proto;
+
 - (nullable FSTFieldValue *)fieldForPath:(const firebase::firestore::model::FieldPath &)path;
 - (BOOL)hasLocalMutations;
 - (BOOL)hasCommittedMutations;
 
 @property(nonatomic, strong, readonly) FSTObjectValue *data;
 
+/**
+ * Memoized serialized form of the document for optimization purposes (avoids repeated
+ * serialization). Might be nil.
+ */
+@property(nonatomic, strong, readonly) GCFSDocument *proto;
+
 @end
 
 @interface FSTDeletedDocument : FSTMaybeDocument

Firestore/Source/Model/FSTDocument.mm

@@ -88,6 +88,18 @@ + (instancetype)documentWithData:(FSTObjectValue *)data
                                      state:state];
 }
 
++ (instancetype)documentWithData:(FSTObjectValue *)data
+                             key:(DocumentKey)key
+                         version:(SnapshotVersion)version
+                           state:(FSTDocumentState)state
+                           proto:(GCFSDocument *)proto {
+  return [[FSTDocument alloc] initWithData:data
+                                       key:std::move(key)
+                                   version:std::move(version)
+                                     state:state
+                                     proto:proto];
+}
+
 - (instancetype)initWithData:(FSTObjectValue *)data
                          key:(DocumentKey)key
                      version:(SnapshotVersion)version
@@ -96,6 +108,21 @@ - (instancetype)initWithData:(FSTObjectValue *)data
   if (self) {
     _data = data;
     _documentState = state;
+    _proto = nil;
+  }
+  return self;
+}
+
+- (instancetype)initWithData:(FSTObjectValue *)data
+                         key:(DocumentKey)key
+                     version:(SnapshotVersion)version
+                       state:(FSTDocumentState)state
+                       proto:(GCFSDocument *)proto {
+  self = [super initWithKey:std::move(key) version:std::move(version)];
+  if (self) {
+    _data = data;
+    _documentState = state;
+    _proto = proto;
   }
   return self;
 }

Firestore/Source/Remote/FSTSerializerBeta.mm

@@ -438,7 +438,11 @@ - (FSTDocument *)decodedFoundDocument:(GCFSBatchGetDocumentsResponse *)response
   HARD_ASSERT(version != SnapshotVersion::None(),
               "Got a document response with no snapshot version");
 
-  return [FSTDocument documentWithData:value key:key version:version state:FSTDocumentStateSynced];
+  return [FSTDocument documentWithData:value
+                                   key:key
+                               version:version
+                                 state:FSTDocumentStateSynced
+                                 proto:response.found];
 }
 
 - (FSTDeletedDocument *)decodedDeletedDocument:(GCFSBatchGetDocumentsResponse *)response {
@@ -1144,8 +1148,13 @@ - (FSTDocumentWatchChange *)decodedDocumentChange:(GCFSDocumentChange *)change {
   const DocumentKey key = [self decodedDocumentKey:change.document.name];
   SnapshotVersion version = [self decodedVersion:change.document.updateTime];
   HARD_ASSERT(version != SnapshotVersion::None(), "Got a document change with no snapshot version");
-  FSTMaybeDocument *document =
-      [FSTDocument documentWithData:value key:key version:version state:FSTDocumentStateSynced];
+  // The document may soon be re-serialized back to protos in order to store it in local
+  // persistence. Memoize the encoded form to avoid encoding it again.
+  FSTMaybeDocument *document = [FSTDocument documentWithData:value
+                                                         key:key
+                                                     version:version
+                                                       state:FSTDocumentStateSynced
+                                                       proto:change.document];
 
   NSArray<NSNumber *> *updatedTargetIds = [self decodedIntegerArray:change.targetIdsArray];
   NSArray<NSNumber *> *removedTargetIds = [self decodedIntegerArray:change.removedTargetIdsArray];

Firestore/core/src/firebase/firestore/remote/watch_stream.mm

@@ -84,8 +84,10 @@
     return status;
   }
 
-  LOG_DEBUG("%s response: %s", GetDebugDescription(),
-            serializer_bridge_.Describe(response));
+  if (bridge::IsLoggingEnabled()) {
+    LOG_DEBUG("%s response: %s", GetDebugDescription(),
+              serializer_bridge_.Describe(response));
+  }
 
   // A successful response means the stream is healthy.
   backoff_.Reset();