fix: add retry logic for readrows v1beta1

google-cloud-bigquerystorage/pom.xml

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

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/stub/EnhancedBigQueryStorageStub.java

@@ -17,9 +17,20 @@
 
 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.v1beta1.BigQueryStorageGrpc;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.BatchCreateReadSessionStreamsRequest;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.BatchCreateReadSessionStreamsResponse;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.CreateReadSessionRequest;
@@ -29,8 +40,12 @@
 import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.SplitReadStreamRequest;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.SplitReadStreamResponse;
+import com.google.cloud.bigquery.storage.v1beta1.stub.readrows.ApiResultRetryAlgorithm;
+import com.google.cloud.bigquery.storage.v1beta1.stub.readrows.ReadRowsRetryingCallable;
+import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Empty;
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -40,7 +55,10 @@
  */
 public class EnhancedBigQueryStorageStub implements BackgroundResource {
 
+  private static final String TRACING_OUTER_CLIENT_NAME = "BigQueryStorage";
   private final GrpcBigQueryStorageStub stub;
+  private final BigQueryStorageStubSettings stubSettings;
+  private final ClientContext context;
 
   public static EnhancedBigQueryStorageStub create(EnhancedBigQueryStorageStubSettings settings)
       throws IOException {
@@ -89,20 +107,67 @@ public static EnhancedBigQueryStorageStub create(EnhancedBigQueryStorageStubSett
     BigQueryStorageStubSettings baseSettings = baseSettingsBuilder.build();
     ClientContext clientContext = ClientContext.create(baseSettings);
     GrpcBigQueryStorageStub stub = new GrpcBigQueryStorageStub(baseSettings, clientContext);
-    return new EnhancedBigQueryStorageStub(stub);
+    return new EnhancedBigQueryStorageStub(stub, baseSettings, clientContext);
   }
 
   @InternalApi("Visible for testing")
-  EnhancedBigQueryStorageStub(GrpcBigQueryStorageStub stub) {
+  EnhancedBigQueryStorageStub(
+      GrpcBigQueryStorageStub stub,
+      BigQueryStorageStubSettings 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(BigQueryStorageGrpc.getReadRowsMethod())
+                .setParamsExtractor(
+                    new RequestParamsExtractor<ReadRowsRequest>() {
+                      @Override
+                      public Map<String, String> extract(ReadRowsRequest request) {
+                        return ImmutableMap.of(
+                            "read_position.stream.name",
+                            String.valueOf(request.getReadPosition().getStream().getName()));
+                      }
+                    })
+                .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<BatchCreateReadSessionStreamsRequest, BatchCreateReadSessionStreamsResponse>

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta1/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.v1beta1.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();
+  }
+}