GC targets without deserializing Target model (#7026)

Author: schmidt-sebastian

Firestore/CHANGELOG.md

@@ -1,5 +1,10 @@
+# Unreleased
+- [fixed] Fixed a crash that could happen when the SDK encountered invalid
+  data during garbage collection (#6721).
+
 # v7.2.0
-- [added] Made emulator connection API consistent between Auth, Database, Firestore, and Functions (#5916).
+- [added] Made emulator connection API consistent between Auth, Database,
+  Firestore, and Functions (#5916).
 
 # v7.1.0
 - [changed] Added the original query data to error messages for Queries that

Firestore/core/src/local/leveldb_lru_reference_delegate.cc

@@ -78,7 +78,7 @@ void LevelDbLruReferenceDelegate::RemoveMutationReference(
 void LevelDbLruReferenceDelegate::RemoveTarget(const TargetData& target_data) {
   TargetData updated =
       target_data.WithSequenceNumber(current_sequence_number());
-  db_->target_cache()->UpdateTarget(std::move(updated));
+  db_->target_cache()->UpdateTarget(updated);
 }
 
 void LevelDbLruReferenceDelegate::UpdateLimboDocument(const DocumentKey& key) {
@@ -119,9 +119,9 @@ size_t LevelDbLruReferenceDelegate::GetSequenceNumberCount() {
   return total_count;
 }
 
-void LevelDbLruReferenceDelegate::EnumerateTargets(
-    const TargetCallback& callback) {
-  db_->target_cache()->EnumerateTargets(callback);
+void LevelDbLruReferenceDelegate::EnumerateTargetSequenceNumbers(
+    const SequenceNumberCallback& callback) {
+  db_->target_cache()->EnumerateSequenceNumbers(callback);
 }
 
 void LevelDbLruReferenceDelegate::EnumerateOrphanedDocuments(

Firestore/core/src/local/leveldb_lru_reference_delegate.h

@@ -60,7 +60,8 @@ class LevelDbLruReferenceDelegate : public LruDelegate {
   util::StatusOr<int64_t> CalculateByteSize() override;
   size_t GetSequenceNumberCount() override;
 
-  void EnumerateTargets(const TargetCallback& callback) override;
+  void EnumerateTargetSequenceNumbers(
+      const SequenceNumberCallback& callback) override;
   void EnumerateOrphanedDocuments(
       const OrphanedDocumentCallback& callback) override;
 

Firestore/core/src/local/leveldb_target_cache.cc

@@ -17,6 +17,7 @@
 #include "Firestore/core/src/local/leveldb_target_cache.h"
 
 #include <string>
+#include <unordered_set>
 #include <utility>
 
 #include "Firestore/core/src/local/leveldb_key.h"
@@ -124,7 +125,7 @@ void LevelDbTargetCache::UpdateTarget(const TargetData& target_data) {
 void LevelDbTargetCache::RemoveTarget(const TargetData& target_data) {
   TargetId target_id = target_data.target_id();
 
-  RemoveAllKeysForTarget(target_id);
+  RemoveAllDocumentKeysForTarget(target_id);
 
   std::string key = LevelDbTargetKey::Key(target_id);
   db_->current_transaction()->Delete(key);
@@ -188,35 +189,57 @@ absl::optional<TargetData> LevelDbTargetCache::GetTarget(const Target& target) {
   return absl::nullopt;
 }
 
-void LevelDbTargetCache::EnumerateTargets(const TargetCallback& callback) {
+void LevelDbTargetCache::EnumerateSequenceNumbers(
+    const SequenceNumberCallback& callback) {
   // Enumerate all targets, give their sequence numbers.
   std::string target_prefix = LevelDbTargetKey::KeyPrefix();
   auto it = db_->current_transaction()->NewIterator();
   it->Seek(target_prefix);
   for (; it->Valid() && absl::StartsWith(it->key(), target_prefix);
        it->Next()) {
-    TargetData target = DecodeTarget(it->value());
-    callback(target);
+    StringReader reader{it->value()};
+    auto target_proto = DecodeTargetProto(&reader);
+    callback(target_proto->last_listen_sequence_number);
   }
 }
 
-int LevelDbTargetCache::RemoveTargets(
+size_t LevelDbTargetCache::RemoveTargets(
     ListenSequenceNumber upper_bound,
     const std::unordered_map<model::TargetId, TargetData>& live_targets) {
-  int count = 0;
   std::string target_prefix = LevelDbTargetKey::KeyPrefix();
   auto it = db_->current_transaction()->NewIterator();
   it->Seek(target_prefix);
+
+  std::unordered_set<TargetId> removed_targets;
+
+  // In https://github.com/firebase/firebase-ios-sdk/issues/6721, a customer
+  // reports that their client crashes when deserializing an invalid Target
+  // during an LRU run. Instead of deserializing the value into a full Target
+  // model, we only convert it into the underlying Protobuf message.
   for (; it->Valid() && absl::StartsWith(it->key(), target_prefix);
        it->Next()) {
-    TargetData target_data = DecodeTarget(it->value());
-    if (target_data.sequence_number() <= upper_bound &&
-        live_targets.find(target_data.target_id()) == live_targets.end()) {
-      RemoveTarget(target_data);
-      count++;
+    StringReader reader{it->value()};
+    auto target_proto = DecodeTargetProto(&reader);
+    if (target_proto->last_listen_sequence_number <= upper_bound &&
+        live_targets.find(target_proto->target_id) == live_targets.end()) {
+      TargetId target_id = target_proto->target_id;
+
+      // Remove the DocumentKey to TargetId mapping
+      RemoveAllDocumentKeysForTarget(target_id);
+      // Remove the TargetId to Target mapping
+      db_->current_transaction()->Delete(it->key());
+
+      removed_targets.insert(target_id);
     }
   }
-  return count;
+
+  // Remove the CanonicalId to TargetId mapping
+  RemoveQueryTargetKeyForTargets(removed_targets);
+
+  metadata_->target_count -= removed_targets.size();
+  SaveMetadata();
+
+  return removed_targets.size();
 }
 
 void LevelDbTargetCache::AddMatchingKeys(const DocumentKeySet& keys,
@@ -247,7 +270,7 @@ void LevelDbTargetCache::RemoveMatchingKeys(const DocumentKeySet& keys,
   }
 }
 
-void LevelDbTargetCache::RemoveAllKeysForTarget(TargetId target_id) {
+void LevelDbTargetCache::RemoveAllDocumentKeysForTarget(TargetId target_id) {
   std::string index_prefix = LevelDbTargetDocumentKey::KeyPrefix(target_id);
   auto index_iterator = db_->current_transaction()->NewIterator();
   index_iterator->Seek(index_prefix);
@@ -269,6 +292,24 @@ void LevelDbTargetCache::RemoveAllKeysForTarget(TargetId target_id) {
   }
 }
 
+void LevelDbTargetCache::RemoveQueryTargetKeyForTargets(
+    const std::unordered_set<TargetId>& target_ids) {
+  std::string index_prefix = LevelDbQueryTargetKey::KeyPrefix();
+  auto index_iterator = db_->current_transaction()->NewIterator();
+  index_iterator->Seek(index_prefix);
+
+  LevelDbQueryTargetKey row_key;
+  for (; index_iterator->Valid(); index_iterator->Next()) {
+    if (!row_key.Decode(index_iterator->key())) {
+      break;
+    }
+
+    if (target_ids.find(row_key.target_id()) != target_ids.end()) {
+      db_->current_transaction()->Delete(index_iterator->key());
+    }
+  }
+}
+
 DocumentKeySet LevelDbTargetCache::GetMatchingKeys(TargetId target_id) {
   std::string index_prefix = LevelDbTargetDocumentKey::KeyPrefix(target_id);
   auto index_iterator = db_->current_transaction()->NewIterator();
@@ -386,12 +427,21 @@ void LevelDbTargetCache::SaveMetadata() {
   db_->current_transaction()->Put(LevelDbTargetGlobalKey::Key(), metadata_);
 }
 
+nanopb::Message<firestore_client_Target> LevelDbTargetCache::DecodeTargetProto(
+    nanopb::Reader* reader) {
+  auto message = Message<firestore_client_Target>::TryParse(reader);
+  if (!reader->ok()) {
+    HARD_FAIL("Target proto failed to parse: %s", reader->status().ToString());
+  }
+  return message;
+}
+
 TargetData LevelDbTargetCache::DecodeTarget(absl::string_view encoded) {
   StringReader reader{encoded};
-  auto message = Message<firestore_client_Target>::TryParse(&reader);
+  auto message = DecodeTargetProto(&reader);
   auto result = serializer_->DecodeTargetData(&reader, *message);
   if (!reader.ok()) {
-    HARD_FAIL("Target proto failed to parse: %s, message: %s",
+    HARD_FAIL("Target failed to parse: %s, message: %s",
               reader.status().ToString(), message.ToString());
   }
 

Firestore/core/src/local/leveldb_target_cache.h

@@ -18,6 +18,7 @@
 #define FIRESTORE_CORE_SRC_LOCAL_LEVELDB_TARGET_CACHE_H_
 
 #include <unordered_map>
+#include <unordered_set>
 
 #include "Firestore/Protos/nanopb/firestore/local/target.nanopb.h"
 #include "Firestore/core/src/local/target_cache.h"
@@ -71,11 +72,12 @@ class LevelDbTargetCache : public TargetCache {
 
   absl::optional<TargetData> GetTarget(const core::Target& target) override;
 
-  void EnumerateTargets(const TargetCallback& callback) override;
+  void EnumerateSequenceNumbers(
+      const SequenceNumberCallback& callback) override;
 
-  int RemoveTargets(model::ListenSequenceNumber upper_bound,
-                    const std::unordered_map<model::TargetId, TargetData>&
-                        live_targets) override;
+  size_t RemoveTargets(model::ListenSequenceNumber upper_bound,
+                       const std::unordered_map<model::TargetId, TargetData>&
+                           live_targets) override;
 
   // Key-related methods
 
@@ -91,8 +93,8 @@ class LevelDbTargetCache : public TargetCache {
   void RemoveMatchingKeys(const model::DocumentKeySet& keys,
                           model::TargetId target_id) override;
 
-  /** Removes all the keys in the query results of the given target ID. */
-  void RemoveAllKeysForTarget(model::TargetId target_id);
+  /** Removes all document keys in the query results of the given target ID. */
+  void RemoveAllDocumentKeysForTarget(model::TargetId target_id);
 
   model::DocumentKeySet GetMatchingKeys(model::TargetId target_id) override;
 
@@ -128,12 +130,20 @@ class LevelDbTargetCache : public TargetCache {
   bool UpdateMetadata(const TargetData& target_data);
   void SaveMetadata();
 
+  /** Parses the given bytes as a `firestore_client_Target` protocol buffer. */
+  nanopb::Message<firestore_client_Target> DecodeTargetProto(
+      nanopb::Reader* reader);
+
   /**
    * Parses the given bytes as a `firestore_client_Target` protocol buffer and
    * then converts to the equivalent target data.
    */
   TargetData DecodeTarget(absl::string_view encoded);
 
+  /** Removes the given targets from the query to target mapping. */
+  void RemoveQueryTargetKeyForTargets(
+      const std::unordered_set<model::TargetId>& target_id);
+
   // The LevelDbTargetCache is owned by LevelDbPersistence.
   LevelDbPersistence* db_;
   // Owned by LevelDbPersistence.

Firestore/core/src/local/lru_garbage_collector.cc

@@ -197,9 +197,10 @@ ListenSequenceNumber LruGarbageCollector::SequenceNumberForQueryCount(
 
   RollingSequenceNumberBuffer buffer(query_count);
 
-  delegate_->EnumerateTargets([&buffer](const TargetData& target_data) {
-    buffer.AddElement(target_data.sequence_number());
-  });
+  delegate_->EnumerateTargetSequenceNumbers(
+      [&buffer](ListenSequenceNumber sequence_number) {
+        buffer.AddElement(sequence_number);
+      });
 
   delegate_->EnumerateOrphanedDocuments(
       [&buffer](const DocumentKey&, ListenSequenceNumber sequence_number) {

Firestore/core/src/local/lru_garbage_collector.h

@@ -78,10 +78,11 @@ class LruDelegate : public ReferenceDelegate {
   virtual size_t GetSequenceNumberCount() = 0;
 
   /**
-   * Enumerates all the targets that the delegate is aware of. This is typically
-   * all of the targets in an TargetCache.
+   * Enumerates the sequence numbers of all the targets that the delegate is
+   * aware of. This is typically all sequence numbers in an TargetCache.
    */
-  virtual void EnumerateTargets(const TargetCallback& callback) = 0;
+  virtual void EnumerateTargetSequenceNumbers(
+      const SequenceNumberCallback& callback) = 0;
 
   /**
    * Enumerates all of the outstanding mutations.

Firestore/core/src/local/memory_lru_reference_delegate.cc

@@ -86,9 +86,9 @@ void MemoryLruReferenceDelegate::OnTransactionCommitted() {
   current_sequence_number_ = kListenSequenceNumberInvalid;
 }
 
-void MemoryLruReferenceDelegate::EnumerateTargets(
-    const TargetCallback& callback) {
-  return persistence_->target_cache()->EnumerateTargets(callback);
+void MemoryLruReferenceDelegate::EnumerateTargetSequenceNumbers(
+    const SequenceNumberCallback& callback) {
+  return persistence_->target_cache()->EnumerateSequenceNumbers(callback);
 }
 
 void MemoryLruReferenceDelegate::EnumerateOrphanedDocuments(

Firestore/core/src/local/memory_lru_reference_delegate.h

@@ -68,7 +68,8 @@ class MemoryLruReferenceDelegate : public LruDelegate {
   util::StatusOr<int64_t> CalculateByteSize() override;
   size_t GetSequenceNumberCount() override;
 
-  void EnumerateTargets(const TargetCallback& callback) override;
+  void EnumerateTargetSequenceNumbers(
+      const SequenceNumberCallback& callback) override;
   void EnumerateOrphanedDocuments(
       const OrphanedDocumentCallback& callback) override;
 

Firestore/core/src/local/memory_target_cache.cc

@@ -68,13 +68,14 @@ absl::optional<TargetData> MemoryTargetCache::GetTarget(const Target& target) {
   return iter == targets_.end() ? absl::optional<TargetData>{} : iter->second;
 }
 
-void MemoryTargetCache::EnumerateTargets(const TargetCallback& callback) {
+void MemoryTargetCache::EnumerateSequenceNumbers(
+    const SequenceNumberCallback& callback) {
   for (const auto& kv : targets_) {
-    callback(kv.second);
+    callback(kv.second.sequence_number());
   }
 }
 
-int MemoryTargetCache::RemoveTargets(
+size_t MemoryTargetCache::RemoveTargets(
     model::ListenSequenceNumber upper_bound,
     const std::unordered_map<TargetId, TargetData>& live_targets) {
   std::vector<const Target*> to_remove;
@@ -93,7 +94,7 @@ int MemoryTargetCache::RemoveTargets(
   for (const Target* element : to_remove) {
     targets_.erase(*element);
   }
-  return static_cast<int>(to_remove.size());
+  return to_remove.size();
 }
 
 void MemoryTargetCache::AddMatchingKeys(const DocumentKeySet& keys,