Implement byteSize methods on lru delegates for use in GC threshold

Firestore/Example/Tests/Local/FSTLRUGarbageCollectorTests.mm

@@ -639,6 +639,28 @@ - (void)testRemoveTargetsThenGC {
   });
   [_persistence shutdown];
 }
+
+- (void)testGetsSize {
+  if ([self isTestBaseClass]) return;
+
+  [self newTestResources];
+
+  size_t initialSize = [_gc byteSize];
+
+  _persistence.run("fill cache", [&]() {
+    // Simulate a bunch of ack'd mutations
+    for (int i = 0; i < 50; i++) {
+      FSTDocument *doc = [self cacheADocumentInTransaction];
+      [self markDocumentEligibleForGCInTransaction:doc.key];
+    }
+  });
+
+  size_t finalSize = [_gc byteSize];
+  XCTAssertGreaterThan(finalSize, initialSize);
+
+  [_persistence shutdown];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.mm

@@ -38,6 +38,14 @@
 
 @implementation FSTPersistenceTestHelpers
 
++ (FSTLocalSerializer *)localSerializer {
+  // This owns the DatabaseIds since we do not have FirestoreClient instance to own them.
+  static DatabaseId database_id{"p", "d"};
+
+  FSTSerializerBeta *remoteSerializer = [[FSTSerializerBeta alloc] initWithDatabaseID:&database_id];
+  return [[FSTLocalSerializer alloc] initWithRemoteSerializer:remoteSerializer];
+}
+
 + (Path)levelDBDir {
   Path dir = util::TempDir().AppendUtf8("FSTPersistenceTestHelpers");
 
@@ -53,12 +61,7 @@ + (Path)levelDBDir {
 }
 
 + (FSTLevelDB *)levelDBPersistenceWithDir:(Path)dir {
-  // This owns the DatabaseIds since we do not have FirestoreClient instance to own them.
-  static DatabaseId database_id{"p", "d"};
-
-  FSTSerializerBeta *remoteSerializer = [[FSTSerializerBeta alloc] initWithDatabaseID:&database_id];
-  FSTLocalSerializer *serializer =
-      [[FSTLocalSerializer alloc] initWithRemoteSerializer:remoteSerializer];
+  FSTLocalSerializer *serializer = [self localSerializer];
   FSTLevelDB *db = [[FSTLevelDB alloc] initWithDirectory:std::move(dir) serializer:serializer];
   Status status = [db start];
   if (!status.ok()) {
@@ -85,7 +88,9 @@ + (FSTMemoryPersistence *)eagerGCMemoryPersistence {
 }
 
 + (FSTMemoryPersistence *)lruMemoryPersistence {
-  FSTMemoryPersistence *persistence = [FSTMemoryPersistence persistenceWithLRUGC];
+  FSTLocalSerializer *serializer = [self localSerializer];
+  FSTMemoryPersistence *persistence =
+      [FSTMemoryPersistence persistenceWithLRUGCAndSerializer:serializer];
   Status status = [persistence start];
   if (!status.ok()) {
     [NSException raise:NSInternalInconsistencyException

Firestore/Source/Local/FSTLRUGarbageCollector.h

@@ -61,6 +61,8 @@ extern const firebase::firestore::model::ListenSequenceNumber kFSTListenSequence
            (firebase::firestore::model::ListenSequenceNumber)sequenceNumber
                               liveQueries:(NSDictionary<NSNumber *, FSTQueryData *> *)liveQueries;
 
+- (size_t)byteSize;
+
 /** Access to the underlying LRU Garbage collector instance. */
 @property(strong, nonatomic, readonly) FSTLRUGarbageCollector *gc;
 
@@ -103,4 +105,6 @@ extern const firebase::firestore::model::ListenSequenceNumber kFSTListenSequence
 - (int)removeOrphanedDocumentsThroughSequenceNumber:
     (firebase::firestore::model::ListenSequenceNumber)sequenceNumber;
 
+- (size_t)byteSize;
+
 @end

Firestore/Source/Local/FSTLRUGarbageCollector.mm

@@ -119,4 +119,8 @@ - (int)removeOrphanedDocumentsThroughSequenceNumber:(ListenSequenceNumber)sequen
   return [_delegate removeOrphanedDocumentsThroughSequenceNumber:sequenceNumber];
 }
 
+- (size_t)byteSize {
+  return [_delegate byteSize];
+}
+
 @end

Firestore/Source/Local/FSTLevelDB.mm

@@ -77,6 +77,15 @@
 
 static const char *kReservedPathComponent = "firestore";
 
+@interface FSTLevelDB ()
+
+- (size_t)byteSize;
+
+@property(nonatomic, assign, getter=isStarted) BOOL started;
+@property(nonatomic, strong, readonly) FSTLocalSerializer *serializer;
+
+@end
+
 /**
  * Provides LRU functionality for leveldb persistence.
  *
@@ -234,12 +243,9 @@ - (void)limboDocumentUpdated:(const DocumentKey &)key {
   [self writeSentinelForKey:key];
 }
 
-@end
-
-@interface FSTLevelDB ()
-
-@property(nonatomic, assign, getter=isStarted) BOOL started;
-@property(nonatomic, strong, readonly) FSTLocalSerializer *serializer;
+- (size_t)byteSize {
+  return [_db byteSize];
+}
 
 @end
 
@@ -290,6 +296,19 @@ - (instancetype)initWithDirectory:(Path)directory serializer:(FSTLocalSerializer
   return self;
 }
 
+- (size_t)byteSize {
+  off_t count = 0;
+  auto iter = util::DirectoryIterator::Create(_directory);
+  for (; iter->Valid(); iter->Next()) {
+    off_t fileSize = util::FileSize(iter->file()).ValueOrDie();
+    count += fileSize;
+  }
+  HARD_ASSERT(iter->status().ok(), "Failed to iterate leveldb directory: %s",
+              iter->status().error_message().c_str());
+  HARD_ASSERT(count <= SIZE_MAX, "Overflowed counting bytes cached");
+  return count;
+}
+
 - (const std::set<std::string> &)users {
   return _users;
 }

Firestore/Source/Local/FSTMemoryMutationQueue.h

@@ -21,6 +21,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class FSTLocalSerializer;
+
 @interface FSTMemoryMutationQueue : NSObject <FSTMutationQueue>
 
 - (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence NS_DESIGNATED_INITIALIZER;
@@ -32,6 +34,8 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (BOOL)containsKey:(const firebase::firestore::model::DocumentKey &)key;
 
+- (size_t)byteSizeWithSerializer:(FSTLocalSerializer *)serializer;
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Source/Local/FSTMemoryMutationQueue.mm

@@ -16,8 +16,11 @@
 
 #import "Firestore/Source/Local/FSTMemoryMutationQueue.h"
 
+#import <Protobuf/GPBProtocolBuffers.h>
+
 #include <set>
 
+#import "Firestore/Protos/objc/firestore/local/Mutation.pbobjc.h"
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTDocumentReference.h"
 #import "Firestore/Source/Local/FSTMemoryPersistence.h"
@@ -466,6 +469,14 @@ - (NSUInteger)indexOfExistingBatchID:(BatchId)batchID action:(NSString *)action
   return (NSUInteger)index;
 }
 
+- (size_t)byteSizeWithSerializer:(FSTLocalSerializer *)serializer {
+  size_t count = 0;
+  for (FSTMutationBatch *batch in self.queue) {
+    count += [[[serializer encodedMutationBatch:batch] data] length];
+  };
+  return count;
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Source/Local/FSTMemoryPersistence.h

@@ -17,6 +17,7 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
+#import "Firestore/Source/Local/FSTLocalSerializer.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/types.h"
@@ -31,7 +32,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (instancetype)persistenceWithEagerGC;
 
-+ (instancetype)persistenceWithLRUGC;
++ (instancetype)persistenceWithLRUGCAndSerializer:(FSTLocalSerializer *)serializer;
 
 @end
 
@@ -50,7 +51,8 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FSTMemoryLRUReferenceDelegate
     : NSObject <FSTReferenceDelegate, FSTLRUDelegate, FSTTransactional>
 
-- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence;
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence
+                         serializer:(FSTLocalSerializer *)serializer;
 
 - (BOOL)isPinnedAtSequenceNumber:(firebase::firestore::model::ListenSequenceNumber)upperBound
                         document:(const firebase::firestore::model::DocumentKey &)key;

Firestore/Source/Local/FSTMemoryPersistence.mm

@@ -44,6 +44,10 @@
 
 @interface FSTMemoryPersistence ()
 
+- (FSTMemoryQueryCache *)queryCache;
+
+- (FSTMemoryRemoteDocumentCache *)remoteDocumentCache;
+
 @property(nonatomic, readonly) MutationQueues &mutationQueues;
 
 @property(nonatomic, assign, getter=isStarted) BOOL started;
@@ -79,10 +83,10 @@ + (instancetype)persistenceWithEagerGC {
   return persistence;
 }
 
-+ (instancetype)persistenceWithLRUGC {
++ (instancetype)persistenceWithLRUGCAndSerializer:(FSTLocalSerializer *)serializer {
   FSTMemoryPersistence *persistence = [[FSTMemoryPersistence alloc] init];
   persistence.referenceDelegate =
-      [[FSTMemoryLRUReferenceDelegate alloc] initWithPersistence:persistence];
+      [[FSTMemoryLRUReferenceDelegate alloc] initWithPersistence:persistence serializer:serializer];
   return persistence;
 }
 
@@ -136,7 +140,7 @@ - (ListenSequenceNumber)currentSequenceNumber {
   return queue;
 }
 
-- (id<FSTQueryCache>)queryCache {
+- (FSTMemoryQueryCache *)queryCache {
   return _queryCache;
 }
 
@@ -155,9 +159,11 @@ @implementation FSTMemoryLRUReferenceDelegate {
   FSTLRUGarbageCollector *_gc;
   FSTListenSequence *_listenSequence;
   ListenSequenceNumber _currentSequenceNumber;
+  FSTLocalSerializer *_serializer;
 }
 
-- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence {
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence
+                         serializer:(FSTLocalSerializer *)serializer {
   if (self = [super init]) {
     _persistence = persistence;
     _gc =
@@ -167,6 +173,7 @@ - (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence {
     ListenSequenceNumber highestSequenceNumber =
         _persistence.queryCache.highestListenSequenceNumber;
     _listenSequence = [[FSTListenSequence alloc] initStartingAfter:highestSequenceNumber];
+    _serializer = serializer;
   }
   return self;
 }
@@ -274,6 +281,20 @@ - (BOOL)isPinnedAtSequenceNumber:(ListenSequenceNumber)upperBound
   return NO;
 }
 
+- (size_t)byteSize {
+  // Note that this method is only used for testing because this delegate is only
+  // used for testing. The algorithm here (loop through everything, serialize it
+  // and count bytes) is inefficient and inexact, but won't run in production.
+  size_t count = 0;
+  count += [_persistence.queryCache byteSizeWithSerializer:_serializer];
+  count += [_persistence.remoteDocumentCache byteSizeWithSerializer:_serializer];
+  const MutationQueues &queues = [_persistence mutationQueues];
+  for (auto it = queues.begin(); it != queues.end(); ++it) {
+    count += [it->second byteSizeWithSerializer:_serializer];
+  }
+  return count;
+}
+
 @end
 
 @implementation FSTMemoryEagerReferenceDelegate {

Firestore/Source/Local/FSTMemoryQueryCache.h

@@ -20,6 +20,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class FSTLocalSerializer;
 @class FSTMemoryPersistence;
 
 /**
@@ -32,6 +33,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)init NS_UNAVAILABLE;
 
+- (size_t)byteSizeWithSerializer:(FSTLocalSerializer *)serializer;
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Source/Local/FSTMemoryQueryCache.mm

@@ -16,8 +16,11 @@
 
 #import "Firestore/Source/Local/FSTMemoryQueryCache.h"
 
+#import <Protobuf/GPBProtocolBuffers.h>
+
 #include <utility>
 
+#import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTMemoryPersistence.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
@@ -168,6 +171,15 @@ - (BOOL)containsKey:(const firebase::firestore::model::DocumentKey &)key {
   return [self.references containsKey:key];
 }
 
+- (size_t)byteSizeWithSerializer:(FSTLocalSerializer *)serializer {
+  __block size_t count = 0;
+  [self.queries
+      enumerateKeysAndObjectsUsingBlock:^(FSTQuery *key, FSTQueryData *queryData, BOOL *stop) {
+        count += [[[serializer encodedQueryData:queryData] data] length];
+      }];
+  return count;
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Source/Local/FSTMemoryRemoteDocumentCache.h

@@ -22,6 +22,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class FSTLocalSerializer;
 @class FSTMemoryLRUReferenceDelegate;
 
 @interface FSTMemoryRemoteDocumentCache : NSObject <FSTRemoteDocumentCache>
@@ -31,6 +32,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (int)removeOrphanedDocuments:(FSTMemoryLRUReferenceDelegate *)referenceDelegate
          throughSequenceNumber:(firebase::firestore::model::ListenSequenceNumber)upperBound;
 
+- (size_t)byteSizeWithSerializer:(FSTLocalSerializer *)serializer;
+
 @end
 
 NS_ASSUME_NONNULL_END