Index Manager handles index states (#9384)

Author: wu-hui

Firestore/core/src/local/leveldb_index_manager.cc

@@ -21,6 +21,7 @@
 #include <utility>
 #include <vector>
 
+#include "Firestore/core/src/credentials/user.h"
 #include "Firestore/core/src/local/leveldb_key.h"
 #include "Firestore/core/src/local/leveldb_persistence.h"
 #include "Firestore/core/src/local/local_serializer.h"
@@ -35,6 +36,7 @@ namespace firebase {
 namespace firestore {
 namespace local {
 
+using credentials::User;
 using model::DocumentKey;
 using model::FieldIndex;
 using model::IndexState;
@@ -49,21 +51,15 @@ struct DbIndexState {
   int32_t nanos;
   std::string key;
   model::ListenSequenceNumber sequence_number;
+  model::BatchId largest_batch_id;
 };
 
-// TODO(wuandy): Uncomment this when needed.
-// void to_json(json& j, const DbIndexState& s) {
-//  j = json{{"seconds", s.seconds},
-//           {"nanos", s.nanos},
-//           {"key", s.key},
-//           {"seq_num", s.sequence_number}};
-//}
-
 void from_json(const json& j, DbIndexState& s) {
   j.at("seconds").get_to(s.seconds);
   j.at("nanos").get_to(s.nanos);
   j.at("key").get_to(s.key);
   j.at("seq_num").get_to(s.sequence_number);
+  j.at("largest_batch").get_to(s.largest_batch_id);
 }
 
 IndexState DecodeIndexState(const std::string& encoded) {
@@ -72,16 +68,34 @@ IndexState DecodeIndexState(const std::string& encoded) {
   auto db_state = j.get<DbIndexState>();
   return {db_state.sequence_number,
           SnapshotVersion(Timestamp(db_state.seconds, db_state.nanos)),
-          DocumentKey::FromPathString(db_state.key)};
+          DocumentKey::FromPathString(db_state.key), db_state.largest_batch_id};
+}
+
+std::string EncodeIndexState(const IndexState& state) {
+  return json{
+      {"seconds", state.index_offset().read_time().timestamp().seconds()},
+      {"nanos", state.index_offset().read_time().timestamp().nanoseconds()},
+      {"key", state.index_offset().document_key().ToString()},
+      {"seq_num", state.sequence_number()},
+      {"largest_batch", state.index_offset().largest_batch_id()}}
+      .dump();
 }
 
 }  // namespace
 
-LevelDbIndexManager::LevelDbIndexManager(LevelDbPersistence* db,
+LevelDbIndexManager::LevelDbIndexManager(const User& user,
+                                         LevelDbPersistence* db,
                                          LocalSerializer* serializer)
-    : db_(db), serializer_(serializer) {
+    : db_(db), serializer_(serializer), uid_(user.uid()) {
+  // The contract for this comparison expected by priority queue is
+  // `std::less`, but std::priority_queue's default order is descending.
+  // We change the order to be ascending by doing left >= right instead.
   auto cmp = [](FieldIndex* left, FieldIndex* right) {
-    return left->index_state().sequence_number() <
+    if (left->index_state().sequence_number() ==
+        right->index_state().sequence_number()) {
+      return left->collection_group() >= right->collection_group();
+    }
+    return left->index_state().sequence_number() >
            right->index_state().sequence_number();
   };
   next_index_to_update_ = std::priority_queue<
@@ -225,7 +239,8 @@ void LevelDbIndexManager::MemoizeIndex(FieldIndex index) {
   }
 
   // Moves `index` into `existing_indexes`.
-  existing_indexes.insert({index_id, std::move(index)});
+  existing_indexes[index_id] = std::move(index);
+
   // next_index_to_update_ holds a pointer to the index owned by
   // `existing_indexes`.
   next_index_to_update_.push(&existing_indexes.find(index_id)->second);
@@ -300,9 +315,11 @@ std::vector<FieldIndex> LevelDbIndexManager::GetFieldIndexes(
   HARD_ASSERT(started_, "IndexManager not started");
 
   std::vector<FieldIndex> result;
-  const auto& indexes = memoized_indexes_[collection_group];
-  for (const auto& entry : indexes) {
-    result.push_back(entry.second);
+  const auto iter = memoized_indexes_.find(collection_group);
+  if (iter != memoized_indexes_.end()) {
+    for (const auto& entry : iter->second) {
+      result.push_back(entry.second);
+    }
   }
 
   return result;
@@ -335,13 +352,30 @@ LevelDbIndexManager::GetDocumentsMatchingTarget(model::FieldIndex field_index,
 
 absl::optional<std::string>
 LevelDbIndexManager::GetNextCollectionGroupToUpdate() {
-  return {};
+  if (next_index_to_update_.empty()) {
+    return absl::nullopt;
+  }
+
+  return next_index_to_update_.top()->collection_group();
 }
 
 void LevelDbIndexManager::UpdateCollectionGroup(
     const std::string& collection_group, model::IndexOffset offset) {
-  (void)collection_group;
-  (void)offset;
+  HARD_ASSERT(started_, "IndexManager not started");
+
+  ++memoized_max_sequence_number_;
+  for (const auto& field_index : GetFieldIndexes(collection_group)) {
+    IndexState updated_state{memoized_max_sequence_number_, offset};
+
+    auto state_key = LevelDbIndexStateKey::Key(uid_, field_index.index_id());
+    auto val = EncodeIndexState(updated_state);
+    db_->current_transaction()->Put(std::move(state_key),
+                                    EncodeIndexState(updated_state));
+
+    MemoizeIndex(FieldIndex{field_index.index_id(),
+                            field_index.collection_group(),
+                            field_index.segments(), std::move(updated_state)});
+  }
 }
 
 void LevelDbIndexManager::UpdateIndexEntries(

Firestore/core/src/local/leveldb_index_manager.h

@@ -28,6 +28,11 @@
 
 namespace firebase {
 namespace firestore {
+
+namespace credentials {
+class User;
+}  // namespace credentials
+
 namespace local {
 
 class LevelDbPersistence;
@@ -36,7 +41,8 @@ class LocalSerializer;
 /** A persisted implementation of IndexManager. */
 class LevelDbIndexManager : public IndexManager {
  public:
-  explicit LevelDbIndexManager(LevelDbPersistence* db,
+  explicit LevelDbIndexManager(const credentials::User& user,
+                               LevelDbPersistence* db,
                                LocalSerializer* serializer);
 
   void Start() override;

Firestore/core/src/local/leveldb_mutation_queue.cc

@@ -126,8 +126,10 @@ BatchId LoadNextBatchIdFromDb(DB* db) {
 
 LevelDbMutationQueue::LevelDbMutationQueue(const User& user,
                                            LevelDbPersistence* db,
+                                           IndexManager* index_manager,
                                            LocalSerializer* serializer)
     : db_(NOT_NULL(db)),
+      index_manager_(NOT_NULL(index_manager)),
       serializer_(NOT_NULL(serializer)),
       user_id_(user.is_authenticated() ? user.uid() : "") {
 }
@@ -177,8 +179,7 @@ MutationBatch LevelDbMutationQueue::AddMutationBatch(
     key = LevelDbDocumentMutationKey::Key(user_id_, mutation.key(), batch_id);
     db_->current_transaction()->Put(key, empty_buffer);
 
-    db_->index_manager()->AddToCollectionParentIndex(
-        mutation.key().path().PopLast());
+    index_manager_->AddToCollectionParentIndex(mutation.key().path().PopLast());
   }
 
   return batch;

Firestore/core/src/local/leveldb_mutation_queue.h

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "Firestore/Protos/nanopb/firestore/local/mutation.nanopb.h"
+#include "Firestore/core/src/local/leveldb_index_manager.h"
 #include "Firestore/core/src/local/mutation_queue.h"
 #include "Firestore/core/src/model/model_fwd.h"
 #include "Firestore/core/src/model/types.h"
@@ -53,6 +54,7 @@ class LevelDbMutationQueue : public MutationQueue {
  public:
   LevelDbMutationQueue(const credentials::User& user,
                        LevelDbPersistence* db,
+                       IndexManager* index_manager,
                        LocalSerializer* serializer);
 
   void Start() override;
@@ -114,6 +116,7 @@ class LevelDbMutationQueue : public MutationQueue {
 
   // The LevelDbMutationQueue instance is owned by LevelDbPersistence.
   LevelDbPersistence* db_;
+  IndexManager* index_manager_;
 
   // Owned by LevelDbPersistence.
   LocalSerializer* serializer_ = nullptr;

Firestore/core/src/local/leveldb_persistence.cc

@@ -114,7 +114,6 @@ LevelDbPersistence::LevelDbPersistence(std::unique_ptr<leveldb::DB> db,
   target_cache_ = absl::make_unique<LevelDbTargetCache>(this, &serializer_);
   document_cache_ =
       absl::make_unique<LevelDbRemoteDocumentCache>(this, &serializer_);
-  index_manager_ = absl::make_unique<LevelDbIndexManager>(this, &serializer_);
   reference_delegate_ =
       absl::make_unique<LevelDbLruReferenceDelegate>(this, lru_params);
   bundle_cache_ = absl::make_unique<LevelDbBundleCache>(this, &serializer_);
@@ -225,11 +224,11 @@ void LevelDbPersistence::Shutdown() {
   db_.reset();
 }
 
-LevelDbMutationQueue* LevelDbPersistence::GetMutationQueueForUser(
-    const credentials::User& user) {
+LevelDbMutationQueue* LevelDbPersistence::GetMutationQueue(
+    const credentials::User& user, IndexManager* manager) {
   users_.insert(user.uid());
-  current_mutation_queue_ =
-      absl::make_unique<LevelDbMutationQueue>(user, this, &serializer_);
+  current_mutation_queue_ = absl::make_unique<LevelDbMutationQueue>(
+      user, this, dynamic_cast<LevelDbIndexManager*>(manager), &serializer_);
   return current_mutation_queue_.get();
 }
 
@@ -241,7 +240,11 @@ LevelDbRemoteDocumentCache* LevelDbPersistence::remote_document_cache() {
   return document_cache_.get();
 }
 
-LevelDbIndexManager* LevelDbPersistence::index_manager() {
+LevelDbIndexManager* LevelDbPersistence::GetIndexManager(
+    const credentials::User& user) {
+  users_.insert(user.uid());
+  index_manager_ =
+      absl::make_unique<LevelDbIndexManager>(user, this, &serializer_);
   return index_manager_.get();
 }
 
@@ -253,7 +256,7 @@ LevelDbBundleCache* LevelDbPersistence::bundle_cache() {
   return bundle_cache_.get();
 }
 
-LevelDbDocumentOverlayCache* LevelDbPersistence::document_overlay_cache(
+LevelDbDocumentOverlayCache* LevelDbPersistence::GetDocumentOverlayCache(
     const User& user) {
   users_.insert(user.uid());
   current_document_overlay_cache_ =

Firestore/core/src/local/leveldb_persistence.h

@@ -81,17 +81,17 @@ class LevelDbPersistence : public Persistence {
 
   LevelDbBundleCache* bundle_cache() override;
 
-  LevelDbDocumentOverlayCache* document_overlay_cache(
+  LevelDbDocumentOverlayCache* GetDocumentOverlayCache(
       const credentials::User& user) override;
 
-  LevelDbMutationQueue* GetMutationQueueForUser(
-      const credentials::User& user) override;
+  LevelDbMutationQueue* GetMutationQueue(const credentials::User& user,
+                                         IndexManager* index_manager) override;
 
   LevelDbTargetCache* target_cache() override;
 
   LevelDbRemoteDocumentCache* remote_document_cache() override;
 
-  LevelDbIndexManager* index_manager() override;
+  LevelDbIndexManager* GetIndexManager(const credentials::User& user) override;
 
   LevelDbLruReferenceDelegate* reference_delegate() override;
 

Firestore/core/src/local/leveldb_remote_document_cache.cc

@@ -115,8 +115,8 @@ void LevelDbRemoteDocumentCache::Add(const MutableDocument& document,
       path.PopLast(), read_time, path.last_segment());
   db_->current_transaction()->Put(ldb_read_time_key, "");
 
-  db_->index_manager()->AddToCollectionParentIndex(
-      document.key().path().PopLast());
+  NOT_NULL(index_manager_);
+  index_manager_->AddToCollectionParentIndex(document.key().path().PopLast());
 }
 
 void LevelDbRemoteDocumentCache::Remove(const DocumentKey& key) {
@@ -285,6 +285,10 @@ MutableDocument LevelDbRemoteDocumentCache::DecodeMaybeDocument(
   return maybe_document;
 }
 
+void LevelDbRemoteDocumentCache::SetIndexManager(IndexManager* manager) {
+  index_manager_ = NOT_NULL(manager);
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase

Firestore/core/src/local/leveldb_remote_document_cache.h

@@ -21,6 +21,7 @@
 #include <thread>  // NOLINT(build/c++11)
 #include <vector>
 
+#include "Firestore/core/src/local/leveldb_index_manager.h"
 #include "Firestore/core/src/local/remote_document_cache.h"
 #include "Firestore/core/src/model/model_fwd.h"
 #include "Firestore/core/src/model/types.h"
@@ -55,6 +56,8 @@ class LevelDbRemoteDocumentCache : public RemoteDocumentCache {
       const core::Query& query,
       const model::SnapshotVersion& since_read_time) override;
 
+  void SetIndexManager(IndexManager* manager) override;
+
  private:
   /**
    * Looks up a set of entries in the cache, returning only existing entries of
@@ -67,6 +70,8 @@ class LevelDbRemoteDocumentCache : public RemoteDocumentCache {
 
   // The LevelDbRemoteDocumentCache instance is owned by LevelDbPersistence.
   LevelDbPersistence* db_;
+  // The LevelDbIndexManager instance is owned by LevelDbPersistence.
+  IndexManager* index_manager_ = nullptr;
   // Owned by LevelDbPersistence.
   LocalSerializer* serializer_ = nullptr;
 

Firestore/core/src/local/local_store.cc

@@ -82,15 +82,16 @@ LocalStore::LocalStore(Persistence* persistence,
                        QueryEngine* query_engine,
                        const User& initial_user)
     : persistence_(persistence),
-      mutation_queue_(persistence->GetMutationQueueForUser(initial_user)),
       remote_document_cache_(persistence->remote_document_cache()),
       target_cache_(persistence->target_cache()),
       bundle_cache_(persistence->bundle_cache()),
-      query_engine_(query_engine),
-      local_documents_(
-          absl::make_unique<LocalDocumentsView>(remote_document_cache_,
-                                                mutation_queue_,
-                                                persistence->index_manager())) {
+      query_engine_(query_engine) {
+  index_manager_ = persistence->GetIndexManager(initial_user);
+  mutation_queue_ = persistence->GetMutationQueue(initial_user, index_manager_);
+  local_documents_ = absl::make_unique<LocalDocumentsView>(
+      remote_document_cache_, mutation_queue_, index_manager_);
+  remote_document_cache_->SetIndexManager(index_manager_);
+
   persistence->reference_delegate()->AddInMemoryPins(&local_view_references_);
   target_id_generator_ = TargetIdGenerator::TargetCacheTargetIdGenerator(0);
   query_engine_->SetLocalDocumentsView(local_documents_.get());
@@ -117,7 +118,9 @@ DocumentMap LocalStore::HandleUserChange(const User& user) {
 
   // The old one has a reference to the mutation queue, so null it out first.
   local_documents_.reset();
-  mutation_queue_ = persistence_->GetMutationQueueForUser(user);
+  index_manager_ = persistence_->GetIndexManager(user);
+  mutation_queue_ = persistence_->GetMutationQueue(user, index_manager_);
+  remote_document_cache_->SetIndexManager(index_manager_);
 
   StartMutationQueue();
 
@@ -127,7 +130,7 @@ DocumentMap LocalStore::HandleUserChange(const User& user) {
 
     // Recreate our LocalDocumentsView using the new MutationQueue.
     local_documents_ = absl::make_unique<LocalDocumentsView>(
-        remote_document_cache_, mutation_queue_, persistence_->index_manager());
+        remote_document_cache_, mutation_queue_, index_manager_);
     query_engine_->SetLocalDocumentsView(local_documents_.get());
 
     // Union the old/new changed keys.

Firestore/core/src/local/local_store.h

@@ -51,6 +51,7 @@ class TargetChange;
 namespace local {
 
 class BundleCache;
+class IndexManager;
 class LocalDocumentsView;
 class LocalViewChanges;
 class LocalWriteResult;
@@ -353,6 +354,11 @@ class LocalStore : public bundle::BundleCallback {
    */
   QueryEngine* query_engine_ = nullptr;
 
+  /**
+   * Manages indexes and support indexed queries.
+   */
+  IndexManager* index_manager_ = nullptr;
+
   /**
    * The "local" view of all documents (layering mutation queue on top of
    * remote_document_cache_).