Support custom SendDecorator chains via context (#417)

* Support custom SendDecorator chains via context

Added `autorest.WithSendDecorators` and `autorest.GetSendDecorators` for
adding and retrieving a custom chain of SendDecorators to the provided
context.
Added `autorest.DoRetryForStatusCodesWithCap` and
`autorest.DelayForBackoffWithCap` to enforce an upper bound on the
duration between retries.
Fixed up some code comments.

* small refactor based on PR feedback

* remove some changes for dev branch

autorest/sender.go

@@ -15,6 +15,7 @@ package autorest
 //  limitations under the License.
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"math"
@@ -25,6 +26,23 @@ import (
 	"github.com/Azure/go-autorest/tracing"
 )
 
+// used as a key type in context.WithValue()
+type ctxSendDecorators struct{}
+
+// WithSendDecorators adds the specified SendDecorators to the provided context.
+func WithSendDecorators(ctx context.Context, sendDecorator []SendDecorator) context.Context {
+	return context.WithValue(ctx, ctxSendDecorators{}, sendDecorator)
+}
+
+// GetSendDecorators returns the SendDecorators in the provided context or the provided default SendDecorators.
+func GetSendDecorators(ctx context.Context, defaultSendDecorators ...SendDecorator) []SendDecorator {
+	inCtx := ctx.Value(ctxSendDecorators{})
+	if sd, ok := inCtx.([]SendDecorator); ok {
+		return sd
+	}
+	return defaultSendDecorators
+}
+
 // Sender is the interface that wraps the Do method to send HTTP requests.
 //
 // The standard http.Client conforms to this interface.
@@ -211,43 +229,59 @@ func DoRetryForAttempts(attempts int, backoff time.Duration) SendDecorator {
 
 // DoRetryForStatusCodes returns a SendDecorator that retries for specified statusCodes for up to the specified
 // number of attempts, exponentially backing off between requests using the supplied backoff
-// time.Duration (which may be zero). Retrying may be canceled by closing the optional channel on
-// the http.Request.
+// time.Duration (which may be zero). Retrying may be canceled by cancelling the context on the http.Request.
+// NOTE: Code http.StatusTooManyRequests (429) will *not* be counted against the number of attempts.
 func DoRetryForStatusCodes(attempts int, backoff time.Duration, codes ...int) SendDecorator {
 	return func(s Sender) Sender {
-		return SenderFunc(func(r *http.Request) (resp *http.Response, err error) {
-			rr := NewRetriableRequest(r)
-			// Increment to add the first call (attempts denotes number of retries)
-			for attempt := 0; attempt < attempts+1; {
-				err = rr.Prepare()
-				if err != nil {
-					return resp, err
-				}
-				resp, err = s.Do(rr.Request())
-				// if the error isn't temporary don't bother retrying
-				if err != nil && !IsTemporaryNetworkError(err) {
-					return nil, err
-				}
-				// we want to retry if err is not nil (e.g. transient network failure).  note that for failed authentication
-				// resp and err will both have a value, so in this case we don't want to retry as it will never succeed.
-				if err == nil && !ResponseHasStatusCode(resp, codes...) || IsTokenRefreshError(err) {
-					return resp, err
-				}
-				delayed := DelayWithRetryAfter(resp, r.Context().Done())
-				if !delayed && !DelayForBackoff(backoff, attempt, r.Context().Done()) {
-					return resp, r.Context().Err()
-				}
-				// don't count a 429 against the number of attempts
-				// so that we continue to retry until it succeeds
-				if resp == nil || resp.StatusCode != http.StatusTooManyRequests {
-					attempt++
-				}
-			}
-			return resp, err
+		return SenderFunc(func(r *http.Request) (*http.Response, error) {
+			return doRetryForStatusCodesImpl(s, r, false, attempts, backoff, 0, codes...)
 		})
 	}
 }
 
+// DoRetryForStatusCodesWithCap returns a SendDecorator that retries for specified statusCodes for up to the
+// specified number of attempts, exponentially backing off between requests using the supplied backoff
+// time.Duration (which may be zero). To cap the maximum possible delay between iterations specify a value greater
+// than zero for cap. Retrying may be canceled by cancelling the context on the http.Request.
+func DoRetryForStatusCodesWithCap(attempts int, backoff, cap time.Duration, codes ...int) SendDecorator {
+	return func(s Sender) Sender {
+		return SenderFunc(func(r *http.Request) (*http.Response, error) {
+			return doRetryForStatusCodesImpl(s, r, true, attempts, backoff, cap, codes...)
+		})
+	}
+}
+
+func doRetryForStatusCodesImpl(s Sender, r *http.Request, count429 bool, attempts int, backoff, cap time.Duration, codes ...int) (resp *http.Response, err error) {
+	rr := NewRetriableRequest(r)
+	// Increment to add the first call (attempts denotes number of retries)
+	for attempt := 0; attempt < attempts+1; {
+		err = rr.Prepare()
+		if err != nil {
+			return
+		}
+		resp, err = s.Do(rr.Request())
+		// if the error isn't temporary don't bother retrying
+		if err != nil && !IsTemporaryNetworkError(err) {
+			return
+		}
+		// we want to retry if err is not nil (e.g. transient network failure).  note that for failed authentication
+		// resp and err will both have a value, so in this case we don't want to retry as it will never succeed.
+		if err == nil && !ResponseHasStatusCode(resp, codes...) || IsTokenRefreshError(err) {
+			return resp, err
+		}
+		delayed := DelayWithRetryAfter(resp, r.Context().Done())
+		if !delayed && !DelayForBackoffWithCap(backoff, cap, attempt, r.Context().Done()) {
+			return resp, r.Context().Err()
+		}
+		// when count429 == false don't count a 429 against the number
+		// of attempts so that we continue to retry until it succeeds
+		if count429 || (resp == nil || resp.StatusCode != http.StatusTooManyRequests) {
+			attempt++
+		}
+	}
+	return resp, err
+}
+
 // DelayWithRetryAfter invokes time.After for the duration specified in the "Retry-After" header.
 // The value of Retry-After can be either the number of seconds or a date in RFC1123 format.
 // The function returns true after successfully waiting for the specified duration.  If there is
@@ -325,8 +359,22 @@ func WithLogging(logger *log.Logger) SendDecorator {
 // Note: Passing attempt 1 will result in doubling "backoff" duration. Treat this as a zero-based attempt
 // count.
 func DelayForBackoff(backoff time.Duration, attempt int, cancel <-chan struct{}) bool {
+	return DelayForBackoffWithCap(backoff, 0, attempt, cancel)
+}
+
+// DelayForBackoffWithCap invokes time.After for the supplied backoff duration raised to the power of
+// passed attempt (i.e., an exponential backoff delay). Backoff duration is in seconds and can set
+// to zero for no delay. To cap the maximum possible delay specify a value greater than zero for cap.
+// The delay may be canceled by closing the passed channel. If terminated early, returns false.
+// Note: Passing attempt 1 will result in doubling "backoff" duration. Treat this as a zero-based attempt
+// count.
+func DelayForBackoffWithCap(backoff, cap time.Duration, attempt int, cancel <-chan struct{}) bool {
+	d := time.Duration(backoff.Seconds()*math.Pow(2, float64(attempt))) * time.Second
+	if cap > 0 && d > cap {
+		d = cap
+	}
 	select {
-	case <-time.After(time.Duration(backoff.Seconds()*math.Pow(2, float64(attempt))) * time.Second):
+	case <-time.After(d):
 		return true
 	case <-cancel:
 		return false

autorest/sender_test.go

@@ -513,6 +513,15 @@ func TestDelayForBackoff(t *testing.T) {
 	}
 }
 
+func TestDelayForBackoffWithCap(t *testing.T) {
+	d := 2 * time.Second
+	start := time.Now()
+	DelayForBackoffWithCap(d, 1*time.Second, 0, nil)
+	if time.Since(start) >= d {
+		t.Fatal("autorest: DelayForBackoffWithCap delayed for too long")
+	}
+}
+
 func TestDelayForBackoff_Cancels(t *testing.T) {
 	cancel := make(chan struct{})
 	delay := 5 * time.Second
@@ -979,3 +988,26 @@ func TestDoRetryForStatusCodes_Race(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestGetSendDecorators(t *testing.T) {
+	sd := GetSendDecorators(context.Background())
+	if l := len(sd); l != 0 {
+		t.Fatalf("expected zero length but got %d", l)
+	}
+	sd = GetSendDecorators(context.Background(), DoCloseIfError(), DoErrorIfStatusCode())
+	if l := len(sd); l != 2 {
+		t.Fatalf("expected length of two but got %d", l)
+	}
+}
+
+func TestWithSendDecorators(t *testing.T) {
+	ctx := WithSendDecorators(context.Background(), []SendDecorator{DoRetryForAttempts(5, 5*time.Second)})
+	sd := GetSendDecorators(ctx)
+	if l := len(sd); l != 1 {
+		t.Fatalf("expected length of one but got %d", l)
+	}
+	sd = GetSendDecorators(ctx, DoCloseIfError(), DoErrorIfStatusCode())
+	if l := len(sd); l != 1 {
+		t.Fatalf("expected length of one but got %d", l)
+	}
+}