Browse Source

Migrate FSTLRUGarbageCollector to C++ LruGarbageCollector (#3685)

* Add ReferenceDelegate

* Remove Objective-C crutches from model/types.h and query_cache.h

* Add C++ LruGarbageCollector

* Add delegate bridges for LruDelegate and ReferenceDelegate

* Migrate from FSTLRUGarbageCollector to C++ LruGarbageCollector
Gil 6 năm trước cách đây
mục cha
commit
aa326a47ce

+ 10 - 0
Firestore/Example/Tests/Local/FSTLRUGarbageCollectorTests.h

@@ -18,6 +18,16 @@
 
 #import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 
+namespace firebase {
+namespace firestore {
+namespace local {
+
+class LruParams;
+
+}  // namespace local
+}  // namespace firestore
+}  // namespace firebase
+
 @protocol FSTPersistence;
 
 NS_ASSUME_NONNULL_BEGIN

+ 23 - 27
Firestore/Example/Tests/Local/FSTLRUGarbageCollectorTests.mm

@@ -24,13 +24,13 @@
 #include <vector>
 
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
-#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
 #import "Firestore/Source/Util/FSTClasses.h"
 
 #include "Firestore/core/include/firebase/firestore/timestamp.h"
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/core/query.h"
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
 #include "Firestore/core/src/firebase/firestore/local/mutation_queue.h"
 #include "Firestore/core/src/firebase/firestore/local/query_cache.h"
 #include "Firestore/core/src/firebase/firestore/local/query_data.h"
@@ -48,6 +48,7 @@ namespace core = firebase::firestore::core;
 namespace testutil = firebase::firestore::testutil;
 using firebase::Timestamp;
 using firebase::firestore::auth::User;
+using firebase::firestore::local::LruGarbageCollector;
 using firebase::firestore::local::LruParams;
 using firebase::firestore::local::LruResults;
 using firebase::firestore::local::MutationQueue;
@@ -84,7 +85,7 @@ NS_ASSUME_NONNULL_BEGIN
   RemoteDocumentCache *_documentCache;
   MutationQueue *_mutationQueue;
   id<FSTLRUDelegate> _lruDelegate;
-  FSTLRUGarbageCollector *_gc;
+  LruGarbageCollector *_gc;
   ListenSequenceNumber _initialSequenceNumber;
   User _user;
   ReferenceSet _additionalReferences;
@@ -140,28 +141,26 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (ListenSequenceNumber)sequenceNumberForQueryCount:(int)queryCount {
   return _persistence.run(
-      "gc", [&]() -> ListenSequenceNumber { return [_gc sequenceNumberForQueryCount:queryCount]; });
+      "gc", [&]() -> ListenSequenceNumber { return _gc->SequenceNumberForQueryCount(queryCount); });
 }
 
 - (int)queryCountForPercentile:(int)percentile {
   return _persistence.run("query count",
-                          [&]() -> int { return [_gc queryCountForPercentile:percentile]; });
+                          [&]() -> int { return _gc->QueryCountForPercentile(percentile); });
 }
 
 - (int)removeQueriesThroughSequenceNumber:(ListenSequenceNumber)sequenceNumber
                               liveQueries:
                                   (const std::unordered_map<TargetId, QueryData> &)liveQueries {
-  return _persistence.run("gc", [&]() -> int {
-    return [_gc removeQueriesUpThroughSequenceNumber:sequenceNumber liveQueries:liveQueries];
-  });
+  return _persistence.run("gc",
+                          [&]() -> int { return _gc->RemoveTargets(sequenceNumber, liveQueries); });
 }
 
 // Removes documents that are not part of a target or a mutation and have a sequence number
 // less than or equal to the given sequence number.
 - (int)removeOrphanedDocumentsThroughSequenceNumber:(ListenSequenceNumber)sequenceNumber {
-  return _persistence.run("gc", [&]() -> int {
-    return [_gc removeOrphanedDocumentsThroughSequenceNumber:sequenceNumber];
-  });
+  return _persistence.run("gc",
+                          [&]() -> int { return _gc->RemoveOrphanedDocuments(sequenceNumber); });
 }
 
 - (QueryData)nextTestQuery {
@@ -284,7 +283,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   // No queries... should get invalid sequence number (-1)
   [self newTestResources];
-  XCTAssertEqual(kFSTListenSequenceNumberInvalid, [self sequenceNumberForQueryCount:0]);
+  XCTAssertEqual(local::kListenSequenceNumberInvalid, [self sequenceNumberForQueryCount:0]);
   [_persistence shutdown];
 }
 
@@ -672,7 +671,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   [self newTestResources];
 
-  size_t initialSize = [_gc byteSize];
+  size_t initialSize = _gc->CalculateByteSize();
 
   _persistence.run("fill cache", [&]() {
     // Simulate a bunch of ack'd mutations
@@ -682,7 +681,7 @@ NS_ASSUME_NONNULL_BEGIN
     }
   });
 
-  size_t finalSize = [_gc byteSize];
+  size_t finalSize = _gc->CalculateByteSize();
   XCTAssertGreaterThan(finalSize, initialSize);
 
   [_persistence shutdown];
@@ -702,9 +701,8 @@ NS_ASSUME_NONNULL_BEGIN
     }
   });
 
-  LruResults results =
-      _persistence.run("GC", [&]() -> LruResults { return [_gc collectWithLiveTargets:{}]; });
-  XCTAssertFalse(results.didRun);
+  LruResults results = _persistence.run("GC", [&]() -> LruResults { return _gc->Collect({}); });
+  XCTAssertFalse(results.did_run);
 
   [_persistence shutdown];
 }
@@ -723,14 +721,13 @@ NS_ASSUME_NONNULL_BEGIN
     }
   });
 
-  int cacheSize = (int)[_gc byteSize];
+  int cacheSize = (int)_gc->CalculateByteSize();
   // Verify that we don't have enough in our cache to warrant collection
-  XCTAssertLessThan(cacheSize, params.minBytesThreshold);
+  XCTAssertLessThan(cacheSize, params.min_bytes_threshold);
 
   // Try collection and verify that it didn't run
-  LruResults results =
-      _persistence.run("GC", [&]() -> LruResults { return [_gc collectWithLiveTargets:{}]; });
-  XCTAssertFalse(results.didRun);
+  LruResults results = _persistence.run("GC", [&]() -> LruResults { return _gc->Collect({}); });
+  XCTAssertFalse(results.did_run);
 
   [_persistence shutdown];
 }
@@ -740,7 +737,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   LruParams params = LruParams::Default();
   // Set a low threshold so we will definitely run
-  params.minBytesThreshold = 100;
+  params.min_bytes_threshold = 100;
   [self newTestResourcesWithLruParams:params];
 
   // Add 100 targets and 10 documents to each
@@ -757,14 +754,13 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   // Mark nothing as live, so everything is eligible.
-  LruResults results =
-      _persistence.run("GC", [&]() -> LruResults { return [_gc collectWithLiveTargets:{}]; });
+  LruResults results = _persistence.run("GC", [&]() -> LruResults { return _gc->Collect({}); });
 
   // By default, we collect 10% of the sequence numbers. Since we added 100 targets,
   // that should be 10 targets with 10 documents each, for a total of 100 documents.
-  XCTAssertTrue(results.didRun);
-  XCTAssertEqual(10, results.targetsRemoved);
-  XCTAssertEqual(100, results.documentsRemoved);
+  XCTAssertTrue(results.did_run);
+  XCTAssertEqual(10, results.targets_removed);
+  XCTAssertEqual(100, results.documents_removed);
   [_persistence shutdown];
 }
 

+ 10 - 1
Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.h

@@ -16,12 +16,21 @@
 
 #import <Foundation/Foundation.h>
 
-#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #include "Firestore/core/src/firebase/firestore/util/path.h"
 
 @class FSTLevelDB;
 @class FSTMemoryPersistence;
 
+namespace firebase {
+namespace firestore {
+namespace local {
+
+class LruParams;
+
+}  // namespace local
+}  // namespace firestore
+}  // namespace firebase
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTPersistenceTestHelpers : NSObject

+ 1 - 1
Firestore/Example/Tests/Util/FSTHelpers.h

@@ -182,7 +182,7 @@ class TestTargetMetadataProvider : public TargetMetadataProvider {
       const std::vector<model::TargetId> &limbo_targets);
 
   /**
-   * Creates an `TestTargetMetadataProvider` that behaves as if there's an established listen for
+   * Creates a `TestTargetMetadataProvider` that behaves as if there's an established listen for
    * each of the given targets, where each target has not seen any previous document.
    *
    * Internally this means that the `GetRemoteKeysForTarget` callback for these targets will return

+ 3 - 83
Firestore/Source/Local/FSTLRUGarbageCollector.h

@@ -25,53 +25,18 @@
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/types.h"
 
-@class FSTLRUGarbageCollector;
-
-namespace local = firebase::firestore::local;
-namespace model = firebase::firestore::model;
-
-extern const model::ListenSequenceNumber kFSTListenSequenceNumberInvalid;
-
 namespace firebase {
 namespace firestore {
 namespace local {
 
-struct LruParams {
-  static LruParams Default() {
-    return LruParams{100 * 1024 * 1024, 10, 1000};
-  }
-
-  static LruParams Disabled() {
-    return LruParams{api::Settings::CacheSizeUnlimited, 0, 0};
-  }
-
-  static LruParams WithCacheSize(int64_t cacheSize) {
-    LruParams params = Default();
-    params.minBytesThreshold = cacheSize;
-    return params;
-  }
-
-  int64_t minBytesThreshold;
-  int percentileToCollect;
-  int maximumSequenceNumbersToCollect;
-};
-
-struct LruResults {
-  static LruResults DidNotRun() {
-    return LruResults{/* didRun= */ false, 0, 0, 0};
-  }
-
-  bool didRun;
-  int sequenceNumbersCollected;
-  int targetsRemoved;
-  int documentsRemoved;
-};
+class LruGarbageCollector;
 
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase
 
 namespace local = firebase::firestore::local;
+namespace model = firebase::firestore::model;
 
 /**
  * Persistence layers intending to use LRU Garbage collection should implement this protocol. This
@@ -111,51 +76,6 @@ namespace local = firebase::firestore::local;
 - (size_t)sequenceNumberCount;
 
 /** Access to the underlying LRU Garbage collector instance. */
-@property(strong, nonatomic, readonly) FSTLRUGarbageCollector *gc;
-
-@end
-
-/**
- * FSTLRUGarbageCollector defines the LRU algorithm used to clean up old documents and targets. It
- * is persistence-agnostic, as long as proper delegate is provided.
- */
-@interface FSTLRUGarbageCollector : NSObject
-
-- (instancetype)initWithDelegate:(id<FSTLRUDelegate>)delegate
-                          params:(local::LruParams)params NS_DESIGNATED_INITIALIZER;
-
-- (instancetype)init NS_UNAVAILABLE;
-
-/**
- * Given a target percentile, return the number of queries that make up that percentage of the
- * queries that are cached. For instance, if 20 queries are cached, and the percentile is 40, the
- * result will be 8.
- */
-- (int)queryCountForPercentile:(NSUInteger)percentile;
-
-/**
- * Given a number of queries n, return the nth sequence number in the cache.
- */
-- (model::ListenSequenceNumber)sequenceNumberForQueryCount:(NSUInteger)queryCount;
-
-/**
- * Removes queries that are not currently live (as indicated by presence in the liveQueries map) and
- * have a sequence number less than or equal to the given sequence number.
- */
-- (int)removeQueriesUpThroughSequenceNumber:(model::ListenSequenceNumber)sequenceNumber
-                                liveQueries:
-                                    (const std::unordered_map<model::TargetId, local::QueryData> &)
-                                        liveQueries;
-
-/**
- * Removes all unreferenced documents from the cache that have a sequence number less than or equal
- * to the given sequence number. Returns the number of documents removed.
- */
-- (int)removeOrphanedDocumentsThroughSequenceNumber:(model::ListenSequenceNumber)sequenceNumber;
-
-- (size_t)byteSize;
-
-- (local::LruResults)collectWithLiveTargets:
-    (const std::unordered_map<model::TargetId, local::QueryData> &)liveTargets;
+- (local::LruGarbageCollector *)gc;
 
 @end

+ 0 - 191
Firestore/Source/Local/FSTLRUGarbageCollector.mm

@@ -1,191 +0,0 @@
-/*
- * Copyright 2018 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
-
-#include <chrono>  //NOLINT(build/c++11)
-#include <queue>
-#include <utility>
-
-#import "Firestore/Source/Local/FSTPersistence.h"
-#include "Firestore/core/include/firebase/firestore/timestamp.h"
-#include "Firestore/core/src/firebase/firestore/api/settings.h"
-#include "Firestore/core/src/firebase/firestore/model/document_key.h"
-#include "Firestore/core/src/firebase/firestore/util/log.h"
-
-namespace api = firebase::firestore::api;
-using Millis = std::chrono::milliseconds;
-using firebase::Timestamp;
-using firebase::firestore::local::LruParams;
-using firebase::firestore::local::LruResults;
-using firebase::firestore::local::QueryData;
-using firebase::firestore::model::DocumentKey;
-using firebase::firestore::model::ListenSequenceNumber;
-using firebase::firestore::model::TargetId;
-
-const ListenSequenceNumber kFSTListenSequenceNumberInvalid = -1;
-
-static Millis::rep millisecondsBetween(const Timestamp &start, const Timestamp &end) {
-  return std::chrono::duration_cast<Millis>(end.ToTimePoint() - start.ToTimePoint()).count();
-}
-
-namespace {
-
-/**
- * RollingSequenceNumberBuffer tracks the nth sequence number in a series. Sequence numbers may be
- * added out of order.
- */
-class RollingSequenceNumberBuffer {
- public:
-  explicit RollingSequenceNumberBuffer(size_t max_elements)
-      : queue_(std::priority_queue<ListenSequenceNumber>()), max_elements_(max_elements) {
-  }
-
-  RollingSequenceNumberBuffer(const RollingSequenceNumberBuffer &other) = delete;
-
-  RollingSequenceNumberBuffer &operator=(const RollingSequenceNumberBuffer &other) = delete;
-
-  void AddElement(ListenSequenceNumber sequence_number) {
-    if (queue_.size() < max_elements_) {
-      queue_.push(sequence_number);
-    } else {
-      ListenSequenceNumber highestValue = queue_.top();
-      if (sequence_number < highestValue) {
-        queue_.pop();
-        queue_.push(sequence_number);
-      }
-    }
-  }
-
-  ListenSequenceNumber max_value() const {
-    return queue_.top();
-  }
-
-  size_t size() const {
-    return queue_.size();
-  }
-
- private:
-  std::priority_queue<ListenSequenceNumber> queue_;
-  const size_t max_elements_;
-};
-
-}  // namespace
-
-@implementation FSTLRUGarbageCollector {
-  __weak id<FSTLRUDelegate> _delegate;
-  LruParams _params;
-}
-
-- (instancetype)initWithDelegate:(id<FSTLRUDelegate>)delegate params:(LruParams)params {
-  self = [super init];
-  if (self) {
-    _delegate = delegate;
-    _params = std::move(params);
-  }
-  return self;
-}
-
-- (LruResults)collectWithLiveTargets:(const std::unordered_map<TargetId, QueryData> &)liveTargets {
-  if (_params.minBytesThreshold == api::Settings::CacheSizeUnlimited) {
-    LOG_DEBUG("Garbage collection skipped; disabled");
-    return LruResults::DidNotRun();
-  }
-
-  size_t currentSize = [self byteSize];
-  if (currentSize < _params.minBytesThreshold) {
-    // Not enough on disk to warrant collection. Wait another timeout cycle.
-    LOG_DEBUG("Garbage collection skipped; Cache size %s is lower than threshold %s", currentSize,
-              _params.minBytesThreshold);
-    return LruResults::DidNotRun();
-  } else {
-    LOG_DEBUG("Running garbage collection on cache of size: %s", currentSize);
-    return [self runGCWithLiveTargets:liveTargets];
-  }
-}
-
-- (LruResults)runGCWithLiveTargets:(const std::unordered_map<TargetId, QueryData> &)liveTargets {
-  Timestamp start = Timestamp::Now();
-  int sequenceNumbers = [self queryCountForPercentile:_params.percentileToCollect];
-  // Cap at the configured max
-  if (sequenceNumbers > _params.maximumSequenceNumbersToCollect) {
-    sequenceNumbers = _params.maximumSequenceNumbersToCollect;
-  }
-  Timestamp countedTargets = Timestamp::Now();
-
-  ListenSequenceNumber upperBound = [self sequenceNumberForQueryCount:sequenceNumbers];
-  Timestamp foundUpperBound = Timestamp::Now();
-
-  int numTargetsRemoved = [self removeQueriesUpThroughSequenceNumber:upperBound
-                                                         liveQueries:liveTargets];
-  Timestamp removedTargets = Timestamp::Now();
-
-  int numDocumentsRemoved = [self removeOrphanedDocumentsThroughSequenceNumber:upperBound];
-  Timestamp removedDocuments = Timestamp::Now();
-
-  std::string desc = "LRU Garbage Collection:\n";
-  absl::StrAppend(&desc, "\tCounted targets in ", millisecondsBetween(start, countedTargets),
-                  "ms\n");
-  absl::StrAppend(&desc, "\tDetermined least recently used ", sequenceNumbers,
-                  " sequence numbers in ", millisecondsBetween(countedTargets, foundUpperBound),
-                  "ms\n");
-  absl::StrAppend(&desc, "\tRemoved ", numTargetsRemoved, " targets in ",
-                  millisecondsBetween(foundUpperBound, removedTargets), "ms\n");
-  absl::StrAppend(&desc, "\tRemoved ", numDocumentsRemoved, " documents in ",
-                  millisecondsBetween(removedTargets, removedDocuments), "ms\n");
-  absl::StrAppend(&desc, "Total duration: ", millisecondsBetween(start, removedDocuments), "ms");
-  LOG_DEBUG(desc.c_str());
-
-  return LruResults{/* didRun= */ true, sequenceNumbers, numTargetsRemoved, numDocumentsRemoved};
-}
-
-- (int)queryCountForPercentile:(NSUInteger)percentile {
-  size_t totalCount = [_delegate sequenceNumberCount];
-  int setSize = (int)((percentile / 100.0f) * totalCount);
-  return setSize;
-}
-
-- (ListenSequenceNumber)sequenceNumberForQueryCount:(NSUInteger)queryCount {
-  if (queryCount == 0) {
-    return kFSTListenSequenceNumberInvalid;
-  }
-  RollingSequenceNumberBuffer buffer(queryCount);
-
-  [_delegate enumerateTargetsUsingCallback:[&buffer](const QueryData &queryData) {
-    buffer.AddElement(queryData.sequence_number());
-  }];
-  [_delegate enumerateMutationsUsingCallback:[&buffer](const DocumentKey &docKey,
-                                                       ListenSequenceNumber sequenceNumber) {
-    buffer.AddElement(sequenceNumber);
-  }];
-  return buffer.max_value();
-}
-
-- (int)removeQueriesUpThroughSequenceNumber:(ListenSequenceNumber)sequenceNumber
-                                liveQueries:
-                                    (const std::unordered_map<TargetId, QueryData> &)liveQueries {
-  return [_delegate removeTargetsThroughSequenceNumber:sequenceNumber liveQueries:liveQueries];
-}
-
-- (int)removeOrphanedDocumentsThroughSequenceNumber:(ListenSequenceNumber)sequenceNumber {
-  return [_delegate removeOrphanedDocumentsThroughSequenceNumber:sequenceNumber];
-}
-
-- (size_t)byteSize {
-  return [_delegate byteSize];
-}
-
-@end

+ 2 - 0
Firestore/Source/Local/FSTLevelDB.h

@@ -22,8 +22,10 @@
 
 #import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
+
 #include "Firestore/core/src/firebase/firestore/core/database_info.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_transaction.h"
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
 #include "Firestore/core/src/firebase/firestore/util/path.h"
 #include "Firestore/core/src/firebase/firestore/util/status.h"
 #include "Firestore/core/src/firebase/firestore/util/statusor.h"

+ 13 - 8
Firestore/Source/Local/FSTLevelDB.mm

@@ -36,6 +36,8 @@
 #include "Firestore/core/src/firebase/firestore/local/leveldb_transaction.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_util.h"
 #include "Firestore/core/src/firebase/firestore/local/listen_sequence.h"
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_delegate.h"
 #include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/local/remote_document_cache.h"
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
@@ -72,6 +74,7 @@ using firebase::firestore::local::LevelDbQueryCache;
 using firebase::firestore::local::LevelDbRemoteDocumentCache;
 using firebase::firestore::local::LevelDbTransaction;
 using firebase::firestore::local::ListenSequence;
+using firebase::firestore::local::LruGarbageCollector;
 using firebase::firestore::local::LruParams;
 using firebase::firestore::local::OrphanedDocumentCallback;
 using firebase::firestore::local::QueryData;
@@ -123,7 +126,8 @@ static const char *kReservedPathComponent = "firestore";
 @end
 
 @implementation FSTLevelDBLRUDelegate {
-  FSTLRUGarbageCollector *_gc;
+  local::LruDelegateBridge _delegateBridge;
+  std::unique_ptr<LruGarbageCollector> _gc;
   // This delegate should have the same lifetime as the persistence layer, but mark as
   // weak to avoid retain cycle.
   __weak FSTLevelDB *_db;
@@ -135,9 +139,10 @@ static const char *kReservedPathComponent = "firestore";
 
 - (instancetype)initWithPersistence:(FSTLevelDB *)persistence lruParams:(LruParams)lruParams {
   if (self = [super init]) {
-    _gc = [[FSTLRUGarbageCollector alloc] initWithDelegate:self params:lruParams];
+    _delegateBridge = local::LruDelegateBridge(self);
+    _gc = absl::make_unique<LruGarbageCollector>(&_delegateBridge, lruParams);
     _db = persistence;
-    _currentSequenceNumber = kFSTListenSequenceNumberInvalid;
+    _currentSequenceNumber = local::kListenSequenceNumberInvalid;
   }
   return self;
 }
@@ -148,17 +153,17 @@ static const char *kReservedPathComponent = "firestore";
 }
 
 - (void)transactionWillStart {
-  HARD_ASSERT(_currentSequenceNumber == kFSTListenSequenceNumberInvalid,
+  HARD_ASSERT(_currentSequenceNumber == local::kListenSequenceNumberInvalid,
               "Previous sequence number is still in effect");
   _currentSequenceNumber = _listenSequence->Next();
 }
 
 - (void)transactionWillCommit {
-  _currentSequenceNumber = kFSTListenSequenceNumberInvalid;
+  _currentSequenceNumber = local::kListenSequenceNumberInvalid;
 }
 
 - (ListenSequenceNumber)currentSequenceNumber {
-  HARD_ASSERT(_currentSequenceNumber != kFSTListenSequenceNumberInvalid,
+  HARD_ASSERT(_currentSequenceNumber != local::kListenSequenceNumberInvalid,
               "Asking for a sequence number outside of a transaction");
   return _currentSequenceNumber;
 }
@@ -252,8 +257,8 @@ static const char *kReservedPathComponent = "firestore";
   return totalCount;
 }
 
-- (FSTLRUGarbageCollector *)gc {
-  return _gc;
+- (local::LruGarbageCollector *)gc {
+  return _gc.get();
 }
 
 - (void)writeSentinelForKey:(const DocumentKey &)key {

+ 2 - 3
Firestore/Source/Local/FSTLocalStore.h

@@ -18,11 +18,10 @@
 
 #include <vector>
 
-#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
-
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/local/local_view_changes.h"
 #include "Firestore/core/src/firebase/firestore/local/local_write_result.h"
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
 #include "Firestore/core/src/firebase/firestore/local/query_data.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
@@ -200,7 +199,7 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (model::BatchId)getHighestUnacknowledgedBatchId;
 
-- (local::LruResults)collectGarbage:(FSTLRUGarbageCollector *)garbageCollector;
+- (local::LruResults)collectGarbage:(local::LruGarbageCollector *)garbageCollector;
 
 @end
 

+ 2 - 3
Firestore/Source/Local/FSTLocalStore.mm

@@ -23,7 +23,6 @@
 #include <vector>
 
 #import "FIRTimestamp.h"
-#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
 
 #include "Firestore/core/include/firebase/firestore/timestamp.h"
@@ -523,9 +522,9 @@ static const int64_t kResumeTokenMaxAgeSeconds = 5 * 60;  // 5 minutes
   _mutationQueue->RemoveMutationBatch(batch);
 }
 
-- (LruResults)collectGarbage:(FSTLRUGarbageCollector *)garbageCollector {
+- (LruResults)collectGarbage:(local::LruGarbageCollector *)garbageCollector {
   return self.persistence.run("Collect garbage",
-                              [&] { return [garbageCollector collectWithLiveTargets:_targetIDs]; });
+                              [&] { return garbageCollector->Collect(_targetIDs); });
 }
 
 @end

+ 1 - 0
Firestore/Source/Local/FSTMemoryPersistence.h

@@ -20,6 +20,7 @@
 #import "Firestore/Source/Local/FSTLocalSerializer.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
 
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/types.h"
 

+ 12 - 8
Firestore/Source/Local/FSTMemoryPersistence.mm

@@ -24,6 +24,7 @@
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/local/index_manager.h"
 #include "Firestore/core/src/firebase/firestore/local/listen_sequence.h"
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
 #include "Firestore/core/src/firebase/firestore/local/memory_index_manager.h"
 #include "Firestore/core/src/firebase/firestore/local/memory_mutation_queue.h"
 #include "Firestore/core/src/firebase/firestore/local/memory_query_cache.h"
@@ -36,6 +37,7 @@
 using firebase::firestore::auth::HashUser;
 using firebase::firestore::auth::User;
 using firebase::firestore::local::ListenSequence;
+using firebase::firestore::local::LruGarbageCollector;
 using firebase::firestore::local::LruParams;
 using firebase::firestore::local::MemoryIndexManager;
 using firebase::firestore::local::MemoryMutationQueue;
@@ -171,6 +173,7 @@ NS_ASSUME_NONNULL_BEGIN
 @end
 
 @implementation FSTMemoryLRUReferenceDelegate {
+  local::LruDelegateBridge _delegateBridge;
   // This delegate should have the same lifetime as the persistence layer, but mark as
   // weak to avoid retain cycle.
   __weak FSTMemoryPersistence *_persistence;
@@ -178,7 +181,7 @@ NS_ASSUME_NONNULL_BEGIN
   // the leveldb implementation.
   std::unordered_map<DocumentKey, ListenSequenceNumber, DocumentKeyHash> _sequenceNumbers;
   ReferenceSet *_additionalReferences;
-  FSTLRUGarbageCollector *_gc;
+  std::unique_ptr<LruGarbageCollector> _gc;
   // PORTING NOTE: when this class is ported to C++, this does not need to be a pointer
   std::unique_ptr<ListenSequence> _listenSequence;
   ListenSequenceNumber _currentSequenceNumber;
@@ -190,8 +193,9 @@ NS_ASSUME_NONNULL_BEGIN
                           lruParams:(firebase::firestore::local::LruParams)lruParams {
   if (self = [super init]) {
     _persistence = persistence;
-    _gc = [[FSTLRUGarbageCollector alloc] initWithDelegate:self params:lruParams];
-    _currentSequenceNumber = kFSTListenSequenceNumberInvalid;
+    _delegateBridge = local::LruDelegateBridge(self);
+    _gc = absl::make_unique<LruGarbageCollector>(&_delegateBridge, lruParams);
+    _currentSequenceNumber = local::kListenSequenceNumberInvalid;
     // Theoretically this is always 0, since this is all in-memory...
     ListenSequenceNumber highestSequenceNumber =
         _persistence.queryCache->highest_listen_sequence_number();
@@ -201,12 +205,12 @@ NS_ASSUME_NONNULL_BEGIN
   return self;
 }
 
-- (FSTLRUGarbageCollector *)gc {
-  return _gc;
+- (local::LruGarbageCollector *)gc {
+  return _gc.get();
 }
 
 - (ListenSequenceNumber)currentSequenceNumber {
-  HARD_ASSERT(_currentSequenceNumber != kFSTListenSequenceNumberInvalid,
+  HARD_ASSERT(_currentSequenceNumber != local::kListenSequenceNumberInvalid,
               "Asking for a sequence number outside of a transaction");
   return _currentSequenceNumber;
 }
@@ -232,7 +236,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)commitTransaction {
-  _currentSequenceNumber = kFSTListenSequenceNumberInvalid;
+  _currentSequenceNumber = local::kListenSequenceNumberInvalid;
 }
 
 - (void)enumerateTargetsUsingCallback:(const TargetCallback &)callback {
@@ -348,7 +352,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (ListenSequenceNumber)currentSequenceNumber {
-  return kFSTListenSequenceNumberInvalid;
+  return local::kListenSequenceNumberInvalid;
 }
 
 - (void)addInMemoryPins:(ReferenceSet *)set {

+ 1 - 1
Firestore/Source/Remote/FSTSerializerBeta.mm

@@ -534,7 +534,7 @@ absl::any Wrap(GCFSDocument *doc) {
       proto.transform.fieldTransformsArray =
           [self encodedFieldTransforms:transform.field_transforms()];
       // NOTE: We set a precondition of exists: true as a safety-check, since we always combine
-      // TransformMutations with an SetMutation or PatchMutation which (if successful) should
+      // TransformMutations with a SetMutation or PatchMutation which (if successful) should
       // end up with an existing document.
       proto.currentDocument.exists = YES;
       return proto;

+ 3 - 0
Firestore/core/src/firebase/firestore/local/CMakeLists.txt

@@ -57,6 +57,8 @@ cc_library(
     local_serializer.h
     local_view_changes.cc
     local_view_changes.h
+    lru_garbage_collector.cc
+    lru_garbage_collector.h
     memory_index_manager.cc
     memory_index_manager.h
     memory_mutation_queue.h
@@ -69,6 +71,7 @@ cc_library(
     query_cache.h
     query_data.cc
     query_data.h
+    reference_delegate.h
     reference_set.cc
     reference_set.h
     remote_document_cache.h

+ 209 - 0
Firestore/core/src/firebase/firestore/local/lru_garbage_collector.cc

@@ -0,0 +1,209 @@
+/*
+ * Copyright 2019 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h"
+
+#include <chrono>  // NOLINT(build/c++11)
+#include <queue>
+#include <string>
+#include <utility>
+
+#include "Firestore/core/include/firebase/firestore/timestamp.h"
+#include "Firestore/core/src/firebase/firestore/api/settings.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/util/log.h"
+
+namespace firebase {
+namespace firestore {
+namespace local {
+namespace {
+
+using firebase::firestore::api::Settings;
+using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::ListenSequenceNumber;
+using firebase::firestore::model::TargetId;
+
+using Millis = std::chrono::milliseconds;
+
+static Millis::rep MillisecondsBetween(const Timestamp& start,
+                                       const Timestamp& end) {
+  return std::chrono::duration_cast<Millis>(end.ToTimePoint() -
+                                            start.ToTimePoint())
+      .count();
+}
+
+/**
+ * RollingSequenceNumberBuffer tracks the nth sequence number in a series.
+ * Sequence numbers may be added out of order.
+ */
+class RollingSequenceNumberBuffer {
+ public:
+  explicit RollingSequenceNumberBuffer(size_t max_elements)
+      : queue_(std::priority_queue<ListenSequenceNumber>()),
+        max_elements_(max_elements) {
+  }
+
+  RollingSequenceNumberBuffer(const RollingSequenceNumberBuffer& other) =
+      delete;
+
+  void AddElement(ListenSequenceNumber sequence_number) {
+    if (queue_.size() < max_elements_) {
+      queue_.push(sequence_number);
+    } else {
+      ListenSequenceNumber highestValue = queue_.top();
+      if (sequence_number < highestValue) {
+        queue_.pop();
+        queue_.push(sequence_number);
+      }
+    }
+  }
+
+  ListenSequenceNumber max_value() const {
+    return queue_.top();
+  }
+
+  size_t size() const {
+    return queue_.size();
+  }
+
+ private:
+  std::priority_queue<ListenSequenceNumber> queue_;
+  const size_t max_elements_;
+};
+
+}  // namespace
+
+const ListenSequenceNumber kListenSequenceNumberInvalid = -1;
+
+LruParams LruParams::Default() {
+  return LruParams{100 * 1024 * 1024, 10, 1000};
+}
+
+LruParams LruParams::Disabled() {
+  return LruParams{api::Settings::CacheSizeUnlimited, 0, 0};
+}
+
+LruParams LruParams::WithCacheSize(int64_t cacheSize) {
+  LruParams params = Default();
+  params.min_bytes_threshold = cacheSize;
+  return params;
+}
+
+LruGarbageCollector::LruGarbageCollector(LruDelegate* delegate,
+                                         LruParams params)
+    : delegate_(delegate), params_(std::move(params)) {
+}
+
+LruResults LruGarbageCollector::Collect(const LiveQueryMap& live_targets) {
+  if (params_.min_bytes_threshold == Settings::CacheSizeUnlimited) {
+    LOG_DEBUG("Garbage collection skipped; disabled");
+    return LruResults::DidNotRun();
+  }
+
+  size_t current_size = CalculateByteSize();
+  if (current_size < static_cast<size_t>(params_.min_bytes_threshold)) {
+    // Not enough on disk to warrant collection. Wait another timeout cycle.
+    LOG_DEBUG(
+        "Garbage collection skipped; Cache size %s is lower than threshold %s",
+        current_size, params_.min_bytes_threshold);
+    return LruResults::DidNotRun();
+  }
+
+  LOG_DEBUG("Running garbage collection on cache of size: %s", current_size);
+  return RunGarbageCollection(live_targets);
+}
+
+LruResults LruGarbageCollector::RunGarbageCollection(
+    const LiveQueryMap& live_targets) {
+  Timestamp start = Timestamp::Now();
+
+  // Cap at the configured max
+  int sequence_numbers = QueryCountForPercentile(params_.percentile_to_collect);
+  if (sequence_numbers > params_.maximum_sequence_numbers_to_collect) {
+    sequence_numbers = params_.maximum_sequence_numbers_to_collect;
+  }
+  Timestamp counted_targets = Timestamp::Now();
+
+  ListenSequenceNumber upper_bound =
+      SequenceNumberForQueryCount(sequence_numbers);
+  Timestamp found_upper_bound = Timestamp::Now();
+
+  int num_targets_removed = RemoveTargets(upper_bound, live_targets);
+  Timestamp removed_targets = Timestamp::Now();
+
+  int num_documents_removed = RemoveOrphanedDocuments(upper_bound);
+  Timestamp removed_documents = Timestamp::Now();
+
+  std::string desc = "LRU Garbage Collection:\n";
+  absl::StrAppend(&desc, "\tCounted targets in ",
+                  MillisecondsBetween(start, counted_targets), "ms\n");
+  absl::StrAppend(&desc, "\tDetermined least recently used ", sequence_numbers,
+                  " sequence numbers in ",
+                  MillisecondsBetween(counted_targets, found_upper_bound),
+                  "ms\n");
+  absl::StrAppend(&desc, "\tRemoved ", num_targets_removed, " targets in ",
+                  MillisecondsBetween(found_upper_bound, removed_targets),
+                  "ms\n");
+  absl::StrAppend(&desc, "\tRemoved ", num_documents_removed, " documents in ",
+                  MillisecondsBetween(removed_targets, removed_documents),
+                  "ms\n");
+  absl::StrAppend(&desc, "Total duration: ",
+                  MillisecondsBetween(start, removed_documents), "ms");
+  LOG_DEBUG(desc.c_str());
+
+  return LruResults{/* did_run= */ true, sequence_numbers, num_targets_removed,
+                    num_documents_removed};
+}
+
+int LruGarbageCollector::QueryCountForPercentile(int percentile) {
+  size_t total_count = delegate_->GetSequenceNumberCount();
+  return static_cast<int>((percentile / 100.0f) * total_count);
+}
+
+ListenSequenceNumber LruGarbageCollector::SequenceNumberForQueryCount(
+    int query_count) {
+  if (query_count == 0) {
+    return kListenSequenceNumberInvalid;
+  }
+
+  RollingSequenceNumberBuffer buffer(query_count);
+
+  delegate_->EnumerateTargets([&buffer](const QueryData& queryData) {
+    buffer.AddElement(queryData.sequence_number());
+  });
+
+  delegate_->EnumerateOrphanedDocuments(
+      [&buffer](const DocumentKey& docKey,
+                ListenSequenceNumber sequence_number) {
+        buffer.AddElement(sequence_number);
+      });
+
+  return buffer.max_value();
+}
+
+int LruGarbageCollector::RemoveTargets(ListenSequenceNumber sequence_number,
+                                       const LiveQueryMap& live_queries) {
+  return delegate_->RemoveTargets(sequence_number, live_queries);
+}
+
+int LruGarbageCollector::RemoveOrphanedDocuments(
+    ListenSequenceNumber sequence_number) {
+  return delegate_->RemoveOrphanedDocuments(sequence_number);
+}
+
+}  // namespace local
+}  // namespace firestore
+}  // namespace firebase

+ 227 - 0
Firestore/core/src/firebase/firestore/local/lru_garbage_collector.h

@@ -0,0 +1,227 @@
+/*
+ * Copyright 2019 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_LRU_GARBAGE_COLLECTOR_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_LRU_GARBAGE_COLLECTOR_H_
+
+#if __OBJC__
+#import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
+#endif
+
+#include <unordered_map>
+
+#include "Firestore/core/src/firebase/firestore/local/query_cache.h"
+#include "Firestore/core/src/firebase/firestore/local/query_data.h"
+#include "Firestore/core/src/firebase/firestore/model/types.h"
+
+namespace firebase {
+namespace firestore {
+namespace local {
+
+class LruGarbageCollector;
+
+ABSL_CONST_INIT extern const model::ListenSequenceNumber
+    kListenSequenceNumberInvalid;
+
+struct LruParams {
+  static LruParams Default();
+
+  static LruParams Disabled();
+
+  static LruParams WithCacheSize(int64_t cacheSize);
+
+  int64_t min_bytes_threshold;
+  int percentile_to_collect;
+  int maximum_sequence_numbers_to_collect;
+};
+
+struct LruResults {
+  static LruResults DidNotRun() {
+    return LruResults{/* did_run= */ false, 0, 0, 0};
+  }
+
+  bool did_run;
+  int sequence_numbers_collected;
+  int targets_removed;
+  int documents_removed;
+};
+
+using LiveQueryMap = std::unordered_map<model::TargetId, QueryData>;
+
+/**
+ * Persistence layers intending to use LRU Garbage collection should implement
+ * this interface. This interface defines the operations that the LRU garbage
+ * collector needs from the persistence layer.
+ */
+class LruDelegate {
+ public:
+  virtual ~LruDelegate() = default;
+
+  /** Access to the underlying LRU Garbage collector instance. */
+  virtual LruGarbageCollector* garbage_collector() = 0;
+
+  virtual size_t CalculateByteSize() const = 0;
+
+  /** Returns the number of targets and orphaned documents cached. */
+  virtual size_t GetSequenceNumberCount() = 0;
+
+  /**
+   * Enumerates all the targets that the delegate is aware of. This is typically
+   * all of the targets in an FSTQueryCache.
+   */
+  virtual void EnumerateTargets(const TargetCallback& callback) = 0;
+
+  /**
+   * Enumerates all of the outstanding mutations.
+   */
+  virtual void EnumerateOrphanedDocuments(
+      const OrphanedDocumentCallback& callback) = 0;
+
+  /**
+   * Removes all unreferenced documents from the cache that have a sequence
+   * number less than or equal to the given sequence number. Returns the number
+   * of documents removed.
+   */
+  virtual int RemoveOrphanedDocuments(
+      model::ListenSequenceNumber sequence_number) = 0;
+
+  /**
+   * Removes all targets that are not currently being listened to and have a
+   * sequence number less than or equal to the given sequence number. Returns
+   * the number of targets removed.
+   */
+  virtual int RemoveTargets(model::ListenSequenceNumber sequence_number,
+                            const LiveQueryMap& live_queries) = 0;
+};
+
+#if __OBJC__
+
+class LruDelegateBridge : public LruDelegate {
+ public:
+  LruDelegateBridge() = default;
+
+  explicit LruDelegateBridge(id<FSTLRUDelegate> target) : target_(target) {
+  }
+
+  /** Access to the underlying LRU Garbage collector instance. */
+  LruGarbageCollector* garbage_collector() override {
+    return target_.gc;
+  }
+
+  size_t CalculateByteSize() const override {
+    return [target_ byteSize];
+  }
+
+  size_t GetSequenceNumberCount() override {
+    return [target_ sequenceNumberCount];
+  }
+
+  void EnumerateTargets(const TargetCallback& callback) override {
+    [target_ enumerateTargetsUsingCallback:callback];
+  }
+
+  /**
+   * Enumerates all of the outstanding mutations.
+   */
+  void EnumerateOrphanedDocuments(
+      const OrphanedDocumentCallback& callback) override {
+    [target_ enumerateMutationsUsingCallback:callback];
+  }
+
+  /**
+   * Removes all unreferenced documents from the cache that have a sequence
+   * number less than or equal to the given sequence number. Returns the number
+   * of documents removed.
+   */
+  int RemoveOrphanedDocuments(
+      model::ListenSequenceNumber sequence_number) override {
+    return
+        [target_ removeOrphanedDocumentsThroughSequenceNumber:sequence_number];
+  }
+
+  /**
+   * Removes all targets that are not currently being listened to and have a
+   * sequence number less than or equal to the given sequence number. Returns
+   * the number of targets removed.
+   */
+  int RemoveTargets(model::ListenSequenceNumber sequence_number,
+                    const LiveQueryMap& live_queries) override {
+    return [target_ removeTargetsThroughSequenceNumber:sequence_number
+                                           liveQueries:live_queries];
+  }
+
+ private:
+  id<FSTLRUDelegate> target_;
+};
+
+#endif  // __OBJC__
+
+/**
+ * LruGarbageCollector defines the LRU algorithm used to clean up old documents
+ * and targets. It is persistence-agnostic, as long as proper delegate is
+ * provided.
+ */
+class LruGarbageCollector {
+ public:
+  LruGarbageCollector(LruDelegate* delegate, LruParams params);
+
+  size_t CalculateByteSize() const {
+    return delegate_->CalculateByteSize();
+  }
+
+  /**
+   * Given a target percentile, return the number of queries that make up that
+   * percentage of the queries that are cached. For instance, if 20 queries are
+   * cached, and the percentile is 40, the result will be 8.
+   */
+  int QueryCountForPercentile(int percentile);
+
+  /**
+   * Given a number of queries n, return the nth sequence number in the cache.
+   */
+  model::ListenSequenceNumber SequenceNumberForQueryCount(int query_count);
+
+  /**
+   * Removes queries that are not currently live (as indicated by presence in
+   * the liveQueries map) and have a sequence number less than or equal to the
+   * given sequence number.
+   */
+  int RemoveTargets(model::ListenSequenceNumber sequence_number,
+                    const LiveQueryMap& live_queries);
+
+  /**
+   * Removes all unreferenced documents from the cache that have a sequence
+   * number less than or equal to the given sequence number. Returns the number
+   * of documents removed.
+   */
+  int RemoveOrphanedDocuments(model::ListenSequenceNumber sequence_number);
+
+  local::LruResults Collect(const LiveQueryMap& live_targets);
+
+ private:
+  LruResults RunGarbageCollection(const LiveQueryMap& liveTargets);
+
+  // Delegate owns the LruGarbageCollector; this is a back pointer.
+  LruDelegate* delegate_;
+
+  LruParams params_ = LruParams::Default();
+};
+
+}  // namespace local
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_LRU_GARBAGE_COLLECTOR_H_

+ 0 - 10
Firestore/core/src/firebase/firestore/local/query_cache.h

@@ -17,12 +17,6 @@
 #ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_QUERY_CACHE_H_
 #define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_QUERY_CACHE_H_
 
-#if !defined(__OBJC__)
-#error "For now, this file must only be included by ObjC source files."
-#endif  // !defined(__OBJC__)
-
-#import <Foundation/Foundation.h>
-
 #include <functional>
 #include <unordered_map>
 
@@ -33,8 +27,6 @@
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 #include "Firestore/core/src/firebase/firestore/model/types.h"
 
-NS_ASSUME_NONNULL_BEGIN
-
 namespace firebase {
 namespace firestore {
 namespace local {
@@ -154,6 +146,4 @@ class QueryCache {
 }  // namespace firestore
 }  // namespace firebase
 
-NS_ASSUME_NONNULL_END
-
 #endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_QUERY_CACHE_H_

+ 161 - 0
Firestore/core/src/firebase/firestore/local/reference_delegate.h

@@ -0,0 +1,161 @@
+/*
+ * Copyright 2019 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_REFERENCE_DELEGATE_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_REFERENCE_DELEGATE_H_
+
+#if __OBJC__
+#import "Firestore/Source/Local/FSTPersistence.h"
+#endif  // __OBJC__
+
+#include "Firestore/core/src/firebase/firestore/model/types.h"
+#include "absl/strings/string_view.h"
+
+namespace firebase {
+namespace firestore {
+namespace model {
+
+class DocumentKey;
+
+}  // namespace model
+
+namespace local {
+
+class QueryData;
+class ReferenceSet;
+
+/**
+ * A ReferenceDelegate instance handles all of the hooks into the
+ * document-reference lifecycle. This includes being added to a target, being
+ * removed from a target, being subject to mutation, and being mutated by the
+ * user.
+ *
+ * Different implementations may do different things with each of these events.
+ * Not every implementation needs to do something with every lifecycle hook.
+ *
+ * Implementations that care about sequence numbers are responsible for
+ * generating them and making them available.
+ */
+class ReferenceDelegate {
+ public:
+  virtual model::ListenSequenceNumber current_sequence_number() const = 0;
+
+  /**
+   * Registers a ReferenceSet of documents that should be considered
+   * 'referenced' and not eligible for removal during garbage collection.
+   */
+  virtual void AddInMemoryPins(ReferenceSet* set) = 0;
+
+  /**
+   * Notifies the delegate that the given document was added to a target.
+   */
+  virtual void AddReference(const model::DocumentKey& key) = 0;
+
+  /**
+   * Notifies the delegate that the given document was removed from a target.
+   */
+  virtual void RemoveReference(const model::DocumentKey& key) = 0;
+
+  /**
+   * Notifies the delegate that a document is no longer being mutated by the
+   * user.
+   */
+  virtual void RemoveMutationReference(const model::DocumentKey& key) = 0;
+
+  /**
+   * Notifies the delegate that a target was removed.
+   */
+  virtual void RemoveTarget(const local::QueryData& query_data) = 0;
+
+  /**
+   * Notifies the delegate that a limbo document was updated.
+   */
+  virtual void UpdateLimboDocument(const model::DocumentKey& key) = 0;
+
+  /**
+   * Lifecycle hook that notifies the delegate that a transaction has started.
+   */
+  virtual void OnTransactionStarted(absl::string_view label) = 0;
+
+  /**
+   * Lifecycle hook that notifies the delegate that a transaction has committed.
+   */
+  virtual void OnTransactionCommitted() = 0;
+};
+
+#if __OBJC__
+
+class ReferenceDelegateBridge : public ReferenceDelegate {
+ public:
+  ReferenceDelegateBridge() = default;
+
+  explicit ReferenceDelegateBridge(id<FSTReferenceDelegate> target)
+      : target_(target) {
+  }
+
+  model::ListenSequenceNumber current_sequence_number() const override {
+    return target_.currentSequenceNumber;
+  }
+
+  void AddInMemoryPins(ReferenceSet* set) override {
+    [target_ addInMemoryPins:set];
+  }
+
+  void AddReference(const model::DocumentKey& key) override {
+    [target_ addReference:key];
+  }
+
+  void RemoveReference(const model::DocumentKey& key) override {
+    [target_ removeReference:key];
+  }
+
+  void RemoveMutationReference(const model::DocumentKey& key) override {
+    [target_ removeMutationReference:key];
+  }
+
+  void RemoveTarget(const local::QueryData& query_data) override {
+    [target_ removeTarget:query_data];
+  }
+
+  void UpdateLimboDocument(const model::DocumentKey& key) override {
+    [target_ limboDocumentUpdated:key];
+  }
+
+  void OnTransactionStarted(absl::string_view label) override {
+    if ([target_ conformsToProtocol:@protocol(FSTTransactional)]) {
+      auto transactional = static_cast<id<FSTTransactional>>(target_);
+      [transactional startTransaction:label];
+    }
+  }
+
+  void OnTransactionCommitted() override {
+    if ([target_ conformsToProtocol:@protocol(FSTTransactional)]) {
+      auto transactional = static_cast<id<FSTTransactional>>(target_);
+      [transactional commitTransaction];
+    }
+  }
+
+ private:
+  __weak id<FSTReferenceDelegate> target_;
+};
+
+#endif  // __OBJC__
+
+}  // namespace local
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_LOCAL_REFERENCE_DELEGATE_H_

+ 0 - 4
Firestore/core/src/firebase/firestore/model/types.h

@@ -19,10 +19,6 @@
 
 #include <cstdint>
 
-#if defined(__OBJC__)
-#import <Foundation/Foundation.h>
-#endif
-
 namespace firebase {
 namespace firestore {
 namespace model {

+ 1 - 1
Firestore/core/src/firebase/firestore/util/status.cc

@@ -52,7 +52,7 @@ Status& Status::CausedBy(const Status& cause) {
   absl::StrAppend(&state_->msg, ": ", cause.error_message());
 
   // If this Status has no accompanying PlatformError but the cause does, create
-  // an PlatformError for this Status ahead of time to preserve the causal chain
+  // a PlatformError for this Status ahead of time to preserve the causal chain
   // that Status doesn't otherwise support.
   if (state_->platform_error == nullptr &&
       cause.state_->platform_error != nullptr) {

+ 1 - 1
Firestore/core/test/firebase/firestore/local/leveldb_key_test.cc

@@ -61,7 +61,7 @@ std::string DocTargetKey(absl::string_view key, TargetId target_id) {
  * description.
  *
  * @param key A StringView of a textual key
- * @param key An string that `Describe(key)` is expected to produce.
+ * @param key A string that `Describe(key)` is expected to produce.
  */
 #define AssertExpectedKeyDescription(expected_description, key) \
   ASSERT_EQ((expected_description), DescribeKey(key))

+ 1 - 1
Firestore/core/test/firebase/firestore/testutil/testutil.h

@@ -322,7 +322,7 @@ inline model::UnknownDocument UnknownDoc(absl::string_view key,
 #if __APPLE__
 
 /**
- * Creates an DocumentComparator that will compare Documents by the given
+ * Creates a DocumentComparator that will compare Documents by the given
  * fieldPath string then by key.
  */
 model::DocumentComparator DocComparator(absl::string_view field_path);