fix(bigquery): wrap default ResultRetryAlgorithm to retry transient HTTP errors

Author: jinseopkim0

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigquery;
 
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.core.ApiClock;
 import com.google.api.gax.retrying.DirectRetryingExecutor;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
@@ -23,6 +24,7 @@
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.retrying.RetryingExecutor;
 import com.google.api.gax.retrying.RetryingFuture;
+import com.google.api.gax.retrying.TimedAttemptSettings;
 import com.google.api.gax.retrying.TimedRetryAlgorithm;
 import com.google.cloud.RetryHelper;
 import io.opentelemetry.api.trace.Span;
@@ -69,6 +71,9 @@ public static <V> V runWithRetries(
       // implementation does not use response at all, so ignoring its type is ok.
       @SuppressWarnings("unchecked")
       ResultRetryAlgorithm<V> algorithm = (ResultRetryAlgorithm<V>) resultRetryAlgorithm;
+      if (algorithm == BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER) {
+        algorithm = wrapDefaultAlgorithm(algorithm);
+      }
       return run(
           callable,
           new ExponentialRetryAlgorithm(retrySettings, clock),
@@ -119,6 +124,29 @@ private static <V> V run(
     return retryingFuture.get();
   }
 
+  private static <V> ResultRetryAlgorithm<V> wrapDefaultAlgorithm(
+      ResultRetryAlgorithm<V> defaultAlgorithm) {
+    return new ResultRetryAlgorithm<V>() {
+      @Override
+      public TimedAttemptSettings createNextAttempt(
+          Throwable previousThrowable, V previousResponse, TimedAttemptSettings previousSettings) {
+        return defaultAlgorithm.createNextAttempt(
+            previousThrowable, previousResponse, previousSettings);
+      }
+
+      @Override
+      public boolean shouldRetry(Throwable previousThrowable, V previousResponse) {
+        if (previousThrowable instanceof HttpResponseException) {
+          int statusCode = ((HttpResponseException) previousThrowable).getStatusCode();
+          if (statusCode == 500 || statusCode == 502 || statusCode == 503 || statusCode == 504) {
+            return true;
+          }
+        }
+        return defaultAlgorithm.shouldRetry(previousThrowable, previousResponse);
+      }
+    };
+  }
+
   public static class BigQueryRetryHelperException extends RuntimeException {
 
     private static final long serialVersionUID = -8519852520090965314L;

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -36,6 +36,10 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.api.client.googleapis.json.GoogleJsonError;
+import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpResponseException;
 import com.google.api.gax.paging.Page;
 import com.google.api.services.bigquery.model.ErrorProto;
 import com.google.api.services.bigquery.model.GetQueryResultsResponse;
@@ -935,6 +939,37 @@ void testGetTable() throws IOException {
         .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
   }
 
+  @Test
+  void testGetTableFailureShouldRetryServerErrors() throws IOException {
+    GoogleJsonError error = new GoogleJsonError();
+    error.setMessage("Visibility check was unavailable. Please retry the request");
+    error.setCode(503);
+    GoogleJsonError.ErrorInfo errorInfo = new GoogleJsonError.ErrorInfo();
+    errorInfo.setReason("backendError");
+    error.setErrors(ImmutableList.of(errorInfo));
+
+    when(bigqueryRpcMock.getTableSkipExceptionTranslation(
+            PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
+        .thenThrow(new GoogleJsonResponseException(serverErrorResponse(), error))
+        .thenReturn(TABLE_INFO_WITH_PROJECT.toPb());
+
+    bigquery =
+        options.toBuilder()
+            .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
+            .build()
+            .getService();
+
+    Table table = bigquery.getTable(DATASET, TABLE);
+
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
+    verify(bigqueryRpcMock, times(2))
+        .getTableSkipExceptionTranslation(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS);
+  }
+
+  private static HttpResponseException.Builder serverErrorResponse() {
+    return new HttpResponseException.Builder(503, "Service Unavailable", new HttpHeaders());
+  }
+
   @Test
   void testGetModel() throws IOException {
     when(bigqueryRpcMock.getModelSkipExceptionTranslation(