Expose operation counts from Persistence layer (#595)

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryMutationQueue.java

@@ -69,9 +69,11 @@ final class MemoryMutationQueue implements MutationQueue {
   private ByteString lastStreamToken;
 
   private final MemoryPersistence persistence;
+  private final StatsCollector statsCollector;
 
-  MemoryMutationQueue(MemoryPersistence persistence) {
+  MemoryMutationQueue(MemoryPersistence persistence, StatsCollector statsCollector) {
     this.persistence = persistence;
+    this.statsCollector = statsCollector;
     queue = new ArrayList<>();
 
     batchesByDocumentKey = new ImmutableSortedSet<>(emptyList(), DocumentReference.BY_KEY);
@@ -154,12 +156,16 @@ public MutationBatch addMutationBatch(
           .addToCollectionParentIndex(mutation.getKey().getPath().popLast());
     }
 
+    statsCollector.recordRowsWritten(STATS_TAG, 1);
+
     return batch;
   }
 
   @Nullable
   @Override
   public MutationBatch lookupMutationBatch(int batchId) {
+    statsCollector.recordRowsRead(STATS_TAG, 1);
+
     int index = indexOfBatchId(batchId);
     if (index < 0 || index >= queue.size()) {
       return null;
@@ -203,6 +209,8 @@ public List<MutationBatch> getAllMutationBatchesAffectingDocumentKey(DocumentKey
       result.add(batch);
     }
 
+    statsCollector.recordRowsRead(STATS_TAG, result.size());
+
     return result;
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryPersistence.java

@@ -35,30 +35,38 @@ public final class MemoryPersistence extends Persistence {
   private final MemoryIndexManager indexManager;
   private final MemoryQueryCache queryCache;
   private final MemoryRemoteDocumentCache remoteDocumentCache;
+  private final StatsCollector statsCollector;
   private ReferenceDelegate referenceDelegate;
 
   private boolean started;
 
   public static MemoryPersistence createEagerGcMemoryPersistence() {
-    MemoryPersistence persistence = new MemoryPersistence();
+    return createEagerGcMemoryPersistence(StatsCollector.NO_OP_STATS_COLLECTOR);
+  }
+
+  public static MemoryPersistence createEagerGcMemoryPersistence(StatsCollector statsCollector) {
+    MemoryPersistence persistence = new MemoryPersistence(statsCollector);
     persistence.setReferenceDelegate(new MemoryEagerReferenceDelegate(persistence));
     return persistence;
   }
 
   public static MemoryPersistence createLruGcMemoryPersistence(
-      LruGarbageCollector.Params params, LocalSerializer serializer) {
-    MemoryPersistence persistence = new MemoryPersistence();
+      LruGarbageCollector.Params params,
+      StatsCollector statsCollector,
+      LocalSerializer serializer) {
+    MemoryPersistence persistence = new MemoryPersistence(statsCollector);
     persistence.setReferenceDelegate(
         new MemoryLruReferenceDelegate(persistence, params, serializer));
     return persistence;
   }
 
   /** Use static helpers to instantiate */
-  private MemoryPersistence() {
+  private MemoryPersistence(StatsCollector statsCollector) {
+    this.statsCollector = statsCollector;
     mutationQueues = new HashMap<>();
     indexManager = new MemoryIndexManager();
     queryCache = new MemoryQueryCache(this);
-    remoteDocumentCache = new MemoryRemoteDocumentCache(this);
+    remoteDocumentCache = new MemoryRemoteDocumentCache(this, statsCollector);
   }
 
   @Override
@@ -93,7 +101,7 @@ private void setReferenceDelegate(ReferenceDelegate delegate) {
   MutationQueue getMutationQueue(User user) {
     MemoryMutationQueue queue = mutationQueues.get(user);
     if (queue == null) {
-      queue = new MemoryMutationQueue(this);
+      queue = new MemoryMutationQueue(this, statsCollector);
       mutationQueues.put(user, queue);
     }
     return queue;

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java

@@ -36,9 +36,11 @@ final class MemoryRemoteDocumentCache implements RemoteDocumentCache {
   private ImmutableSortedMap<DocumentKey, MaybeDocument> docs;
 
   private final MemoryPersistence persistence;
+  private StatsCollector statsCollector;
 
-  MemoryRemoteDocumentCache(MemoryPersistence persistence) {
+  MemoryRemoteDocumentCache(MemoryPersistence persistence, StatsCollector statsCollector) {
     docs = emptyMaybeDocumentMap();
+    this.statsCollector = statsCollector;
     this.persistence = persistence;
   }
 
@@ -51,12 +53,14 @@ public void add(MaybeDocument document) {
 
   @Override
   public void remove(DocumentKey key) {
+    statsCollector.recordRowsDeleted(STATS_TAG, 1);
     docs = docs.remove(key);
   }
 
   @Nullable
   @Override
   public MaybeDocument get(DocumentKey key) {
+    statsCollector.recordRowsRead(STATS_TAG, 1);
     return docs.get(key);
   }
 
@@ -70,6 +74,7 @@ public Map<DocumentKey, MaybeDocument> getAll(Iterable<DocumentKey> keys) {
       result.put(key, get(key));
     }
 
+    statsCollector.recordRowsRead(STATS_TAG, result.size());
     return result;
   }
 
@@ -85,8 +90,14 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
     ResourcePath queryPath = query.getPath();
     DocumentKey prefix = DocumentKey.fromPath(queryPath.append(""));
     Iterator<Map.Entry<DocumentKey, MaybeDocument>> iterator = docs.iteratorFrom(prefix);
+
+    int rowsRead = 0;
+
     while (iterator.hasNext()) {
       Map.Entry<DocumentKey, MaybeDocument> entry = iterator.next();
+
+      ++rowsRead;
+
       DocumentKey key = entry.getKey();
       if (!queryPath.isPrefixOf(key.getPath())) {
         break;
@@ -103,6 +114,8 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
       }
     }
 
+    statsCollector.recordRowsRead(STATS_TAG, rowsRead);
+
     return result;
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MutationQueue.java

@@ -25,6 +25,9 @@
 
 /** A queue of mutations to apply to the remote store. */
 interface MutationQueue {
+  /** The tag used by the StatsCollector. */
+  String STATS_TAG = "mutations";
+
   /**
    * Starts the mutation queue, performing any initial reads that might be required to establish
    * invariants, etc.

firebase-firestore/src/main/java/com/google/firebase/firestore/local/RemoteDocumentCache.java

@@ -30,6 +30,9 @@
  * instances (indicating that the document is known to not exist).
  */
 interface RemoteDocumentCache {
+  /** The tag used by the StatsCollector. */
+  String STATS_TAG = "remote_documents";
+
   /**
    * Adds or replaces an entry in the cache.
    *