Add more retry after headers to check during retries

sdk/azcore/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+* The `retry-after-ms` and `x-ms-retry-after-ms` headers weren't being checked during retries.
+
 ### Other Changes
 
 ## 1.9.0 (2023-11-06)

sdk/azcore/internal/shared/constants.go

@@ -22,11 +22,13 @@ const (
 	HeaderLocation               = "Location"
 	HeaderOperationLocation      = "Operation-Location"
 	HeaderRetryAfter             = "Retry-After"
+	HeaderRetryAfterMS           = "Retry-After-Ms"
 	HeaderUserAgent              = "User-Agent"
 	HeaderWWWAuthenticate        = "WWW-Authenticate"
 	HeaderXMSClientRequestID     = "x-ms-client-request-id"
 	HeaderXMSRequestID           = "x-ms-request-id"
 	HeaderXMSErrorCode           = "x-ms-error-code"
+	HeaderXMSRetryAfterMS        = "x-ms-retry-after-ms"
 )
 
 const BearerTokenPrefix = "Bearer "

sdk/azcore/internal/shared/shared.go

@@ -44,22 +44,64 @@ func Delay(ctx context.Context, delay time.Duration) error {
 	}
 }
 
-// RetryAfter returns non-zero if the response contains a Retry-After header value.
+// RetryAfter returns non-zero if the response contains one of the headers with a "retry after" value.
+// Headers are checked in the following order: retry-after-ms, x-ms-retry-after-ms, retry-after
 func RetryAfter(resp *http.Response) time.Duration {
 	if resp == nil {
 		return 0
 	}
-	ra := resp.Header.Get(HeaderRetryAfter)
-	if ra == "" {
-		return 0
+
+	type retryData struct {
+		header string
+		units  time.Duration
+
+		// custom is used when the regular algorithm failed and is optional.
+		// the returned duration is used verbatim (units is not applied).
+		custom func(string) time.Duration
 	}
-	// retry-after values are expressed in either number of
-	// seconds or an HTTP-date indicating when to try again
-	if retryAfter, _ := strconv.Atoi(ra); retryAfter > 0 {
-		return time.Duration(retryAfter) * time.Second
-	} else if t, err := time.Parse(time.RFC1123, ra); err == nil {
-		return time.Until(t)
+
+	nop := func(string) time.Duration { return 0 }
+
+	// the headers are listed in order of preference
+	retries := []retryData{
+		{
+			header: HeaderRetryAfterMS,
+			units:  time.Millisecond,
+			custom: nop,
+		},
+		{
+			header: HeaderXMSRetryAfterMS,
+			units:  time.Millisecond,
+			custom: nop,
+		},
+		{
+			header: HeaderRetryAfter,
+			units:  time.Second,
+
+			// retry-after values are expressed in either number of
+			// seconds or an HTTP-date indicating when to try again
+			custom: func(ra string) time.Duration {
+				t, err := time.Parse(time.RFC1123, ra)
+				if err != nil {
+					return 0
+				}
+				return time.Until(t)
+			},
+		},
 	}
+
+	for _, retry := range retries {
+		v := resp.Header.Get(retry.header)
+		if v == "" {
+			continue
+		}
+		if retryAfter, _ := strconv.Atoi(v); retryAfter > 0 {
+			return time.Duration(retryAfter) * retry.units
+		} else if d := retry.custom(v); d > 0 {
+			return d
+		}
+	}
+
 	return 0
 }
 

sdk/azcore/internal/shared/shared_test.go

@@ -59,6 +59,28 @@ func TestRetryAfter(t *testing.T) {
 	if d = RetryAfter(resp); d != 0 {
 		t.Fatalf("expected zero for invalid value, got %d", d)
 	}
+	// verify that the ms-granularity headers are preferred
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderRetryAfterMS, "500")
+	require.Equal(t, time.Duration(500)*time.Millisecond, RetryAfter(resp))
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderXMSRetryAfterMS, "400")
+	require.Equal(t, time.Duration(400)*time.Millisecond, RetryAfter(resp))
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderRetryAfterMS, "500")
+	resp.Header.Set(HeaderXMSRetryAfterMS, "400")
+	resp.Header.Set(HeaderRetryAfter, "300")
+	require.Equal(t, time.Duration(500)*time.Millisecond, RetryAfter(resp))
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderXMSRetryAfterMS, "400")
+	resp.Header.Set(HeaderRetryAfter, "300")
+	require.Equal(t, time.Duration(400)*time.Millisecond, RetryAfter(resp))
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderRetryAfterMS, "invalid")
+	require.Zero(t, RetryAfter(resp))
+	resp.Header = http.Header{}
+	resp.Header.Set(HeaderXMSRetryAfterMS, "invalid")
+	require.Zero(t, RetryAfter(resp))
 }
 
 func TestTypeOfT(t *testing.T) {