feat: Add retry logic to COUNT queries (#1062)

Author: dconeybe

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

@@ -21,14 +21,17 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.api.gax.rpc.StreamController;
 import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.v1.FirestoreSettings;
 import com.google.firestore.v1.RunAggregationQueryRequest;
 import com.google.firestore.v1.RunAggregationQueryResponse;
 import com.google.firestore.v1.RunQueryRequest;
 import com.google.firestore.v1.StructuredAggregationQuery;
 import com.google.firestore.v1.Value;
 import com.google.protobuf.ByteString;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -68,25 +71,68 @@ public ApiFuture<AggregateQuerySnapshot> get() {
 
  @Nonnull
  ApiFuture<AggregateQuerySnapshot> get(@Nullable final ByteString transactionId) {
- RunAggregationQueryRequest request = toProto(transactionId);
- AggregateQueryResponseObserver responseObserver = new AggregateQueryResponseObserver();
+ AggregateQueryResponseDeliverer responseDeliverer =
+ new AggregateQueryResponseDeliverer(
+ transactionId, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
+ runQuery(responseDeliverer);
+ return responseDeliverer.getFuture();
+ }
+
+ private void runQuery(AggregateQueryResponseDeliverer responseDeliverer) {
+ RunAggregationQueryRequest request = toProto(responseDeliverer.getTransactionId());
+ AggregateQueryResponseObserver responseObserver =
+ new AggregateQueryResponseObserver(responseDeliverer);
  ServerStreamingCallable<RunAggregationQueryRequest, RunAggregationQueryResponse> callable =
  query.rpcContext.getClient().runAggregationQueryCallable();
-
  query.rpcContext.streamRequest(request, responseObserver, callable);
+ }
+
+ private final class AggregateQueryResponseDeliverer {
+
+ @Nullable private final ByteString transactionId;
+ private final long startTimeNanos;
+ private final SettableApiFuture<AggregateQuerySnapshot> future = SettableApiFuture.create();
+ private final AtomicBoolean isFutureCompleted = new AtomicBoolean(false);
+
+ AggregateQueryResponseDeliverer(@Nullable ByteString transactionId, long startTimeNanos) {
+ this.transactionId = transactionId;
+ this.startTimeNanos = startTimeNanos;
+ }
+
+ ApiFuture<AggregateQuerySnapshot> getFuture() {
+ return future;
+ }
+
+ @Nullable
+ ByteString getTransactionId() {
+ return transactionId;
+ }
+
+ long getStartTimeNanos() {
+ return startTimeNanos;
+ }
+
+ void deliverResult(long count, Timestamp readTime) {
+ if (isFutureCompleted.compareAndSet(false, true)) {
+ future.set(new AggregateQuerySnapshot(AggregateQuery.this, readTime, count));
+ }
+ }
 
- return responseObserver.getFuture();
+ void deliverError(Throwable throwable) {
+ if (isFutureCompleted.compareAndSet(false, true)) {
+ future.setException(throwable);
+ }
+ }
  }
 
  private final class AggregateQueryResponseObserver
  implements ResponseObserver<RunAggregationQueryResponse> {
 
- private final SettableApiFuture<AggregateQuerySnapshot> future = SettableApiFuture.create();
- private final AtomicBoolean isFutureNotified = new AtomicBoolean(false);
+ private final AggregateQueryResponseDeliverer responseDeliverer;
  private StreamController streamController;
 
- SettableApiFuture<AggregateQuerySnapshot> getFuture() {
- return future;
+ AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer) {
+ this.responseDeliverer = responseDeliverer;
  }
 
  @Override
@@ -96,14 +142,10 @@ public void onStart(StreamController streamController) {
 
  @Override
  public void onResponse(RunAggregationQueryResponse response) {
- // Ignore subsequent response messages. The RunAggregationQuery RPC returns a stream of
- // responses (rather than just a single response); however, only the first response of the
- // stream is actually used. Any more responses are technically errors, but since the Future
- // will have already been notified, we just drop any unexpected responses.
- if (!isFutureNotified.compareAndSet(false, true)) {
- return;
- }
+ // Close the stream to avoid it dangling, since we're not expecting any more responses.
+ streamController.cancel();
 
+ // Extract the count and read time from the RunAggregationQueryResponse.
  Timestamp readTime = Timestamp.fromProto(response.getReadTime());
  Value value = response.getResult().getAggregateFieldsMap().get(ALIAS_COUNT);
  if (value == null) {
@@ -118,19 +160,30 @@ public void onResponse(RunAggregationQueryResponse response) {
  }
  long count = value.getIntegerValue();
 
- future.set(new AggregateQuerySnapshot(AggregateQuery.this, readTime, count));
-
- // Close the stream to avoid it dangling, since we're not expecting any more responses.
- streamController.cancel();
+ // Deliver the result; even though the `RunAggregationQuery` RPC is a "streaming" RPC, meaning
+ // that `onResponse()` can be called multiple times, it _should_ only be called once for count
+ // queries. But even if it is called more than once, `responseDeliverer` will drop superfluous
+ // results.
+ responseDeliverer.deliverResult(count, readTime);
  }
 
  @Override
  public void onError(Throwable throwable) {
- if (!isFutureNotified.compareAndSet(false, true)) {
- return;
+ if (shouldRetry(throwable)) {
+ runQuery(responseDeliverer);
+ } else {
+ responseDeliverer.deliverError(throwable);
  }
+ }
 
- future.setException(throwable);
+ private boolean shouldRetry(Throwable throwable) {
+ Set<StatusCode.Code> retryableCodes =
+ FirestoreSettings.newBuilder().runAggregationQuerySettings().getRetryableCodes();
+ return query.shouldRetryQuery(
+ throwable,
+ responseDeliverer.getTransactionId(),
+ responseDeliverer.getStartTimeNanos(),
+ retryableCodes);
  }
 
  @Override

google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java

@@ -1672,27 +1672,15 @@ public void onComplete() {
  }
 
  boolean shouldRetry(DocumentSnapshot lastDocument, Throwable t) {
- if (transactionId != null) {
- // Transactional queries are retried via the transaction runner.
- return false;
- }
-
  if (lastDocument == null) {
  // Only retry if we have received a single result. Retries for RPCs with initial
  // failure are handled by Google Gax, which also implements backoff.
  return false;
  }
 
- if (!isRetryableError(t)) {
- return false;
- }
-
- if (rpcContext.getTotalRequestTimeout().isZero()) {
- return true;
- }
-
- Duration duration = Duration.ofNanos(rpcContext.getClock().nanoTime() - startTimeNanos);
- return duration.compareTo(rpcContext.getTotalRequestTimeout()) < 0;
+ Set<StatusCode.Code> retryableCodes =
+ FirestoreSettings.newBuilder().runQuerySettings().getRetryableCodes();
+ return shouldRetryQuery(t, transactionId, startTimeNanos, retryableCodes);
  }
  };
 
@@ -1831,21 +1819,42 @@ private <T> ImmutableList<T> append(ImmutableList<T> existingList, T newElement)
  }
 
  /** Verifies whether the given exception is retryable based on the RunQuery configuration. */
- private boolean isRetryableError(Throwable throwable) {
+ private boolean isRetryableError(Throwable throwable, Set<StatusCode.Code> retryableCodes) {
  if (!(throwable instanceof FirestoreException)) {
  return false;
  }
- Set<StatusCode.Code> codes =
- FirestoreSettings.newBuilder().runQuerySettings().getRetryableCodes();
  Status status = ((FirestoreException) throwable).getStatus();
- for (StatusCode.Code code : codes) {
+ for (StatusCode.Code code : retryableCodes) {
  if (code.equals(StatusCode.Code.valueOf(status.getCode().name()))) {
  return true;
  }
  }
  return false;
  }
 
+ /** Returns whether a query that failed in the given scenario should be retried. */
+ boolean shouldRetryQuery(
+ Throwable throwable,
+ @Nullable ByteString transactionId,
+ long startTimeNanos,
+ Set<StatusCode.Code> retryableCodes) {
+ if (transactionId != null) {
+ // Transactional queries are retried via the transaction runner.
+ return false;
+ }
+
+ if (!isRetryableError(throwable, retryableCodes)) {
+ return false;
+ }
+
+ if (rpcContext.getTotalRequestTimeout().isZero()) {
+ return true;
+ }
+
+ Duration duration = Duration.ofNanos(rpcContext.getClock().nanoTime() - startTimeNanos);
+ return duration.compareTo(rpcContext.getTotalRequestTimeout()) < 0;
+ }
+
  /**
  * Returns a query that counts the documents in the result set of this query.
  *

google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryCountTest.java

@@ -30,7 +30,9 @@
 import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 
+import com.google.api.core.ApiClock;
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
@@ -39,14 +41,15 @@
 import com.google.firestore.v1.RunAggregationQueryRequest;
 import com.google.firestore.v1.RunAggregationQueryResponse;
 import com.google.firestore.v1.StructuredQuery;
+import io.grpc.Status;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Matchers;
-import org.mockito.Mockito;
 import org.mockito.Spy;
 import org.mockito.runners.MockitoJUnitRunner;
 import org.threeten.bp.Duration;
@@ -58,7 +61,7 @@ public class QueryCountTest {
  private final FirestoreImpl firestoreMock =
  new FirestoreImpl(
  FirestoreOptions.newBuilder().setProjectId("test-project").build(),
- Mockito.mock(FirestoreRpc.class));
+ mock(FirestoreRpc.class));
 
  @Captor private ArgumentCaptor<RunAggregationQueryRequest> runAggregationQuery;
 
@@ -211,6 +214,117 @@ public void aggregateQuerySnapshotGetQueryShouldReturnCorrectValue() throws Exce
 
  AggregateQuery aggregateQuery = query.count();
  AggregateQuerySnapshot snapshot = aggregateQuery.get().get();
+
  assertThat(snapshot.getQuery()).isSameInstanceAs(aggregateQuery);
  }
+
+ @Test
+ public void shouldNotRetryIfExceptionIsNotFirestoreException() {
+ doAnswer(aggregationQueryResponse(new NotFirestoreException()))
+ .doAnswer(aggregationQueryResponse())
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+
+ assertThrows(ExecutionException.class, future::get);
+ }
+
+ @Test
+ public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatus() throws Exception {
+ doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
+ .doAnswer(aggregationQueryResponse(42))
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+ AggregateQuerySnapshot snapshot = future.get();
+
+ assertThat(snapshot.getCount()).isEqualTo(42);
+ }
+
+ @Test
+ public void shouldNotRetryIfExceptionIsFirestoreExceptionWithNonRetryableStatus() {
+ doReturn(Duration.ZERO).when(firestoreMock).getTotalRequestTimeout();
+ doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INVALID_ARGUMENT)))
+ .doAnswer(aggregationQueryResponse())
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+
+ assertThrows(ExecutionException.class, future::get);
+ }
+
+ @Test
+ public void
+ shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatusWithInfiniteTimeoutWindow()
+ throws Exception {
+ doReturn(Duration.ZERO).when(firestoreMock).getTotalRequestTimeout();
+ doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
+ .doAnswer(aggregationQueryResponse(42))
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+ AggregateQuerySnapshot snapshot = future.get();
+
+ assertThat(snapshot.getCount()).isEqualTo(42);
+ }
+
+ @Test
+ public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatusWithinTimeoutWindow()
+ throws Exception {
+ doReturn(Duration.ofDays(999)).when(firestoreMock).getTotalRequestTimeout();
+ doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
+ .doAnswer(aggregationQueryResponse(42))
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+ AggregateQuerySnapshot snapshot = future.get();
+
+ assertThat(snapshot.getCount()).isEqualTo(42);
+ }
+
+ @Test
+ public void
+ shouldNotRetryIfExceptionIsFirestoreExceptionWithRetryableStatusBeyondTimeoutWindow() {
+ ApiClock clockMock = mock(ApiClock.class);
+ doReturn(clockMock).when(firestoreMock).getClock();
+ doReturn(TimeUnit.SECONDS.toNanos(10))
+ .doReturn(TimeUnit.SECONDS.toNanos(20))
+ .doReturn(TimeUnit.SECONDS.toNanos(30))
+ .when(clockMock)
+ .nanoTime();
+ doReturn(Duration.ofSeconds(5)).when(firestoreMock).getTotalRequestTimeout();
+ doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
+ .doAnswer(aggregationQueryResponse(42))
+ .when(firestoreMock)
+ .streamRequest(
+ runAggregationQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ ApiFuture<AggregateQuerySnapshot> future = query.count().get();
+
+ assertThrows(ExecutionException.class, future::get);
+ }
+
+ private static final class NotFirestoreException extends Exception {}
 }