From ba7c5348363ab6c33e1cee3c03c0be68a46ca07c Mon Sep 17 00:00:00 2001
From: Noah Dietz <noahdietz@users.noreply.github.com>
Date: Thu, 5 May 2022 09:49:25 -0700
Subject: [PATCH] feat(v2): add OnHTTPCodes CallOption (#188)

---
 v2/call_option.go      | 37 +++++++++++++++++++++++++++++++++++++
 v2/call_option_test.go | 21 +++++++++++++++++++++
 v2/example_test.go     | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/v2/call_option.go b/v2/call_option.go
index 3011ba4..e092005 100644
--- a/v2/call_option.go
+++ b/v2/call_option.go
@@ -30,9 +30,11 @@
 package gax
 
 import (
+	"errors"
 	"math/rand"
 	"time"
 
+	"google.golang.org/api/googleapi"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -119,6 +121,41 @@ func (r *boRetryer) Retry(err error) (time.Duration, bool) {
 	return 0, false
 }
 
+// OnHTTPCodes returns a Retryer that retries if and only if
+// the previous attempt returns a googleapi.Error whose status code is stored in
+// cc. Pause times between retries are specified by bo.
+//
+// bo is only used for its parameters; each Retryer has its own copy.
+func OnHTTPCodes(bo Backoff, cc ...int) Retryer {
+	codes := make(map[int]bool, len(cc))
+	for _, c := range cc {
+		codes[c] = true
+	}
+
+	return &httpRetryer{
+		backoff: bo,
+		codes:   codes,
+	}
+}
+
+type httpRetryer struct {
+	backoff Backoff
+	codes   map[int]bool
+}
+
+func (r *httpRetryer) Retry(err error) (time.Duration, bool) {
+	var gerr *googleapi.Error
+	if !errors.As(err, &gerr) {
+		return 0, false
+	}
+
+	if r.codes[gerr.Code] {
+		return r.backoff.Pause(), true
+	}
+
+	return 0, false
+}
+
 // Backoff implements exponential backoff. The wait time between retries is a
 // random value between 0 and the "retry period" - the time between retries. The
 // retry period starts at Initial and increases by the factor of Multiplier
diff --git a/v2/call_option_test.go b/v2/call_option_test.go
index 6a1a3da..03a8609 100644
--- a/v2/call_option_test.go
+++ b/v2/call_option_test.go
@@ -31,9 +31,11 @@ package gax
 
 import (
 	"context"
+	"net/http"
 	"testing"
 	"time"
 
+	"google.golang.org/api/googleapi"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -108,3 +110,22 @@ func TestOnErrorFunc(t *testing.T) {
 		}
 	}
 }
+
+func TestOnHTTPCodes(t *testing.T) {
+	apiErr := &googleapi.Error{Code: http.StatusBadGateway}
+	tests := []struct {
+		c     []int
+		retry bool
+	}{
+		{nil, false},
+		{[]int{http.StatusConflict}, false},
+		{[]int{http.StatusConflict, http.StatusBadGateway}, true},
+		{[]int{http.StatusBadGateway}, true},
+	}
+	for _, tst := range tests {
+		b := OnHTTPCodes(Backoff{}, tst.c...)
+		if _, retry := b.Retry(apiErr); retry != tst.retry {
+			t.Errorf("retriable codes: %v, error: %s, retry: %t, want %t", tst.c, apiErr, retry, tst.retry)
+		}
+	}
+}
diff --git a/v2/example_test.go b/v2/example_test.go
index 8cb071d..2fa47dd 100644
--- a/v2/example_test.go
+++ b/v2/example_test.go
@@ -145,6 +145,47 @@ func ExampleOnCodes() {
 	_ = resp // TODO: use resp if err is nil
 }
 
+func ExampleOnHTTPCodes() {
+	ctx := context.Background()
+	c := &fakeClient{}
+
+	retryer := gax.OnHTTPCodes(gax.Backoff{
+		Initial:    time.Second,
+		Max:        32 * time.Second,
+		Multiplier: 2,
+	}, http.StatusBadGateway, http.StatusServiceUnavailable)
+
+	performSomeRPCWithRetry := func(ctx context.Context) (*fakeResponse, error) {
+		for {
+			resp, err := c.PerformSomeRPC(ctx)
+			if err != nil {
+				if delay, shouldRetry := retryer.Retry(err); shouldRetry {
+					if err := gax.Sleep(ctx, delay); err != nil {
+						return nil, err
+					}
+					continue
+				}
+				return nil, err
+			}
+			return resp, err
+		}
+	}
+
+	// It's recommended to set deadlines on RPCs and around retrying. This is
+	// also usually preferred over setting some fixed number of retries: one
+	// advantage this has is that backoff settings can be changed independently
+	// of the deadline, whereas with a fixed number of retries the deadline
+	// would be a constantly-shifting goalpost.
+	ctxWithTimeout, cancel := context.WithDeadline(ctx, time.Now().Add(5*time.Minute))
+	defer cancel()
+
+	resp, err := performSomeRPCWithRetry(ctxWithTimeout)
+	if err != nil {
+		// TODO: handle err
+	}
+	_ = resp // TODO: use resp if err is nil
+}
+
 func ExampleBackoff() {
 	ctx := context.Background()