Keep track of number of queries in the query cache (#776)

* Implement schema versions

* Style fixes

* newlines, copyrights, assumptions

* Fix nullability

* Raw ptr -> shared_ptr

* kVersionTableGlobal -> kVersionGlobalTable

* Drop utils, move into static methods

* Drop extra include

* Add a few more comments

* Move version constant into migrations file

* formatting?

* Fix comment

* Split add and update queryData

* Work on adding targetCount

* More work on count

* Using shared_ptr

* Implement count for query cache

* use quotes

* Add cast

* Styling

* Revert year bump in copyright

* Add adversarial key to migration test

* Add comment

* Fix style

11 files changed, 159 insertions, 24 deletions

diff --git a/Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm b/Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm
index 8ef0e94..3559d5d 100644
--- a/Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm
+++ b/Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm
@@ -18,13 +18,18 @@
 #include <leveldb/db.h>
 
 #import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
+#import "Firestore/Source/Local/FSTLevelDBKey.h"
 #import "Firestore/Source/Local/FSTLevelDBMigrations.h"
 #import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
+#import "Firestore/Source/Local/FSTWriteGroup.h"
+
+#include "Firestore/core/src/firebase/firestore/util/ordered_code.h"
 
 #import "Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
+using firebase::firestore::util::OrderedCode;
 using leveldb::DB;
 using leveldb::Options;
 using leveldb::Status;
@@ -70,6 +75,30 @@ using leveldb::Status;
   XCTAssertGreaterThan(actual, 0, @"Expected to migrate to a schema version > 0");
 }
 
+- (void)testCountsQueries {
+  NSUInteger expected = 50;
+  FSTWriteGroup *group = [FSTWriteGroup groupWithAction:@"Setup"];
+  for (int i = 0; i < expected; i++) {
+    std::string key = [FSTLevelDBTargetKey keyWithTargetID:i];
+    [group setData:"dummy" forKey:key];
+  }
+  // Add a dummy entry after the targets to make sure the iteration is correctly bounded.
+  // Use a table that would sort logically right after that table 'target'.
+  std::string dummyKey;
+  // Magic number that indicates a table name follows. Needed to mimic the prefix to the target
+  // table.
+  OrderedCode::WriteSignedNumIncreasing(&dummyKey, 5);
+  OrderedCode::WriteString(&dummyKey, "targetA");
+  [group setData:"dummy" forKey:dummyKey];
+
+  Status status = [group writeToDB:_db];
+  XCTAssertTrue(status.ok(), @"Failed to write targets");
+
+  [FSTLevelDBMigrations runMigrationsOnDB:_db];
+  FSTPBTargetGlobal *metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db];
+  XCTAssertEqual(expected, metadata.targetCount, @"Failed to count all of the targets we added");
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
index 0c6a2a4..6ab655a 100644
--- a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
+++ b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
@@ -89,6 +89,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   FSTQueryData *data2 = [self queryDataWithQuery:q2];
   [self addQueryData:data2];
+  XCTAssertEqual(2, [self.queryCache count]);
 
   XCTAssertEqualObjects([self.queryCache queryDataForQuery:q1], data1);
   XCTAssertEqualObjects([self.queryCache queryDataForQuery:q2], data2);
@@ -96,10 +97,12 @@ NS_ASSUME_NONNULL_BEGIN
   [self removeQueryData:data1];
   XCTAssertNil([self.queryCache queryDataForQuery:q1]);
   XCTAssertEqualObjects([self.queryCache queryDataForQuery:q2], data2);
+  XCTAssertEqual(1, [self.queryCache count]);
 
   [self removeQueryData:data2];
   XCTAssertNil([self.queryCache queryDataForQuery:q1]);
   XCTAssertNil([self.queryCache queryDataForQuery:q2]);
+  XCTAssertEqual(0, [self.queryCache count]);
 }
 
 - (void)testSetQueryToNewValue {
diff --git a/Firestore/Protos/FrameworkMaker.xcodeproj/project.pbxproj b/Firestore/Protos/FrameworkMaker.xcodeproj/project.pbxproj
index 51a61b8..2efcb21 100644
--- a/Firestore/Protos/FrameworkMaker.xcodeproj/project.pbxproj
+++ b/Firestore/Protos/FrameworkMaker.xcodeproj/project.pbxproj
@@ -201,7 +201,7 @@
 			);
 			inputPaths = (
 				"${SRCROOT}/Pods/Target Support Files/Pods-FrameworkMaker_iOS/Pods-FrameworkMaker_iOS-resources.sh",
-				"$PODS_CONFIGURATION_BUILD_DIR/gRPC/gRPCCertificates.bundle",
+				"${PODS_CONFIGURATION_BUILD_DIR}/gRPC/gRPCCertificates.bundle",
 			);
 			name = "[CP] Copy Pods Resources";
 			outputPaths = (
diff --git a/Firestore/Protos/objc/firestore/local/Target.pbobjc.h b/Firestore/Protos/objc/firestore/local/Target.pbobjc.h
index d8bf49c..0672a6e 100644
--- a/Firestore/Protos/objc/firestore/local/Target.pbobjc.h
+++ b/Firestore/Protos/objc/firestore/local/Target.pbobjc.h
@@ -160,6 +160,7 @@ typedef GPB_ENUM(FSTPBTargetGlobal_FieldNumber) {
   FSTPBTargetGlobal_FieldNumber_HighestTargetId = 1,
   FSTPBTargetGlobal_FieldNumber_HighestListenSequenceNumber = 2,
   FSTPBTargetGlobal_FieldNumber_LastRemoteSnapshotVersion = 3,
+  FSTPBTargetGlobal_FieldNumber_TargetCount = 4,
 };
 
 /**
@@ -197,6 +198,9 @@ typedef GPB_ENUM(FSTPBTargetGlobal_FieldNumber) {
 /** Test to see if @c lastRemoteSnapshotVersion has been set. */
 @property(nonatomic, readwrite) BOOL hasLastRemoteSnapshotVersion;
 
+/** On platforms that need it, holds the number of targets persisted. */
+@property(nonatomic, readwrite) int32_t targetCount;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Protos/objc/firestore/local/Target.pbobjc.m b/Firestore/Protos/objc/firestore/local/Target.pbobjc.m
index 6f6ccf2..567c86d 100644
--- a/Firestore/Protos/objc/firestore/local/Target.pbobjc.m
+++ b/Firestore/Protos/objc/firestore/local/Target.pbobjc.m
@@ -183,10 +183,12 @@ void FSTPBTarget_ClearTargetTypeOneOfCase(FSTPBTarget *message) {
 @dynamic highestTargetId;
 @dynamic highestListenSequenceNumber;
 @dynamic hasLastRemoteSnapshotVersion, lastRemoteSnapshotVersion;
+@dynamic targetCount;
 
 typedef struct FSTPBTargetGlobal__storage_ {
   uint32_t _has_storage_[1];
   int32_t highestTargetId;
+  int32_t targetCount;
   GPBTimestamp *lastRemoteSnapshotVersion;
   int64_t highestListenSequenceNumber;
 } FSTPBTargetGlobal__storage_;
@@ -224,6 +226,15 @@ typedef struct FSTPBTargetGlobal__storage_ {
         .flags = GPBFieldOptional,
         .dataType = GPBDataTypeMessage,
       },
+      {
+        .name = "targetCount",
+        .dataTypeSpecific.className = NULL,
+        .number = FSTPBTargetGlobal_FieldNumber_TargetCount,
+        .hasIndex = 3,
+        .offset = (uint32_t)offsetof(FSTPBTargetGlobal__storage_, targetCount),
+        .flags = GPBFieldOptional,
+        .dataType = GPBDataTypeInt32,
+      },
     };
     GPBDescriptor *localDescriptor =
         [GPBDescriptor allocDescriptorForClass:[FSTPBTargetGlobal class]
diff --git a/Firestore/Protos/protos/firestore/local/target.proto b/Firestore/Protos/protos/firestore/local/target.proto
index 7f34515..7f0a886 100644
--- a/Firestore/Protos/protos/firestore/local/target.proto
+++ b/Firestore/Protos/protos/firestore/local/target.proto
@@ -87,4 +87,7 @@ message TargetGlobal {
   // This is updated whenever our we get a TargetChange with a read_time and
   // empty target_ids.
   google.protobuf.Timestamp last_remote_snapshot_version = 3;
+
+  // On platforms that need it, holds the number of targets persisted.
+  int32 target_count = 4;
 }
diff --git a/Firestore/Source/Local/FSTLevelDBMigrations.mm b/Firestore/Source/Local/FSTLevelDBMigrations.mm
index 49af893..7595c53 100644
--- a/Firestore/Source/Local/FSTLevelDBMigrations.mm
+++ b/Firestore/Source/Local/FSTLevelDBMigrations.mm
@@ -16,21 +16,22 @@
 
 #include "Firestore/Source/Local/FSTLevelDBMigrations.h"
 
-#include <leveldb/db.h>
-#include <leveldb/write_batch.h>
+#include "leveldb/write_batch.h"
 
 #import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
 #import "Firestore/Source/Local/FSTLevelDB.h"
 #import "Firestore/Source/Local/FSTLevelDBKey.h"
 #import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
+#import "Firestore/Source/Util/FSTAssert.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
 // Current version of the schema defined in this file.
-static FSTLevelDBSchemaVersion kSchemaVersion = 1;
+static FSTLevelDBSchemaVersion kSchemaVersion = 2;
 
 using leveldb::DB;
+using leveldb::Iterator;
 using leveldb::Status;
 using leveldb::Slice;
 using leveldb::WriteOptions;
@@ -57,6 +58,31 @@ static void SaveVersion(FSTLevelDBSchemaVersion version, FSTWriteGroup *group) {
   [group setData:version_string forKey:key];
 }
 
+/**
+ * This function counts the number of targets that currently exist in the given db. It
+ * then reads the target global row, adds the count to the metadata from that row, and writes
+ * the metadata back.
+ *
+ * It assumes the metadata has already been written and is able to be read in this transaction.
+ */
+static void AddTargetCount(std::shared_ptr<DB> db, FSTWriteGroup *group) {
+  std::unique_ptr<Iterator> it(db->NewIterator([FSTLevelDB standardReadOptions]));
+  Slice start_key = [FSTLevelDBTargetKey keyPrefix];
+  it->Seek(start_key);
+
+  int32_t count = 0;
+  while (it->Valid() && it->key().starts_with(start_key)) {
+    count++;
+    it->Next();
+  }
+
+  FSTPBTargetGlobal *targetGlobal = [FSTLevelDBQueryCache readTargetMetadataFromDB:db];
+  FSTCAssert(targetGlobal != nil,
+             @"We should have a metadata row as it was added in an earlier migration");
+  targetGlobal.targetCount = count;
+  [group setMessage:targetGlobal forKey:[FSTLevelDBTargetGlobalKey key]];
+}
+
 @implementation FSTLevelDBMigrations
 
 + (FSTLevelDBSchemaVersion)schemaVersionForDB:(std::shared_ptr<DB>)db {
@@ -80,6 +106,16 @@ static void SaveVersion(FSTLevelDBSchemaVersion version, FSTWriteGroup *group) {
     case 0:
       EnsureTargetGlobal(db, group);
       // Fallthrough
+    case 1:
+      // We need to make sure we have metadata, since we're going to read and modify it
+      // in this migration. Commit the current transaction and start a new one. Since we're
+      // committing, we need to save a version. It's safe to save this one, if we crash
+      // after saving we'll resume from this step when we try to migrate.
+      SaveVersion(1, group);
+      [group writeToDB:db];
+      group = [FSTWriteGroup groupWithAction:@"Migrations"];
+      AddTargetCount(db, group);
+      // Fallthrough
     default:
       if (currentVersion < kSchemaVersion) {
         SaveVersion(kSchemaVersion, group);
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.mm b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
index b3f4822..fe1bf19 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
@@ -18,7 +18,6 @@
 
 #include <leveldb/db.h>
 #include <leveldb/write_batch.h>
-#include <string>
 
 #import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
 #import "Firestore/Source/Core/FSTQuery.h"
@@ -127,31 +126,50 @@ using leveldb::WriteOptions;
   _db.reset();
 }
 
-- (void)addQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group {
+- (void)saveQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group {
   FSTTargetID targetID = queryData.targetID;
   std::string key = [FSTLevelDBTargetKey keyWithTargetID:targetID];
   [group setMessage:[self.serializer encodedQueryData:queryData] forKey:key];
+}
+
+- (void)saveMetadataInGroup:(FSTWriteGroup *)group {
+  [group setMessage:self.metadata forKey:[FSTLevelDBTargetGlobalKey key]];
+}
+
+- (BOOL)updateMetadataForQueryData:(FSTQueryData *)queryData {
+  BOOL updatedMetadata = NO;
+
+  if (queryData.targetID > self.metadata.highestTargetId) {
+    self.metadata.highestTargetId = queryData.targetID;
+    updatedMetadata = YES;
+  }
+
+  if (queryData.sequenceNumber > self.metadata.highestListenSequenceNumber) {
+    self.metadata.highestListenSequenceNumber = queryData.sequenceNumber;
+    updatedMetadata = YES;
+  }
+  return updatedMetadata;
+}
+
+- (void)addQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group {
+  [self saveQueryData:queryData group:group];
 
   NSString *canonicalID = queryData.query.canonicalID;
   std::string indexKey =
-      [FSTLevelDBQueryTargetKey keyWithCanonicalID:canonicalID targetID:targetID];
+      [FSTLevelDBQueryTargetKey keyWithCanonicalID:canonicalID targetID:queryData.targetID];
   std::string emptyBuffer;
   [group setData:emptyBuffer forKey:indexKey];
 
-  BOOL saveMetadata = NO;
-  FSTPBTargetGlobal *metadata = self.metadata;
-  if (targetID > metadata.highestTargetId) {
-    metadata.highestTargetId = targetID;
-    saveMetadata = YES;
-  }
+  self.metadata.targetCount += 1;
+  [self updateMetadataForQueryData:queryData];
+  [self saveMetadataInGroup:group];
+}
 
-  if (queryData.sequenceNumber > metadata.highestListenSequenceNumber) {
-    metadata.highestListenSequenceNumber = queryData.sequenceNumber;
-    saveMetadata = YES;
-  }
+- (void)updateQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group {
+  [self saveQueryData:queryData group:group];
 
-  if (saveMetadata) {
-    [group setMessage:metadata forKey:[FSTLevelDBTargetGlobalKey key]];
+  if ([self updateMetadataForQueryData:queryData]) {
+    [self saveMetadataInGroup:group];
   }
 }
 
@@ -166,6 +184,12 @@ using leveldb::WriteOptions;
   std::string indexKey =
       [FSTLevelDBQueryTargetKey keyWithCanonicalID:queryData.query.canonicalID targetID:targetID];
   [group removeMessageForKey:indexKey];
+  self.metadata.targetCount -= 1;
+  [self saveMetadataInGroup:group];
+}
+
+- (int32_t)count {
+  return self.metadata.targetCount;
 }
 
 /**
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index d30177a..8a383e5 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -312,7 +312,7 @@ NS_ASSUME_NONNULL_BEGIN
       queryData = [queryData queryDataByReplacingSnapshotVersion:change.snapshotVersion
                                                      resumeToken:resumeToken];
       self.targetIDs[targetIDNumber] = queryData;
-      [self.queryCache addQueryData:queryData group:group];
+      [self.queryCache updateQueryData:queryData group:group];
     }
   }];
 
diff --git a/Firestore/Source/Local/FSTMemoryQueryCache.mm b/Firestore/Source/Local/FSTMemoryQueryCache.mm
index bcab174..56d5699 100644
--- a/Firestore/Source/Local/FSTMemoryQueryCache.mm
+++ b/Firestore/Source/Local/FSTMemoryQueryCache.mm
@@ -90,6 +90,20 @@ NS_ASSUME_NONNULL_BEGIN
   }
 }
 
+- (void)updateQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group {
+  self.queries[queryData.query] = queryData;
+  if (queryData.targetID > self.highestTargetID) {
+    self.highestTargetID = queryData.targetID;
+  }
+  if (queryData.sequenceNumber > self.highestListenSequenceNumber) {
+    self.highestListenSequenceNumber = queryData.sequenceNumber;
+  }
+}
+
+- (int32_t)count {
+  return (int32_t)[self.queries count];
+}
+
 - (void)removeQueryData:(FSTQueryData *)queryData group:(__unused FSTWriteGroup *)group {
   [self.queries removeObjectForKey:queryData.query];
   [self.references removeReferencesForID:queryData.targetID];
diff --git a/Firestore/Source/Local/FSTQueryCache.h b/Firestore/Source/Local/FSTQueryCache.h
index 88c9df9..5c43de4 100644
--- a/Firestore/Source/Local/FSTQueryCache.h
+++ b/Firestore/Source/Local/FSTQueryCache.h
@@ -78,18 +78,29 @@ NS_ASSUME_NONNULL_BEGIN
                                group:(FSTWriteGroup *)group;
 
 /**
- * Adds or replaces an entry in the cache.
+ * Adds an entry in the cache.
  *
- * The cache key is extracted from `queryData.query`. If there is already a cache entry for the
- * key, it will be replaced.
+ * The cache key is extracted from `queryData.query`. The key must not already exist in the cache.
  *
- * @param queryData An FSTQueryData instance to put in the cache.
+ * @param queryData A new FSTQueryData instance to put in the cache.
  */
 - (void)addQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group;
 
+/**
+ * Updates an entry in the cache.
+ *
+ * The cache key is extracted from `queryData.query`. The entry must already exist in the cache,
+ * and it will be replaced.
+ * @param queryData An FSTQueryData instance to replace an existing entry in the cache
+ */
+- (void)updateQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group;
+
 /** Removes the cached entry for the given query data (no-op if no entry exists). */
 - (void)removeQueryData:(FSTQueryData *)queryData group:(FSTWriteGroup *)group;
 
+/** Returns the number of targets cached. */
+- (int32_t)count;
+
 /**
  * Looks up an FSTQueryData entry in the cache.
  *