fix: retry Query streams (#426)

Author: schmidt-sebastian

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

@@ -32,9 +32,11 @@
 import com.google.api.core.InternalExtensionOnly;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.ApiStreamObserver;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.auto.value.AutoValue;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.Query.QueryOptions.Builder;
+import com.google.cloud.firestore.v1.FirestoreSettings;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -52,14 +54,17 @@
 import com.google.firestore.v1.Value;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Int32Value;
+import io.grpc.Status;
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -1297,7 +1302,8 @@ public void onCompleted() {
  responseObserver.onCompleted();
  }
  },
- null);
+ /* transactionId= */ null,
+ /* readTime= */ null);
  }
 
  /**
@@ -1431,13 +1437,18 @@ Timestamp getReadTime() {
  }
 
  private void internalStream(
- final QuerySnapshotObserver documentObserver, @Nullable ByteString transactionId) {
+ final QuerySnapshotObserver documentObserver,
+ @Nullable final ByteString transactionId,
+ @Nullable final Timestamp readTime) {
  RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
  request.setStructuredQuery(buildQuery()).setParent(options.getParentPath().toString());
 
  if (transactionId != null) {
  request.setTransaction(transactionId);
  }
+ if (readTime != null) {
+ request.setReadTime(readTime.toProto());
+ }
 
  Tracing.getTracer()
  .getCurrentSpan()
@@ -1446,6 +1457,8 @@ private void internalStream(
  ImmutableMap.of(
  "transactional", AttributeValue.booleanAttributeValue(transactionId != null)));
 
+ final AtomicReference<QueryDocumentSnapshot> lastReceivedDocument = new AtomicReference<>();
+
  ApiStreamObserver<RunQueryResponse> observer =
  new ApiStreamObserver<RunQueryResponse>() {
  Timestamp readTime;
@@ -1470,6 +1483,7 @@ public void onNext(RunQueryResponse response) {
  QueryDocumentSnapshot.fromDocument(
  rpcContext, Timestamp.fromProto(response.getReadTime()), document);
  documentObserver.onNext(documentSnapshot);
+ lastReceivedDocument.set(documentSnapshot);
  }
 
  if (readTime == null) {
@@ -1479,8 +1493,27 @@ public void onNext(RunQueryResponse response) {
 
  @Override
  public void onError(Throwable throwable) {
- Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Error");
- documentObserver.onError(throwable);
+ // If a non-transactional query failed, attempt to restart.
+ // Transactional queries are retried via the transaction runner.
+ if (transactionId == null && isRetryableError(throwable)) {
+ Tracing.getTracer()
+ .getCurrentSpan()
+ .addAnnotation("Firestore.Query: Retryable Error");
+
+ // Restart the query but use the last document we received as the
+ // query cursor. Note that this it is ok to not use backoff here
+ // since we are requiring at least a single document result.
+ QueryDocumentSnapshot cursor = lastReceivedDocument.get();
+ if (cursor != null) {
+ Query.this
+ .startAfter(cursor)
+ .internalStream(
+ documentObserver, /* transactionId= */ null, cursor.getReadTime());
+ }
+ } else {
+ Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Error");
+ documentObserver.onError(throwable);
+ }
  }
 
  @Override
@@ -1562,7 +1595,8 @@ public void onCompleted() {
  result.set(querySnapshot);
  }
  },
- transactionId);
+ transactionId,
+ /* readTime= */ null);
 
  return result;
  }
@@ -1624,6 +1658,22 @@ private <T> ImmutableList<T> append(ImmutableList<T> existingList, T newElement)
  return builder.build();
  }
 
+ /** Verifies whether the given exception is retryable based on the RunQuery configuration. */
+ private boolean isRetryableError(Throwable throwable) {
+ if (!(throwable instanceof FirestoreException)) {
+ return false;
+ }
+ Set<StatusCode.Code> codes =
+ FirestoreSettings.newBuilder().runQuerySettings().getRetryableCodes();
+ Status status = ((FirestoreException) throwable).getStatus();
+ for (StatusCode.Code code : codes) {
+ if (code.equals(StatusCode.Code.valueOf(status.getCode().name()))) {
+ return true;
+ }
+ }
+ return false;
+ }
+
  /**
  * Returns true if this Query is equal to the provided object.
  *

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

@@ -347,7 +347,10 @@ public void notFound() throws Exception {
  getDocumentResponse.setReadTime(
  com.google.protobuf.Timestamp.newBuilder().setSeconds(5).setNanos(6));
 
- doAnswer(streamingResponse(getDocumentResponse.build()))
+ doAnswer(
+ streamingResponse(
+ new BatchGetDocumentsResponse[] {getDocumentResponse.build()},
+ /* throwable= */ null))
  .when(firestoreMock)
  .streamRequest(
  getAllCapture.capture(),

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

@@ -279,7 +279,7 @@ public static Answer<BatchGetDocumentsResponse> getAllResponse(
  responses[i] = response.build();
  }
 
- return streamingResponse(responses);
+ return streamingResponse(responses, null);
  }
 
  public static ApiFuture<Empty> rollbackResponse() {
@@ -291,6 +291,12 @@ public static Answer<RunQueryResponse> queryResponse() {
  }
 
  public static Answer<RunQueryResponse> queryResponse(String... documentNames) {
+ return queryResponse(/* throwable= */ null, documentNames);
+ }
+
+ /** Returns a stream of documents followed by an optional exception. */
+ public static Answer<RunQueryResponse> queryResponse(
+ @Nullable Throwable throwable, String... documentNames) {
  RunQueryResponse[] responses = new RunQueryResponse[documentNames.length];
 
  for (int i = 0; i < documentNames.length; ++i) {
@@ -301,17 +307,23 @@ public static Answer<RunQueryResponse> queryResponse(String... documentNames) {
  com.google.protobuf.Timestamp.newBuilder().setSeconds(1).setNanos(2));
  responses[i] = runQueryResponse.build();
  }
- return streamingResponse(responses);
+
+ return streamingResponse(responses, throwable);
  }
 
- public static <T> Answer<T> streamingResponse(final T... response) {
+ /** Returns a stream of responses followed by an optional exception. */
+ public static <T> Answer<T> streamingResponse(
+ final T[] response, @Nullable final Throwable throwable) {
  return new Answer<T>() {
  public T answer(InvocationOnMock invocation) {
  Object[] args = invocation.getArguments();
  ApiStreamObserver<T> observer = (ApiStreamObserver<T>) args[1];
  for (T resp : response) {
  observer.onNext(resp);
  }
+ if (throwable != null) {
+ observer.onError(throwable);
+ }
  observer.onCompleted();
  return null;
  }

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

@@ -46,17 +46,20 @@
 import com.google.common.io.BaseEncoding;
 import com.google.firestore.v1.ArrayValue;
 import com.google.firestore.v1.RunQueryRequest;
+import com.google.firestore.v1.RunQueryResponse;
 import com.google.firestore.v1.StructuredQuery;
 import com.google.firestore.v1.StructuredQuery.Direction;
 import com.google.firestore.v1.StructuredQuery.FieldFilter.Operator;
 import com.google.firestore.v1.Value;
 import com.google.protobuf.InvalidProtocolBufferException;
+import io.grpc.Status;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
 import java.util.concurrent.Semaphore;
 import org.junit.Before;
 import org.junit.Test;
@@ -66,7 +69,9 @@
 import org.mockito.Matchers;
 import org.mockito.Mockito;
 import org.mockito.Spy;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.stubbing.Answer;
 
 @RunWith(MockitoJUnitRunner.class)
 public class QueryTest {
@@ -902,6 +907,113 @@ public void onCompleted() {
  semaphore.acquire();
  }
 
+ @Test
+ public void retriesAfterRetryableError() throws Exception {
+ final boolean[] returnError = new boolean[] {true};
+
+ doAnswer(
+ new Answer<RunQueryResponse>() {
+ public RunQueryResponse answer(InvocationOnMock invocation) throws Throwable {
+ if (returnError[0]) {
+ returnError[0] = false;
+ return queryResponse(
+ FirestoreException.serverRejected(
+ Status.DEADLINE_EXCEEDED, "Simulated test failure"),
+ DOCUMENT_NAME + "1",
+ DOCUMENT_NAME + "2")
+ .answer(invocation);
+ } else {
+ return queryResponse(DOCUMENT_NAME + "3").answer(invocation);
+ }
+ }
+ })
+ .when(firestoreMock)
+ .streamRequest(
+ runQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ // Verify the responses
+ final Semaphore semaphore = new Semaphore(0);
+ final Iterator<String> iterator = Arrays.asList("doc1", "doc2", "doc3").iterator();
+
+ query.stream(
+ new ApiStreamObserver<DocumentSnapshot>() {
+ @Override
+ public void onNext(DocumentSnapshot documentSnapshot) {
+ assertEquals(iterator.next(), documentSnapshot.getId());
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ fail();
+ }
+
+ @Override
+ public void onCompleted() {
+ semaphore.release();
+ }
+ });
+
+ semaphore.acquire();
+
+ // Verify the requests
+ List<RunQueryRequest> requests = runQuery.getAllValues();
+ assertEquals(2, requests.size());
+
+ assertFalse(requests.get(0).hasReadTime());
+ assertFalse(requests.get(0).getStructuredQuery().hasStartAt());
+
+ assertEquals(
+ com.google.protobuf.Timestamp.newBuilder().setSeconds(1).setNanos(2).build(),
+ requests.get(1).getReadTime());
+ assertFalse(requests.get(1).getStructuredQuery().getStartAt().getBefore());
+ assertEquals(
+ DOCUMENT_NAME + "2",
+ requests.get(1).getStructuredQuery().getStartAt().getValues(0).getReferenceValue());
+ }
+
+ @Test
+ public void doesNotRetryAfterNonRetryableError() throws Exception {
+ doAnswer(
+ queryResponse(
+ FirestoreException.serverRejected(
+ Status.PERMISSION_DENIED, "Simulated test failure"),
+ DOCUMENT_NAME + "1",
+ DOCUMENT_NAME + "2"))
+ .when(firestoreMock)
+ .streamRequest(
+ runQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ // Verify the responses
+ final Semaphore semaphore = new Semaphore(0);
+ final Iterator<String> iterator = Arrays.asList("doc1", "doc2").iterator();
+
+ query.stream(
+ new ApiStreamObserver<DocumentSnapshot>() {
+ @Override
+ public void onNext(DocumentSnapshot documentSnapshot) {
+ assertEquals(iterator.next(), documentSnapshot.getId());
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ semaphore.release();
+ }
+
+ @Override
+ public void onCompleted() {}
+ });
+
+ semaphore.acquire();
+
+ // Verify the request count
+ List<RunQueryRequest> requests = runQuery.getAllValues();
+ assertEquals(1, runQuery.getAllValues().size());
+ }
+
  @Test
  public void equalsTest() {
  assertEquals(query.limit(42).offset(1337), query.offset(1337).limit(42));