In FSTLocalDocumentsView, search for all batches affecting a set of keys in one go

Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm

@@ -17,6 +17,9 @@
 #import <FirebaseFirestore/FirebaseFirestore.h>
 
 #import <XCTest/XCTest.h>
+#include <mach/mach.h>
+
+#include <cstdint>
 
 #import "Firestore/Example/Tests/Util/FSTEventAccumulator.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
@@ -323,4 +326,54 @@ - (void)testUpdateNestedFields {
   [self awaitExpectations];
 }
 
+// Returns how much memory the test application is currently using, in megabytes (fractional part is
+// truncated), or -1 if the OS call fails.
+// TODO(varconst): move the helper function and the test into a new test target for performance
+// testing.
+std::int64_t GetCurrentMemoryUsedInMb() {
+  mach_task_basic_info taskInfo;
+  mach_msg_type_number_t taskInfoSize = MACH_TASK_BASIC_INFO_COUNT;
+  const auto errorCode =
+      task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&taskInfo, &taskInfoSize);
+  if (errorCode == KERN_SUCCESS) {
+    const int bytesInMegabyte = 1000 * 1000;
+    return taskInfo.resident_size / bytesInMegabyte;
+  }
+  return -1;
+}
+
+- (void)testReasonableMemoryUsageForLotsOfMutations {
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"testReasonableMemoryUsageForLotsOfMutations"];
+
+  FIRDocumentReference *mainDoc = [self documentRef];
+  FIRWriteBatch *batch = [mainDoc.firestore batch];
+
+  // >= 500 mutations will be rejected, so use 500-1 mutations
+  for (int i = 0; i != 500 - 1; ++i) {
+    FIRDocumentReference *nestedDoc = [[mainDoc collectionWithPath:@"nested"] documentWithAutoID];
+    // The exact data doesn't matter; what is important is the large number of mutations.
+    [batch setData:@{
+      @"a" : @"foo",
+      @"b" : @"bar",
+    }
+        forDocument:nestedDoc];
+  }
+
+  const std::int64_t memoryUsedBeforeCommitMb = GetCurrentMemoryUsedInMb();
+  XCTAssertNotEqual(memoryUsedBeforeCommitMb, -1);
+  [batch commitWithCompletion:^(NSError *_Nullable error) {
+    XCTAssertNil(error);
+    const std::int64_t memoryUsedAfterCommitMb = GetCurrentMemoryUsedInMb();
+    XCTAssertNotEqual(memoryUsedAfterCommitMb, -1);
+    const std::int64_t memoryDeltaMb = memoryUsedAfterCommitMb - memoryUsedBeforeCommitMb;
+    // This by its nature cannot be a precise value. In debug mode, the increase on simulator
+    // seems to be around 90 MB. A regression would be on the scale of 500Mb.
+    XCTAssertLessThan(memoryDeltaMb, 150);
+
+    [expectation fulfill];
+  }];
+  [self awaitExpectations];
+}
+
 @end

Firestore/Source/Local/FSTLocalDocumentsView.mm

@@ -62,21 +62,32 @@ - (instancetype)initWithRemoteDocumentCache:(id<FSTRemoteDocumentCache>)remoteDo
 }
 
 - (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key {
+  NSArray<FSTMutationBatch *> *batches =
+      [self.mutationQueue allMutationBatchesAffectingDocumentKey:key];
+  return [self documentForKey:key inBatches:batches];
+}
+
+// Internal version of documentForKey: which allows reusing `batches`.
+- (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key
+                                    inBatches:(NSArray<FSTMutationBatch *> *)batches {
   FSTMaybeDocument *_Nullable remoteDoc = [self.remoteDocumentCache entryForKey:key];
-  return [self localDocument:remoteDoc key:key];
+  return [self localDocument:remoteDoc key:key inBatches:batches];
 }
 
 - (FSTMaybeDocumentDictionary *)documentsForKeys:(const DocumentKeySet &)keys {
   FSTMaybeDocumentDictionary *results = [FSTMaybeDocumentDictionary maybeDocumentDictionary];
+  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];
+    FSTMaybeDocument *maybeDoc = [self documentForKey:key inBatches:batches];
     // TODO(http://b/32275378): Don't conflate missing / deleted.
     if (!maybeDoc) {
       maybeDoc = [FSTDeletedDocument documentWithKey:key version:SnapshotVersion::None()];
     }
     results = [results dictionaryBySettingObject:maybeDoc forKey:key];
   }
+
   return results;
 }
 
@@ -122,12 +133,13 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
   }
 
   // Now add in results for the matchingKeys.
-  for (const DocumentKey &key : matchingKeys) {
-    FSTMaybeDocument *doc = [self documentForKey:key];
-    if ([doc isKindOfClass:[FSTDocument class]]) {
-      results = [results dictionaryBySettingObject:(FSTDocument *)doc forKey:key];
-    }
-  }
+  FSTMaybeDocumentDictionary *matchingKeysDocs = [self documentsForKeys:matchingKeys];
+  [matchingKeysDocs
+      enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTMaybeDocument *doc, BOOL *stop) {
+        if ([doc isKindOfClass:[FSTDocument class]]) {
+          results = [results dictionaryBySettingObject:(FSTDocument *)doc forKey:key];
+        }
+      }];
 
   // Finally, filter out any documents that don't actually match the query. Note that the extra
   // reference here prevents ARC from deallocating the initial unfiltered results while we're
@@ -151,9 +163,8 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
  * @param documentKey The key of the document (necessary when remoteDocument is nil).
  */
 - (nullable FSTMaybeDocument *)localDocument:(nullable FSTMaybeDocument *)document
-                                         key:(const DocumentKey &)documentKey {
-  NSArray<FSTMutationBatch *> *batches =
-      [self.mutationQueue allMutationBatchesAffectingDocumentKey:documentKey];
+                                         key:(const DocumentKey &)documentKey
+                                   inBatches:(NSArray<FSTMutationBatch *> *)batches {
   for (FSTMutationBatch *batch in batches) {
     document = [batch applyTo:document documentKey:documentKey];
   }
@@ -169,10 +180,18 @@ - (nullable FSTMaybeDocument *)localDocument:(nullable FSTMaybeDocument *)docume
  * @return The local view of the documents.
  */
 - (FSTDocumentDictionary *)localDocuments:(FSTDocumentDictionary *)documents {
+  __block DocumentKeySet keySet;
+  [documents
+      enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTDocument *doc, BOOL *stop) {
+        keySet = keySet.insert(doc.key);
+      }];
+  NSArray<FSTMutationBatch *> *batches =
+      [self.mutationQueue allMutationBatchesAffectingDocumentKeys:keySet];
+
   __block FSTDocumentDictionary *result = documents;
   [documents enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTDocument *remoteDocument,
                                                  BOOL *stop) {
-    FSTMaybeDocument *mutatedDoc = [self localDocument:remoteDocument key:key];
+    FSTMaybeDocument *mutatedDoc = [self localDocument:remoteDocument key:key inBatches:batches];
     if ([mutatedDoc isKindOfClass:[FSTDeletedDocument class]]) {
       result = [result dictionaryByRemovingObjectForKey:key];
     } else if ([mutatedDoc isKindOfClass:[FSTDocument class]]) {