Firestore: preserve batch IDs after queue restart

Firestore/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [fixed] Preserve pending document overlays for indexed queries after the
+  mutation queue becomes empty and is recreated. (#16751)
 - [feature] Added support for Pipeline DML stages (`delete`, `update`, `insert`, `upsert`), `literals` source stage, and `Pipeline.ExecuteOptions(isAtomic:)`. (#16859)
 
 # 13.0.0

Firestore/core/src/local/leveldb_index_manager.cc

@@ -440,6 +440,26 @@ std::vector<model::FieldIndex> LevelDbIndexManager::GetFieldIndexes() const {
   return result;
 }
 
+model::BatchId LevelDbIndexManager::GetHighestBatchIdAcrossUsers() const {
+  model::BatchId highest_batch_id = model::IndexOffset::InitialLargestBatchId();
+  const auto state_key_prefix = LevelDbIndexStateKey::KeyPrefix();
+  auto state_iter = db_->current_transaction()->NewIterator();
+  LevelDbIndexStateKey state_key;
+  for (state_iter->Seek(state_key_prefix);
+       state_iter->Valid() &&
+       absl::StartsWith(state_iter->key(), state_key_prefix);
+       state_iter->Next()) {
+    if (!state_key.Decode(state_iter->key())) {
+      break;
+    }
+    highest_batch_id =
+        std::max(highest_batch_id, DecodeIndexState(state_iter->value())
+                                       .index_offset()
+                                       .largest_batch_id());
+  }
+  return highest_batch_id;
+}
+
 std::optional<model::FieldIndex> LevelDbIndexManager::GetFieldIndex(
     const core::Target& target) const {
   HARD_ASSERT(started_, "IndexManager not started");

Firestore/core/src/local/leveldb_index_manager.h

@@ -69,6 +69,9 @@ class LevelDbIndexManager : public IndexManager {
 
   std::vector<model::FieldIndex> GetFieldIndexes() const override;
 
+  /** Returns the greatest persisted index batch ID across all users. */
+  model::BatchId GetHighestBatchIdAcrossUsers() const;
+
   void DeleteAllFieldIndexes() override;
 
   void CreateTargetIndexes(const core::Target& target) override;

Firestore/core/src/local/leveldb_mutation_queue.cc

@@ -16,6 +16,8 @@
 
 #include "Firestore/core/src/local/leveldb_mutation_queue.h"
 
+#include <algorithm>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <utility>
@@ -32,6 +34,7 @@
 #include "Firestore/core/src/model/resource_path.h"
 #include "Firestore/core/src/nanopb/nanopb_util.h"
 #include "Firestore/core/src/nanopb/reader.h"
+#include "Firestore/core/src/util/hard_assert.h"
 #include "Firestore/core/src/util/string_util.h"
 #include "Firestore/core/src/util/to_string.h"
 #include "absl/strings/match.h"
@@ -127,7 +130,7 @@ BatchId LoadNextBatchIdFromDb(DB* db) {
 
 LevelDbMutationQueue::LevelDbMutationQueue(const User& user,
                                            LevelDbPersistence* db,
-                                           IndexManager* index_manager,
+                                           LevelDbIndexManager* index_manager,
                                            LocalSerializer* serializer)
     : db_(NOT_NULL(db)),
       index_manager_(NOT_NULL(index_manager)),
@@ -137,6 +140,17 @@ LevelDbMutationQueue::LevelDbMutationQueue(const User& user,
 
 void LevelDbMutationQueue::Start() {
   next_batch_id_ = LoadNextBatchIdFromDb(db_->ptr());
+  // Batch IDs are global. Index offsets are persisted per user and can outlive
+  // all acknowledged batches in the mutation queue, so include their global
+  // high-water mark when restoring the next ID.
+  const BatchId highest_indexed_batch_id =
+      index_manager_->GetHighestBatchIdAcrossUsers();
+  if (highest_indexed_batch_id >= 0) {
+    // Do not fail startup on an unexpectedly large persisted index offset.
+    const BatchId indexed_batch_id = std::min(
+        highest_indexed_batch_id, std::numeric_limits<BatchId>::max() - 1);
+    next_batch_id_ = std::max(next_batch_id_, indexed_batch_id + 1);
+  }
   metadata_ = MetadataForKey(mutation_queue_key());
 }
 

Firestore/core/src/local/leveldb_mutation_queue.h

@@ -45,16 +45,16 @@ class LevelDbPersistence;
 class LocalSerializer;
 
 /**
- * Returns one larger than the largest batch ID that has been stored. If there
- * are no mutations returns 0. Note that batch IDs are global.
+ * Returns one larger than the largest batch ID stored in the mutation table.
+ * Returns 1 if there are no mutations. Batch IDs are global.
  */
 model::BatchId LoadNextBatchIdFromDb(leveldb::DB* db);
 
 class LevelDbMutationQueue : public MutationQueue {
  public:
   LevelDbMutationQueue(const credentials::User& user,
                        LevelDbPersistence* db,
-                       IndexManager* index_manager,
+                       LevelDbIndexManager* index_manager,
                        LocalSerializer* serializer);
 
   void Start() override;
@@ -116,7 +116,7 @@ class LevelDbMutationQueue : public MutationQueue {
 
   // The LevelDbMutationQueue instance is owned by LevelDbPersistence.
   LevelDbPersistence* db_;
-  IndexManager* index_manager_;
+  LevelDbIndexManager* index_manager_;
 
   // Owned by LevelDbPersistence.
   LocalSerializer* serializer_ = nullptr;

Firestore/core/test/unit/local/leveldb_local_store_test.cc

@@ -17,9 +17,13 @@
 #include "Firestore/core/include/firebase/firestore/geo_point.h"
 #include "Firestore/core/src/core/filter.h"
 #include "Firestore/core/src/core/query.h"
+#include "Firestore/core/src/local/leveldb_index_manager.h"
 #include "Firestore/core/src/local/leveldb_persistence.h"
+#include "Firestore/core/src/local/local_write_result.h"
+#include "Firestore/core/src/local/target_data.h"
 #include "Firestore/core/src/model/delete_mutation.h"
 #include "Firestore/core/src/model/field_index.h"
+#include "Firestore/core/src/model/mutation_batch_result.h"
 #include "Firestore/core/src/model/set_mutation.h"
 #include "Firestore/core/src/remote/remote_event.h"
 #include "Firestore/core/test/unit/local/local_store_test.h"
@@ -221,6 +225,56 @@ TEST_F(LevelDbLocalStoreTest, UsesIndexes) {
   FSTAssertQueryReturned("coll/a");
 }
 
+TEST(LevelDbLocalStoreRestartTest, IndexedCacheQueryIncludesNewPendingWrite) {
+  const auto directory = LevelDbDir();
+  const auto user = credentials::User::Unauthenticated();
+  const core::Query query =
+      testutil::Query("coll").AddingFilter(Filter("matches", "==", true));
+
+  {
+    auto persistence = LevelDbPersistenceForTesting(directory);
+    CountingQueryEngine query_engine;
+    LocalStore store(persistence.get(), &query_engine, user);
+    store.Start();
+    store.ConfigureFieldIndexes(
+        {MakeFieldIndex("coll", 0, FieldIndex::InitialState(), "matches",
+                        model::Segment::Kind::kAscending)});
+    store.AllocateTarget(core::QueryOrPipeline(query).ToTargetOrPipeline());
+
+    auto mutation = SetMutation("coll/a", Map("matches", true));
+    const auto write = store.WriteLocally({mutation});
+    store.Backfill();
+    const auto indexed_batch_id = persistence->Run("ReadIndexedBatchId", [&] {
+      return persistence->GetIndexManager(user)
+          ->GetMinOffset("coll")
+          .largest_batch_id();
+    });
+    ASSERT_GT(indexed_batch_id, 0);
+
+    model::MutationBatch batch(write.batch_id(), Timestamp::Now(), {},
+                               {std::move(mutation)});
+    std::vector<model::MutationResult> results;
+    results.emplace_back(Version(10), Array());
+    store.AcknowledgeBatch(
+        model::MutationBatchResult(batch, Version(10), std::move(results), {}));
+    // Include the acknowledged remote document in the index as well. This
+    // advances its read-time offset without discarding the batch high-water.
+    store.Backfill();
+  }
+
+  // Reopen the same database after acknowledging every mutation. The persisted
+  // index offset must not hide the first pending write from the new queue.
+  auto persistence = LevelDbPersistenceForTesting(directory);
+  CountingQueryEngine query_engine;
+  LocalStore store(persistence.get(), &query_engine, user);
+  store.Start();
+  store.WriteLocally({SetMutation("coll/b", Map("matches", true))});
+  const auto result = store.ExecuteQuery(core::QueryOrPipeline(query), false);
+  ASSERT_EQ(result.documents().size(), 2);
+  ASSERT_NE(result.documents().find(Key("coll/b")), result.documents().end());
+  ASSERT_EQ(query_engine.documents_read_by_query(), 0);
+}
+
 TEST_F(LevelDbLocalStoreTest,
        UsesPartiallyIndexedRemoteDocumentsWhenAvailable) {
   FieldIndex index =

Firestore/core/test/unit/local/leveldb_mutation_queue_test.cc

@@ -25,6 +25,7 @@
 #include "Firestore/core/src/local/leveldb_key.h"
 #include "Firestore/core/src/local/leveldb_persistence.h"
 #include "Firestore/core/src/local/reference_set.h"
+#include "Firestore/core/src/model/field_index.h"
 #include "Firestore/core/src/nanopb/byte_string.h"
 #include "Firestore/core/src/nanopb/message.h"
 #include "Firestore/core/src/nanopb/reader.h"
@@ -33,6 +34,7 @@
 #include "Firestore/core/test/unit/local/mutation_queue_test.h"
 #include "Firestore/core/test/unit/local/persistence_testing.h"
 #include "Firestore/core/test/unit/testutil/status_testing.h"
+#include "Firestore/core/test/unit/testutil/testutil.h"
 #include "absl/strings/string_view.h"
 #include "gtest/gtest.h"
 #include "leveldb/db.h"
@@ -52,6 +54,7 @@ using nanopb::ByteString;
 using nanopb::Message;
 using nanopb::StringReader;
 using nanopb::StringWriter;
+using testutil::MakeFieldIndex;
 using util::OrderedCode;
 
 // A dummy mutation value, useful for testing code that's known to examine only
@@ -158,6 +161,31 @@ TEST_F(LevelDbMutationQueueTest, LoadNextBatchIdOnlyFindsMutations) {
   ASSERT_EQ(LoadNextBatchIdFromDb(db_), 4);
 }
 
+TEST_F(LevelDbMutationQueueTest,
+       StartIncludesLargestIndexedBatchIdAcrossUsers) {
+  persistence_->Run("AddIndexAndRestartMutationQueue", [&] {
+    auto* index_manager = persistence_->GetIndexManager(User("user"));
+    index_manager->Start();
+    index_manager->AddFieldIndex(
+        MakeFieldIndex("coll", "a", model::Segment::Kind::kAscending));
+    index_manager->UpdateCollectionGroup(
+        "coll", model::IndexOffset(model::SnapshotVersion::None(),
+                                   model::DocumentKey::Empty(), 20));
+
+    auto* other_user_index_manager =
+        persistence_->GetIndexManager(User("other-user"));
+    other_user_index_manager->Start();
+    other_user_index_manager->UpdateCollectionGroup(
+        "coll", model::IndexOffset(model::SnapshotVersion::None(),
+                                   model::DocumentKey::Empty(), 40));
+
+    // Re-starting models recreating the queue after the database has been
+    // reopened with no outstanding mutation batches.
+    mutation_queue_->Start();
+    ASSERT_EQ(AddMutationBatch().batch_id(), 41);
+  });
+}
+
 TEST_F(LevelDbMutationQueueTest, EmptyProtoCanBeUpgraded) {
   // An empty protocol buffer serializes to a zero-length byte buffer.
   google_protobuf_Empty empty{};