runtime.Poller.Result won't be done on non-terminal error (#22675)

* runtime.Poller.Result won't be done on non-terminal error

If the underlying error when retrieving the final result is non
terminal, don't mark the poller as done so callers can retry.

* expand list of transient HTTP status codes

sdk/azcore/CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Bugs Fixed
 
 * Pollers that use the `Location` header won't consider `http.StatusRequestTimeout` a terminal failure.
+* `runtime.Poller[T].Result` won't consider non-terminal error responses as terminal.
 
 ### Other Changes
 

sdk/azcore/internal/pollers/loc/loc.go

@@ -103,7 +103,7 @@ func (p *Poller[T]) Poll(ctx context.Context) (*http.Response, error) {
 		} else if resp.StatusCode > 199 && resp.StatusCode < 300 {
 			// any 2xx other than a 202 indicates success
 			p.CurState = poller.StatusSucceeded
-		} else if resp.StatusCode == http.StatusRequestTimeout || resp.StatusCode == http.StatusTooManyRequests {
+		} else if pollers.IsNonTerminalHTTPStatusCode(resp) {
 			// the request timed out or is being throttled.
 			// DO NOT include this as a terminal failure. preserve
 			// the existing state and return the response.

sdk/azcore/internal/pollers/util.go

@@ -185,3 +185,16 @@ func ResultHelper[T any](resp *http.Response, failed bool, out *T) error {
 	}
 	return nil
 }
+
+// IsNonTerminalHTTPStatusCode returns true if the HTTP status code should be
+// considered non-terminal thus eligible for retry.
+func IsNonTerminalHTTPStatusCode(resp *http.Response) bool {
+	return exported.HasStatusCode(resp,
+		http.StatusRequestTimeout,      // 408
+		http.StatusTooManyRequests,     // 429
+		http.StatusInternalServerError, // 500
+		http.StatusBadGateway,          // 502
+		http.StatusServiceUnavailable,  // 503
+		http.StatusGatewayTimeout,      // 504
+	)
+}

sdk/azcore/internal/pollers/util_test.go

@@ -184,3 +184,14 @@ func TestResultHelper(t *testing.T) {
 	require.Equal(t, "happy", widgetResult.Result)
 	require.Equal(t, 123, widgetResult.Precalculated)
 }
+
+func TestIsNonTerminalHTTPStatusCode(t *testing.T) {
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusRequestTimeout}))
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusTooManyRequests}))
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusInternalServerError}))
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusBadGateway}))
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusServiceUnavailable}))
+	require.True(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusGatewayTimeout}))
+	require.False(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusBadRequest}))
+	require.False(t, IsNonTerminalHTTPStatusCode(&http.Response{StatusCode: http.StatusNotImplemented}))
+}

sdk/azcore/runtime/poller.go

@@ -334,6 +334,11 @@ func (p *Poller[T]) Result(ctx context.Context) (res T, err error) {
 	err = p.op.Result(ctx, p.result)
 	var respErr *exported.ResponseError
 	if errors.As(err, &respErr) {
+		if pollers.IsNonTerminalHTTPStatusCode(respErr.RawResponse) {
+			// the request failed in a non-terminal way.
+			// don't cache the error or mark the Poller as done
+			return
+		}
 		// the LRO failed. record the error
 		p.err = err
 	} else if err != nil {

sdk/azcore/runtime/poller_test.go

@@ -758,7 +758,12 @@ func getPipeline(srv *mock.Server) Pipeline {
 		"test",
 		"v0.1.0",
 		PipelineOptions{PerRetry: []policy.Policy{NewLogPolicy(nil)}},
-		&policy.ClientOptions{Transport: srv},
+		&policy.ClientOptions{
+			Retry: policy.RetryOptions{
+				MaxRetryDelay: 1 * time.Second,
+			},
+			Transport: srv,
+		},
 	)
 }
 
@@ -1206,3 +1211,48 @@ func TestNewFakePoller(t *testing.T) {
 	require.NoError(t, err)
 	require.Nil(t, result.Field)
 }
+
+func TestNewPollerWithThrottling(t *testing.T) {
+	srv, close := mock.NewServer()
+	defer close()
+	srv.AppendResponse(mock.WithBody([]byte(statusInProgress)))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithBody([]byte(statusSucceeded)))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests))
+	srv.AppendResponse(mock.WithBody([]byte(successResp)))
+	resp, closed := initialResponse(context.Background(), http.MethodPatch, srv.URL(), strings.NewReader(provStateStarted))
+	resp.Header.Set(shared.HeaderAzureAsync, srv.URL())
+	resp.StatusCode = http.StatusCreated
+	pl := getPipeline(srv)
+	poller, err := NewPoller[mockType](resp, pl, nil)
+	require.NoError(t, err)
+	require.True(t, closed())
+	if pt := typeOfOpField(poller); pt != reflect.TypeOf((*async.Poller[mockType])(nil)) {
+		t.Fatalf("unexpected poller type %s", pt.String())
+	}
+	tk, err := poller.ResumeToken()
+	require.NoError(t, err)
+	poller, err = NewPollerFromResumeToken[mockType](tk, pl, nil)
+	require.NoError(t, err)
+	result, err := poller.PollUntilDone(context.Background(), &PollUntilDoneOptions{Frequency: time.Millisecond})
+	require.Zero(t, result)
+	var respErr *exported.ResponseError
+	require.ErrorAs(t, err, &respErr)
+	require.EqualValues(t, http.StatusTooManyRequests, respErr.StatusCode)
+	result, err = poller.PollUntilDone(context.Background(), &PollUntilDoneOptions{Frequency: time.Millisecond})
+	require.Zero(t, result)
+	require.ErrorAs(t, err, &respErr)
+	require.EqualValues(t, http.StatusTooManyRequests, respErr.StatusCode)
+	result, err = poller.PollUntilDone(context.Background(), &PollUntilDoneOptions{Frequency: time.Millisecond})
+	require.NoError(t, err)
+	require.NotNil(t, result.Field)
+	require.EqualValues(t, "value", *result.Field)
+}