Fix bug where SaveOverlay() would not correctly update a previous entry for the given document

Author: dconeybe

Firestore/core/src/local/leveldb_document_overlay_cache.cc

@@ -168,12 +168,24 @@ Overlay LevelDbDocumentOverlayCache::ParseOverlay(
 void LevelDbDocumentOverlayCache::SaveOverlay(int largest_batch_id,
                                               const DocumentKey& key,
                                               const Mutation& mutation) {
+  DeleteOverlay(key);
   const std::string leveldb_key =
       LevelDbDocumentOverlayKey::Key(user_id_, key, largest_batch_id);
   auto serialized_mutation = serializer_->EncodeMutation(mutation);
   db_->current_transaction()->Put(leveldb_key, serialized_mutation);
 }
 
+void LevelDbDocumentOverlayCache::DeleteOverlay(const model::DocumentKey& key) {
+  const std::string leveldb_key_prefix =
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, key);
+  auto it = db_->current_transaction()->NewIterator();
+  for (it->Seek(leveldb_key_prefix);
+       it->Valid() && absl::StartsWith(it->key(), leveldb_key_prefix);
+       it->Next()) {
+    db_->current_transaction()->Delete(it->key());
+  }
+}
+
 void LevelDbDocumentOverlayCache::ForEachOverlay(
     std::function<void(absl::string_view encoded_key,
                        const LevelDbDocumentOverlayKey& decoded_key,

Firestore/core/src/local/leveldb_document_overlay_cache.h

@@ -75,6 +75,8 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
                    const model::DocumentKey& key,
                    const model::Mutation& mutation);
 
+  void DeleteOverlay(const model::DocumentKey& key);
+
   void ForEachOverlay(
       std::function<void(absl::string_view encoded_key,
                          const LevelDbDocumentOverlayKey& decoded_key,

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

@@ -266,6 +266,28 @@ TEST_P(DocumentOverlayCacheTest,
   });
 }
 
+TEST_P(DocumentOverlayCacheTest, UpdateDocumentOverlay) {
+  this->persistence_->Run("Test", [&] {
+    Mutation mutation1 = PatchMutation("coll/doc", Map("foo", "bar1"));
+    Mutation mutation2 = PatchMutation("coll/doc", Map("foo", "bar2"));
+    this->SaveOverlaysWithMutations(1, {mutation1});
+    this->SaveOverlaysWithMutations(2, {mutation2});
+
+    // Verify that `GetOverlay()` returns the updated mutation.
+    auto overlay_opt =
+        this->cache_->GetOverlay(DocumentKey::FromPathString("coll/doc"));
+    EXPECT_TRUE(overlay_opt);
+    if (overlay_opt) {
+      EXPECT_EQ(mutation2, overlay_opt.value().mutation());
+    }
+
+    // Verify that `RemoveOverlaysForBatchId()` removes the overlay completely.
+    this->cache_->RemoveOverlaysForBatchId(2);
+    EXPECT_FALSE(
+        this->cache_->GetOverlay(DocumentKey::FromPathString("coll/doc")));
+  });
+}
+
 }  // namespace
 }  // namespace local
 }  // namespace firestore