fix: add retry logic for readrows v1beta2 (#315)

google-cloud-bigquerystorage/pom.xml

@@ -174,7 +174,6 @@
     <dependency>
       <groupId>com.google.api.grpc</groupId>
       <artifactId>grpc-google-cloud-bigquerystorage-v1beta2</artifactId>
-      <scope>test</scope>
     </dependency>
     <dependency>
       <groupId>com.google.api.grpc</groupId>

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStub.java

@@ -17,16 +17,31 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.grpc.GrpcCallSettings;
+import com.google.api.gax.grpc.GrpcRawCallableFactory;
+import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.ScheduledRetryingExecutor;
+import com.google.api.gax.retrying.StreamingRetryAlgorithm;
+import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.RequestParamsExtractor;
+import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.SpanName;
+import com.google.api.gax.tracing.TracedServerStreamingCallable;
+import com.google.cloud.bigquery.storage.v1beta2.BigQueryReadGrpc;
 import com.google.cloud.bigquery.storage.v1beta2.CreateReadSessionRequest;
 import com.google.cloud.bigquery.storage.v1beta2.ReadRowsRequest;
 import com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse;
 import com.google.cloud.bigquery.storage.v1beta2.ReadSession;
 import com.google.cloud.bigquery.storage.v1beta2.SplitReadStreamRequest;
 import com.google.cloud.bigquery.storage.v1beta2.SplitReadStreamResponse;
+import com.google.cloud.bigquery.storage.v1beta2.stub.readrows.ApiResultRetryAlgorithm;
+import com.google.cloud.bigquery.storage.v1beta2.stub.readrows.ReadRowsRetryingCallable;
+import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -35,7 +50,11 @@
  * <p>This class is for advanced usage and reflects the underlying API directly.
  */
 public class EnhancedBigQueryReadStub implements BackgroundResource {
+
+  private static final String TRACING_OUTER_CLIENT_NAME = "BigQueryStorage";
   private final GrpcBigQueryReadStub stub;
+  private final BigQueryReadStubSettings stubSettings;
+  private final ClientContext context;
 
   public static EnhancedBigQueryReadStub create(EnhancedBigQueryReadStubSettings settings)
       throws IOException {
@@ -69,20 +88,64 @@ public static EnhancedBigQueryReadStub create(EnhancedBigQueryReadStubSettings s
     BigQueryReadStubSettings baseSettings = baseSettingsBuilder.build();
     ClientContext clientContext = ClientContext.create(baseSettings);
     GrpcBigQueryReadStub stub = new GrpcBigQueryReadStub(baseSettings, clientContext);
-    return new EnhancedBigQueryReadStub(stub);
+    return new EnhancedBigQueryReadStub(stub, baseSettings, clientContext);
   }
 
   @InternalApi("Visible for testing")
-  EnhancedBigQueryReadStub(GrpcBigQueryReadStub stub) {
+  EnhancedBigQueryReadStub(
+      GrpcBigQueryReadStub stub, BigQueryReadStubSettings stubSettings, ClientContext context) {
     this.stub = stub;
+    this.stubSettings = stubSettings;
+    this.context = context;
   }
 
   public UnaryCallable<CreateReadSessionRequest, ReadSession> createReadSessionCallable() {
     return stub.createReadSessionCallable();
   }
 
   public ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> readRowsCallable() {
-    return stub.readRowsCallable();
+    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> innerCallable =
+        GrpcRawCallableFactory.createServerStreamingCallable(
+            GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder()
+                .setMethodDescriptor(BigQueryReadGrpc.getReadRowsMethod())
+                .setParamsExtractor(
+                    new RequestParamsExtractor<ReadRowsRequest>() {
+                      @Override
+                      public Map<String, String> extract(ReadRowsRequest request) {
+                        return ImmutableMap.of(
+                            "read_stream", String.valueOf(request.getReadStream()));
+                      }
+                    })
+                .build(),
+            stubSettings.readRowsSettings().getRetryableCodes());
+    ServerStreamingCallSettings<ReadRowsRequest, ReadRowsResponse> callSettings =
+        stubSettings.readRowsSettings();
+
+    StreamingRetryAlgorithm<Void> retryAlgorithm =
+        new StreamingRetryAlgorithm<>(
+            new ApiResultRetryAlgorithm<Void>(),
+            new ExponentialRetryAlgorithm(callSettings.getRetrySettings(), context.getClock()));
+
+    ScheduledRetryingExecutor<Void> retryingExecutor =
+        new ScheduledRetryingExecutor<>(retryAlgorithm, context.getExecutor());
+
+    if (context.getStreamWatchdog() != null) {
+      innerCallable = Callables.watched(innerCallable, callSettings, context);
+    }
+
+    ReadRowsRetryingCallable outerCallable =
+        new ReadRowsRetryingCallable(
+            context.getDefaultCallContext(),
+            innerCallable,
+            retryingExecutor,
+            callSettings.getResumptionStrategy());
+
+    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> traced =
+        new TracedServerStreamingCallable<>(
+            outerCallable,
+            context.getTracerFactory(),
+            SpanName.of(TRACING_OUTER_CLIENT_NAME, "ReadRows"));
+    return traced.withDefaultCallContext(context.getDefaultCallContext());
   }
 
   public UnaryCallable<SplitReadStreamRequest, SplitReadStreamResponse> splitReadStreamCallable() {

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/readrows/ApiResultRetryAlgorithm.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery.storage.v1beta2.stub.readrows;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
+import com.google.api.gax.rpc.ApiException;
+import io.grpc.Status;
+import org.threeten.bp.Duration;
+
+/** For internal use, public for technical reasons. */
+@InternalApi
+public class ApiResultRetryAlgorithm<ResponseT> implements ResultRetryAlgorithm<ResponseT> {
+  // Duration to sleep on if the error is DEADLINE_EXCEEDED.
+  public static final Duration DEADLINE_SLEEP_DURATION = Duration.ofMillis(1);
+
+  @Override
+  public TimedAttemptSettings createNextAttempt(
+      Throwable prevThrowable, ResponseT prevResponse, TimedAttemptSettings prevSettings) {
+    if (prevThrowable != null) {
+      Status status = Status.fromThrowable(prevThrowable);
+      if (status.getCode() == Status.Code.INTERNAL
+          && status.getDescription() != null
+          && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
+        return TimedAttemptSettings.newBuilder()
+            .setGlobalSettings(prevSettings.getGlobalSettings())
+            .setRetryDelay(prevSettings.getRetryDelay())
+            .setRpcTimeout(prevSettings.getRpcTimeout())
+            .setRandomizedRetryDelay(DEADLINE_SLEEP_DURATION)
+            .setAttemptCount(prevSettings.getAttemptCount() + 1)
+            .setFirstAttemptStartTimeNanos(prevSettings.getFirstAttemptStartTimeNanos())
+            .build();
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public boolean shouldRetry(Throwable prevThrowable, ResponseT prevResponse) {
+    if (prevThrowable != null) {
+      Status status = Status.fromThrowable(prevThrowable);
+      if (status.getCode() == Status.Code.INTERNAL
+          && status.getDescription() != null
+          && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
+        return true;
+      }
+    }
+    return (prevThrowable instanceof ApiException) && ((ApiException) prevThrowable).isRetryable();
+  }
+}