From 753b75139f4b9e8593db5d45d8ab1e0cc8969350 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Thu, 3 Nov 2022 15:26:39 -0700
Subject: [PATCH] feat(bigquery): widen retry predicate (#6976)

This PR adds 500,504 http response codes for the default retry predicate on unary retries.

This doesn't introduce job-level retries (jobs must be recreated wholly, they can't be effectively restarted), so the primary risk of this change is terminal job state propagating into the HTTP response code of one of the polling methods (e.g. job.getQueryResults, jobs.get, etc).  In this situation, job execution may appear to hang indefinitely or until context cancellation/expiry.

Related: https://togithub.com/googleapis/google-cloud-go/issues/5248
---
 bigquery/bigquery.go      | 18 ++++++++++++------
 bigquery/bigquery_test.go |  2 +-
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/bigquery/bigquery.go b/bigquery/bigquery.go
index 9b4735722231..8bb92ba683ca 100644
--- a/bigquery/bigquery.go
+++ b/bigquery/bigquery.go
@@ -203,12 +203,16 @@ func runWithRetryExplicit(ctx context.Context, call func() error, allowedReasons
 var (
 	defaultRetryReasons = []string{"backendError", "rateLimitExceeded"}
 	jobRetryReasons     = []string{"backendError", "rateLimitExceeded", "internalError"}
+	retry5xxCodes       = []int{
+		http.StatusInternalServerError,
+		http.StatusBadGateway,
+		http.StatusServiceUnavailable,
+		http.StatusGatewayTimeout,
+	}
 )
 
-// This is the correct definition of retryable according to the BigQuery team. It
-// also considers 502 ("Bad Gateway") and 503 ("Service Unavailable") errors
-// retryable; these are returned by systems between the client and the BigQuery
-// service.
+// retryableError is the unary retry predicate for this library.  In addition to structured error
+// reasons, it specifies some HTTP codes (500, 502, 503, 504) and network/transport reasons.
 func retryableError(err error, allowedReasons []string) bool {
 	if err == nil {
 		return false
@@ -237,8 +241,10 @@ func retryableError(err error, allowedReasons []string) bool {
 				}
 			}
 		}
-		if e.Code == http.StatusServiceUnavailable || e.Code == http.StatusBadGateway {
-			return true
+		for _, code := range retry5xxCodes {
+			if e.Code == code {
+				return true
+			}
 		}
 	case *url.Error:
 		retryable := []string{"connection refused", "connection reset"}
diff --git a/bigquery/bigquery_test.go b/bigquery/bigquery_test.go
index c5207c0ab180..f1fd24745ac3 100644
--- a/bigquery/bigquery_test.go
+++ b/bigquery/bigquery_test.go
@@ -86,7 +86,7 @@ func TestRetryableErrors(t *testing.T) {
 			&googleapi.Error{
 				Code: http.StatusInternalServerError,
 			},
-			false,
+			true,
 		},
 		{
 			"internal w/backend reason",