Enable custom RetryOptions to be specified per API call

sdk/azcore/policy_retry.go

@@ -49,9 +49,9 @@ type RetryOptions struct {
 
 var (
 	// StatusCodesForRetry is the default set of HTTP status code for which the policy will retry.
-	StatusCodesForRetry = [6]int{
+	// Changing the value of StatusCodesForRetry will affect all clients that use the default values.
+	StatusCodesForRetry = []int{
 		http.StatusRequestTimeout,      // 408
-		http.StatusTooManyRequests,     // 429
 		http.StatusInternalServerError, // 500
 		http.StatusBadGateway,          // 502
 		http.StatusServiceUnavailable,  // 503
@@ -62,14 +62,23 @@ var (
 // DefaultRetryOptions returns an instance of RetryOptions initialized with default values.
 func DefaultRetryOptions() RetryOptions {
 	return RetryOptions{
-		StatusCodes:   StatusCodesForRetry[:],
+		StatusCodes:   StatusCodesForRetry,
 		MaxTries:      defaultMaxTries,
 		TryTimeout:    1 * time.Minute,
 		RetryDelay:    4 * time.Second,
 		MaxRetryDelay: 120 * time.Second,
 	}
 }
 
+// used as a context key for adding/retrieving RetryOptions
+type ctxWithRetryOptionsKey struct{}
+
+// WithRetryOptions adds the specified RetryOptions to the parent context.
+// Use this to specify custom RetryOptions at the API-call level.
+func WithRetryOptions(parent context.Context, options RetryOptions) context.Context {
+	return context.WithValue(parent, ctxWithRetryOptionsKey{}, options)
+}
+
 func (o RetryOptions) calcDelay(try int32) time.Duration { // try is >=1; never 0
 	pow := func(number int64, exponent int32) int64 { // pow is nested helper function
 		var result int64 = 1
@@ -105,6 +114,11 @@ type retryPolicy struct {
 }
 
 func (p *retryPolicy) Do(ctx context.Context, req *Request) (resp *Response, err error) {
+	options := p.options
+	// check if the retry options have been overridden for this call
+	if override := ctx.Value(ctxWithRetryOptionsKey{}); override != nil {
+		options = override.(RetryOptions)
+	}
 	// Exponential retry algorithm: ((2 ^ attempt) - 1) * delay * random(0.8, 1.2)
 	// When to retry: connection failure or temporary/timeout.
 	if req.Body != nil {
@@ -134,14 +148,14 @@ func (p *retryPolicy) Do(ctx context.Context, req *Request) (resp *Response, err
 		}
 
 		// Set the time for this particular retry operation and then Do the operation.
-		tryCtx, tryCancel := context.WithTimeout(ctx, p.options.TryTimeout)
+		tryCtx, tryCancel := context.WithTimeout(ctx, options.TryTimeout)
 		resp, err = req.Next(tryCtx) // Make the request
 		tryCancel()
 		if shouldLog {
 			Log().Write(LogRetryPolicy, fmt.Sprintf("Err=%v, response=%v\n", err, resp))
 		}
 
-		if err == nil && !resp.HasStatusCode(p.options.StatusCodes...) {
+		if err == nil && !resp.HasStatusCode(options.StatusCodes...) {
 			// if there is no error and the response code isn't in the list of retry codes then we're done.
 			return
 		} else if ctx.Err() != nil {
@@ -155,15 +169,15 @@ func (p *retryPolicy) Do(ctx context.Context, req *Request) (resp *Response, err
 		// drain before retrying so nothing is leaked
 		resp.Drain()
 
-		if try == p.options.MaxTries {
+		if try == options.MaxTries {
 			// max number of tries has been reached, don't sleep again
 			return
 		}
 
 		// use the delay from retry-after if available
 		delay, ok := resp.RetryAfter()
 		if !ok {
-			delay = p.options.calcDelay(try)
+			delay = options.calcDelay(try)
 		}
 		if shouldLog {
 			Log().Write(LogRetryPolicy, fmt.Sprintf("Try=%d, Delay=%v\n", try, delay))

sdk/azcore/policy_retry_test.go

@@ -212,6 +212,35 @@ func TestRetryPolicyIsNotRetriable(t *testing.T) {
 	}
 }
 
+func TestWithRetryOptions(t *testing.T) {
+	srv, close := mock.NewServer()
+	defer close()
+	srv.RepeatResponse(9, mock.WithStatusCode(http.StatusRequestTimeout))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK))
+	defaultOptions := testRetryOptions()
+	pl := NewPipeline(srv, NewRetryPolicy(defaultOptions))
+	customOptions := *defaultOptions
+	customOptions.MaxTries = 10
+	customOptions.MaxRetryDelay = 200 * time.Millisecond
+	retryCtx := WithRetryOptions(context.Background(), customOptions)
+	req := NewRequest(http.MethodGet, srv.URL())
+	body := newRewindTrackingBody("stuff")
+	req.SetBody(body)
+	resp, err := pl.Do(retryCtx, req)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("unexpected status code: %d", resp.StatusCode)
+	}
+	if body.rcount != int(customOptions.MaxTries-1) {
+		t.Fatalf("unexpected rewind count: %d", body.rcount)
+	}
+	if !body.closed {
+		t.Fatal("request body wasn't closed")
+	}
+}
+
 // TODO: add test for retry failing to read response body
 
 // TODO: add test for per-retry timeout failed but e2e succeeded