Online Count with tentative package private api

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/CountTest.java

@@ -17,8 +17,10 @@
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollectionWithDocs;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testFirestore;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitForException;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
@@ -108,4 +110,25 @@ public void testCanRunCountOnNonExistentCollection() {
         waitFor(collection.whereEqualTo("k", 100).count().get(AggregateSource.SERVER_DIRECT));
     assertEquals(Long.valueOf(0), snapshot.getCount());
   }
+
+  @Test
+  public void testFailWithoutNetwork() {
+    CollectionReference collection =
+        testCollectionWithDocs(
+            map(
+                "a", map("k", "a"),
+                "b", map("k", "b"),
+                "c", map("k", "c")));
+    waitFor(collection.getFirestore().disableNetwork());
+
+    Exception e = waitForException(collection.count().get(AggregateSource.SERVER_DIRECT));
+    assertTrue(e instanceof FirebaseFirestoreException);
+    assertEquals(
+        FirebaseFirestoreException.Code.UNAVAILABLE, ((FirebaseFirestoreException) e).getCode());
+
+    waitFor(collection.getFirestore().enableNetwork());
+    AggregateQuerySnapshot snapshot =
+        waitFor(collection.count().get(AggregateSource.SERVER_DIRECT));
+    assertEquals(Long.valueOf(3), snapshot.getCount());
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuery.java

@@ -18,7 +18,6 @@
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.firestore.util.Executors;
-import java.util.Objects;
 
 /**
  * A {@code AggregateQuery} computes some aggregation statistics from the result set of a base
@@ -79,7 +78,7 @@ public boolean equals(Object o) {
 
   @Override
   public int hashCode() {
-    return Objects.hash(query);
+    return query.hashCode();
   }
 
   @Override

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuerySnapshot.java

@@ -17,6 +17,13 @@
 import androidx.annotation.Nullable;
 import java.util.Objects;
 
+/**
+ * A {@code AggregateQuerySnapshot} contains results of a {@link AggregateQuery}.
+ *
+ * <p><b>Subclassing Note</b>: Cloud Firestore classes are not meant to be subclassed except for use
+ * in test mocks. Subclassing is not supported in production code and new SDK releases may break
+ * code that does so.
+ */
 public class AggregateQuerySnapshot {
 
   private final long count;
@@ -25,6 +32,10 @@ public class AggregateQuerySnapshot {
     this.count = count;
   }
 
+  /**
+   * @return The result of a document count aggregation. Returns null if no count aggregation is
+   *     available in the result.
+   */
   @Nullable
   public Long getCount() {
     return count;

firebase-firestore/src/main/java/com/google/firebase/firestore/core/OnlineQueryRunner.java

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.util.Util.isRetryableBackendError;
+
 import com.google.android.gms.tasks.OnCompleteListener;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
@@ -31,7 +33,7 @@ public class OnlineQueryRunner {
   private TaskCompletionSource<Long> taskSource = new TaskCompletionSource<>();
 
   public OnlineQueryRunner(AsyncQueue asyncQueue, RemoteStore remoteStore) {
-    backoff = new ExponentialBackoff(asyncQueue, AsyncQueue.TimerId.RETRY_TRANSACTION);
+    backoff = new ExponentialBackoff(asyncQueue, AsyncQueue.TimerId.RETRY_ONLINE_QUERY);
     this.asyncQueue = asyncQueue;
     this.remoteStore = remoteStore;
   }
@@ -61,7 +63,7 @@ private void runWithBackoff(Query query) {
   }
 
   private void handleError(Task<Long> result, Query query) {
-    if (attemptsRemaining > 0) {
+    if (attemptsRemaining > 0 && isRetryableBackendError(result.getException())) {
       runWithBackoff(query);
     } else {
       taskSource.setException(result.getException());

firebase-firestore/src/main/java/com/google/firebase/firestore/core/TransactionRunner.java

@@ -14,12 +14,12 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.util.Util.isRetryableBackendError;
+
 import androidx.annotation.NonNull;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
-import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.TransactionOptions;
-import com.google.firebase.firestore.remote.Datastore;
 import com.google.firebase.firestore.remote.RemoteStore;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.AsyncQueue.TimerId;
@@ -86,22 +86,10 @@ private void runWithBackoff() {
   }
 
   private void handleTransactionError(Task task) {
-    if (attemptsRemaining > 0 && isRetryableTransactionError(task.getException())) {
+    if (attemptsRemaining > 0 && isRetryableBackendError(task.getException())) {
       runWithBackoff();
     } else {
       taskSource.setException(task.getException());
     }
   }
-
-  private static boolean isRetryableTransactionError(Exception e) {
-    if (e instanceof FirebaseFirestoreException) {
-      // In transactions, the backend will fail outdated reads with FAILED_PRECONDITION and
-      // non-matching document versions with ABORTED. These errors should be retried.
-      FirebaseFirestoreException.Code code = ((FirebaseFirestoreException) e).getCode();
-      return code == FirebaseFirestoreException.Code.ABORTED
-          || code == FirebaseFirestoreException.Code.FAILED_PRECONDITION
-          || !Datastore.isPermanentError(((FirebaseFirestoreException) e).getCode());
-    }
-    return false;
-  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteStore.java

@@ -19,7 +19,9 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.database.collection.ImmutableSortedSet;
+import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.core.OnlineState;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.core.Transaction;
@@ -52,7 +54,7 @@
  * RemoteStore handles all interaction with the backend through a simple, clean interface. This
  * class is not thread safe and should be only called from the worker AsyncQueue.
  */
-public final class RemoteStore<TResult> implements WatchChangeAggregator.TargetMetadataProvider {
+public final class RemoteStore implements WatchChangeAggregator.TargetMetadataProvider {
 
   /** The maximum number of pending writes to allow. TODO: Negotiate this value with the backend. */
   private static final int MAX_PENDING_WRITES = 10;
@@ -751,6 +753,12 @@ public TargetData getTargetDataForTarget(int targetId) {
   }
 
   public Task<Long> runCountQuery(Query query) {
-    return datastore.runCountQuery(query);
+    if (canUseNetwork()) {
+      return datastore.runCountQuery(query);
+    }
+
+    return Tasks.forException(
+        new FirebaseFirestoreException(
+            "Failed to get result from server.", FirebaseFirestoreException.Code.UNAVAILABLE));
   }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/util/AsyncQueue.java

@@ -77,6 +77,11 @@ public enum TimerId {
      * multiple of these may be in the queue at a given time.
      */
     RETRY_TRANSACTION,
+    /**
+     * A timer used to retry queries sent to DFE. Since there can be multiple active queries,
+     * multiple of these may be in the queue at a given time.
+     */
+    RETRY_ONLINE_QUERY,
     /**
      * A timer used to monitor when a connection attempt in gRPC is unsuccessful and retry
      * accordingly.

firebase-firestore/src/main/java/com/google/firebase/firestore/util/Util.java

@@ -22,6 +22,7 @@
 import com.google.firebase.firestore.FieldPath;
 import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
+import com.google.firebase.firestore.remote.Datastore;
 import com.google.protobuf.ByteString;
 import io.grpc.Status;
 import io.grpc.StatusException;
@@ -137,6 +138,18 @@ public static Exception convertThrowableToException(Throwable t) {
     }
   }
 
+  public static boolean isRetryableBackendError(Exception e) {
+    if (e instanceof FirebaseFirestoreException) {
+      // In transactions, the backend will fail outdated reads with FAILED_PRECONDITION and
+      // non-matching document versions with ABORTED. These errors should be retried.
+      FirebaseFirestoreException.Code code = ((FirebaseFirestoreException) e).getCode();
+      return code == FirebaseFirestoreException.Code.ABORTED
+          || code == FirebaseFirestoreException.Code.FAILED_PRECONDITION
+          || !Datastore.isPermanentError(((FirebaseFirestoreException) e).getCode());
+    }
+    return false;
+  }
+
   private static final Continuation<Void, Void> VOID_ERROR_TRANSFORMER =
       task -> {
         if (task.isSuccessful()) {