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

5 files changed, 108 insertions, 23 deletions

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.
  *