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

Author: Praful Makani

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();
+ }
+}