Schema migration to ensure sentinel rows exist (#1911)

* Add schema migration to ensure each document has a sentinel row

Author: Greg Soltis

Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm

@@ -27,6 +27,8 @@
 
 #include "Firestore/core/src/firebase/firestore/local/leveldb_key.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_migrations.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/types.h"
 #include "Firestore/core/src/firebase/firestore/util/ordered_code.h"
 #include "Firestore/core/src/firebase/firestore/util/status.h"
 #include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
@@ -43,10 +45,14 @@
 using firebase::firestore::local::LevelDbMutationKey;
 using firebase::firestore::local::LevelDbMutationQueueKey;
 using firebase::firestore::local::LevelDbQueryTargetKey;
+using firebase::firestore::local::LevelDbRemoteDocumentKey;
 using firebase::firestore::local::LevelDbTargetDocumentKey;
+using firebase::firestore::local::LevelDbTargetGlobalKey;
 using firebase::firestore::local::LevelDbTargetKey;
 using firebase::firestore::local::LevelDbTransaction;
 using firebase::firestore::model::BatchId;
+using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::ListenSequenceNumber;
 using firebase::firestore::model::TargetId;
 using firebase::firestore::testutil::Key;
 using firebase::firestore::util::OrderedCode;
@@ -199,6 +205,64 @@ - (void)testDropsTheQueryCacheWithThousandsOfEntries {
   }
 }
 
+- (void)testAddsSentinelRows {
+  ListenSequenceNumber old_sequence_number = 1;
+  ListenSequenceNumber new_sequence_number = 2;
+  std::string encoded_old_sequence_number =
+      LevelDbDocumentTargetKey::EncodeSentinelValue(old_sequence_number);
+  LevelDbMigrations::RunMigrations(_db.get(), 3);
+  {
+    std::string empty_buffer;
+    LevelDbTransaction transaction(_db.get(), "Setup");
+
+    // Set up target global
+    FSTPBTargetGlobal *metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db.get()];
+    // Expect that documents missing a row will get the new number
+    metadata.highestListenSequenceNumber = new_sequence_number;
+    transaction.Put(LevelDbTargetGlobalKey::Key(), metadata);
+
+    // Set up some documents (we only need the keys)
+    // For the odd ones, add sentinel rows.
+    for (int i = 0; i < 10; i++) {
+      DocumentKey key = DocumentKey::FromSegments({"docs", std::to_string(i)});
+      transaction.Put(LevelDbRemoteDocumentKey::Key(key), empty_buffer);
+      if (i % 2 == 1) {
+        std::string sentinel_key = LevelDbDocumentTargetKey::SentinelKey(key);
+        transaction.Put(sentinel_key, encoded_old_sequence_number);
+      }
+    }
+
+    transaction.Commit();
+  }
+
+  LevelDbMigrations::RunMigrations(_db.get(), 4);
+  {
+    LevelDbTransaction transaction(_db.get(), "Verify");
+    auto it = transaction.NewIterator();
+    std::string documents_prefix = LevelDbRemoteDocumentKey::KeyPrefix();
+    it->Seek(documents_prefix);
+    int count = 0;
+    LevelDbRemoteDocumentKey document_key;
+    std::string buffer;
+    for (; it->Valid() && absl::StartsWith(it->key(), documents_prefix); it->Next()) {
+      count++;
+      XCTAssertTrue(document_key.Decode(it->key()));
+      const DocumentKey &key = document_key.document_key();
+      std::string sentinel_key = LevelDbDocumentTargetKey::SentinelKey(key);
+      XCTAssertTrue(transaction.Get(sentinel_key, &buffer).ok());
+      int doc_number = atoi(key.path().last_segment().c_str());
+      // If the document number is odd, we expect the original old sequence number that we wrote.
+      // If it's even, we expect that the migration added the new sequence number from the target
+      // global
+      ListenSequenceNumber expected_sequence_number =
+          doc_number % 2 == 1 ? old_sequence_number : new_sequence_number;
+      ListenSequenceNumber sequence_number = LevelDbDocumentTargetKey::DecodeSentinelValue(buffer);
+      XCTAssertEqual(expected_sequence_number, sequence_number);
+    }
+    XCTAssertEqual(10, count);
+  }
+}
+
 /**
  * Creates the name of a dummy entry to make sure the iteration is correctly bounded.
  */

Firestore/Source/Local/FSTLevelDB.mm

@@ -234,9 +234,9 @@ - (FSTLRUGarbageCollector *)gc {
 }
 
 - (void)writeSentinelForKey:(const DocumentKey &)key {
-  std::string encodedSequenceNumber;
-  OrderedCode::WriteSignedNumIncreasing(&encodedSequenceNumber, [self currentSequenceNumber]);
   std::string sentinelKey = LevelDbDocumentTargetKey::SentinelKey(key);
+  std::string encodedSequenceNumber =
+      LevelDbDocumentTargetKey::EncodeSentinelValue([self currentSequenceNumber]);
   _db.currentTransaction->Put(sentinelKey, encodedSequenceNumber);
 }
 

Firestore/Source/Local/FSTLevelDBQueryCache.mm

@@ -30,7 +30,6 @@
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 #include "Firestore/core/src/firebase/firestore/util/hard_assert.h"
-#include "Firestore/core/src/firebase/firestore/util/ordered_code.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 #include "absl/strings/match.h"
 
@@ -49,22 +48,10 @@
 using firebase::firestore::model::SnapshotVersion;
 using firebase::firestore::model::TargetId;
 using firebase::firestore::util::MakeString;
-using firebase::firestore::util::OrderedCode;
 using leveldb::DB;
 using leveldb::Slice;
 using leveldb::Status;
 
-namespace {
-
-ListenSequenceNumber ReadSequenceNumber(absl::string_view slice) {
-  ListenSequenceNumber decoded;
-  if (!OrderedCode::ReadSignedNumIncreasing(&slice, &decoded)) {
-    HARD_FAIL("Failed to read sequence number from a sentinel row");
-  }
-  return decoded;
-}
-}  // namespace
-
 @interface FSTLevelDBQueryCache ()
 
 /** A write-through cached copy of the metadata for the query cache. */
@@ -201,7 +188,7 @@ - (void)enumerateOrphanedDocumentsUsingBlock:
       }
       // set nextToReport to be this sequence number. It's the next one we might
       // report, if we don't find any targets for this document.
-      nextToReport = ReadSequenceNumber(it->value());
+      nextToReport = LevelDbDocumentTargetKey::DecodeSentinelValue(it->value());
       keyToReport = key.document_key();
     } else {
       // set nextToReport to be 0, we know we don't need to report this one since

Firestore/core/src/firebase/firestore/local/leveldb_key.cc

@@ -793,6 +793,22 @@ std::string LevelDbDocumentTargetKey::SentinelKey(
   return Key(document_key, kInvalidTargetId);
 }
 
+std::string LevelDbDocumentTargetKey::EncodeSentinelValue(
+    model::ListenSequenceNumber sequence_number) {
+  std::string encoded;
+  OrderedCode::WriteSignedNumIncreasing(&encoded, sequence_number);
+  return encoded;
+}
+
+model::ListenSequenceNumber LevelDbDocumentTargetKey::DecodeSentinelValue(
+    absl::string_view slice) {
+  model::ListenSequenceNumber decoded;
+  if (!OrderedCode::ReadSignedNumIncreasing(&slice, &decoded)) {
+    HARD_FAIL("Failed to read sequence number from a sentinel row");
+  }
+  return decoded;
+}
+
 bool LevelDbDocumentTargetKey::Decode(absl::string_view key) {
   Reader reader{key};
   reader.ReadTableNameMatching(kDocumentTargetsTable);

Firestore/core/src/firebase/firestore/local/leveldb_key.h

@@ -416,6 +416,18 @@ class LevelDbDocumentTargetKey {
    */
   static std::string SentinelKey(const model::DocumentKey& document_key);
 
+  /**
+   * Given a sequence number, encodes it for storage in a sentinel row.
+   */
+  static std::string EncodeSentinelValue(
+      model::ListenSequenceNumber sequence_number);
+
+  /**
+   * Given an encoded sentinel row, return the sequence number.
+   */
+  static model::ListenSequenceNumber DecodeSentinelValue(
+      absl::string_view slice);
+
   /**
    * Decodes the contents of a document target key, storing the decoded values
    * in this instance.

Firestore/core/src/firebase/firestore/local/leveldb_migrations.cc

@@ -21,6 +21,8 @@
 
 #include "Firestore/Protos/nanopb/firestore/local/target.nanopb.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/types.h"
 #include "Firestore/core/src/firebase/firestore/nanopb/writer.h"
 #include "absl/strings/match.h"
 
@@ -32,6 +34,7 @@ using leveldb::Iterator;
 using leveldb::Slice;
 using leveldb::Status;
 using leveldb::WriteOptions;
+using nanopb::Reader;
 using nanopb::Writer;
 
 namespace {
@@ -52,8 +55,10 @@ namespace {
  *   * Migration 3 deletes the entire query cache to deal with cache corruption
  *     related to limbo resolution. Addresses
  *     https://github.com/firebase/firebase-ios-sdk/issues/1548.
+ *   * Migration 4 ensures that every document in the remote document cache
+ *     has a sentinel row with a sequence number.
  */
-const LevelDbMigrations::SchemaVersion kSchemaVersion = 3;
+const LevelDbMigrations::SchemaVersion kSchemaVersion = 4;
 
 /**
  * Save the given version number as the current version of the schema of the
@@ -110,6 +115,63 @@ void ClearQueryCache(leveldb::DB* db) {
   transaction.Commit();
 }
 
+/**
+ * Reads the highest sequence number from the target global row.
+ */
+model::ListenSequenceNumber GetHighestSequenceNumber(
+    LevelDbTransaction* transaction) {
+  std::string bytes;
+  transaction->Get(LevelDbTargetGlobalKey::Key(), &bytes);
+
+  firestore_client_TargetGlobal target_global{};
+  Reader reader = Reader::Wrap(bytes);
+  reader.ReadNanopbMessage(firestore_client_TargetGlobal_fields,
+                           &target_global);
+  return target_global.highest_listen_sequence_number;
+}
+
+/**
+ * Given a document key, ensure it has a sentinel row. If it doesn't have one,
+ * add it with the given value.
+ */
+void EnsureSentinelRow(LevelDbTransaction* transaction,
+                       const model::DocumentKey& key,
+                       const std::string& sentinel_value) {
+  std::string sentinel_key = LevelDbDocumentTargetKey::SentinelKey(key);
+  std::string unused_value;
+  if (transaction->Get(sentinel_key, &unused_value).IsNotFound()) {
+    transaction->Put(sentinel_key, sentinel_value);
+  }
+}
+
+/**
+ * Ensure each document in the remote document table has a corresponding
+ * sentinel row in the document target index.
+ */
+void EnsureSentinelRows(leveldb::DB* db) {
+  LevelDbTransaction transaction(db, "Ensure sentinel rows");
+
+  // Get the value we'll use for anything that's missing a row.
+  model::ListenSequenceNumber sequence_number =
+      GetHighestSequenceNumber(&transaction);
+  std::string sentinel_value =
+      LevelDbDocumentTargetKey::EncodeSentinelValue(sequence_number);
+
+  std::string documents_prefix = LevelDbRemoteDocumentKey::KeyPrefix();
+  auto it = transaction.NewIterator();
+  it->Seek(documents_prefix);
+  LevelDbRemoteDocumentKey document_key;
+  for (; it->Valid() && absl::StartsWith(it->key(), documents_prefix);
+       it->Next()) {
+    HARD_ASSERT(document_key.Decode(it->key()),
+                "Failed to decode document key");
+    EnsureSentinelRow(&transaction, document_key.document_key(),
+                      sentinel_value);
+  }
+  SaveVersion(4, &transaction);
+  transaction.Commit();
+}
+
 }  // namespace
 
 LevelDbMigrations::SchemaVersion LevelDbMigrations::ReadSchemaVersion(
@@ -139,6 +201,10 @@ void LevelDbMigrations::RunMigrations(leveldb::DB* db,
   if (from_version < 3 && to_version >= 3) {
     ClearQueryCache(db);
   }
+
+  if (from_version < 4 && to_version >= 4) {
+    EnsureSentinelRows(db);
+  }
 }
 
 }  // namespace local

Firestore/core/src/firebase/firestore/nanopb/reader.cc

@@ -30,6 +30,12 @@ Reader Reader::Wrap(const uint8_t* bytes, size_t length) {
   return Reader{pb_istream_from_buffer(bytes, length)};
 }
 
+Reader Reader::Wrap(absl::string_view string_view) {
+  return Reader{pb_istream_from_buffer(
+      reinterpret_cast<const uint8_t*>(string_view.data()),
+      string_view.size())};
+}
+
 uint32_t Reader::ReadTag() {
   Tag tag;
   if (!status_.ok()) return 0;

Firestore/core/src/firebase/firestore/nanopb/reader.h

@@ -59,6 +59,15 @@ class Reader {
    */
   static Reader Wrap(const uint8_t* bytes, size_t length);
 
+  /**
+   * Creates an input stream from bytes backing the string_view. Note that
+   * the backing buffer must remain valid for the lifetime of this Reader.
+   *
+   * (This is roughly equivalent to the nanopb function
+   * pb_istream_from_buffer())
+   */
+  static Reader Wrap(absl::string_view);
+
   /**
    * Reads a message type from the input stream.
    *