From 507aa2ed61ac5dc343dd7f2486900fe45ba8f1b0 Mon Sep 17 00:00:00 2001 From: Joel Hendrix Date: Thu, 28 Mar 2024 15:34:19 -0700 Subject: [PATCH] Don't consider 429 as terminal failure for Location poller If the polling request is being throttled, preserve the current state and return the *http.Response. --- sdk/azcore/CHANGELOG.md | 1 + sdk/azcore/internal/pollers/loc/loc.go | 4 ++++ sdk/azcore/internal/pollers/loc/loc_test.go | 24 +++++++++++++++++++++ 3 files changed, 29 insertions(+) diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md index 7494c33b54ce..bab76c055cb6 100644 --- a/sdk/azcore/CHANGELOG.md +++ b/sdk/azcore/CHANGELOG.md @@ -13,6 +13,7 @@ ### Bugs Fixed * `runtime.SetMultipartFormData` won't try to stringify `[]byte` values. +* Pollers that use the `Location` header won't consider `http.StatusTooManyRequests` a terminal failure. ### Other Changes diff --git a/sdk/azcore/internal/pollers/loc/loc.go b/sdk/azcore/internal/pollers/loc/loc.go index e83c55bf042b..022abaa4f40f 100644 --- a/sdk/azcore/internal/pollers/loc/loc.go +++ b/sdk/azcore/internal/pollers/loc/loc.go @@ -103,6 +103,10 @@ 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.StatusTooManyRequests { + // the request is being throttled. we DO NOT want to + // include this as terminal failure so preserve the + // existing state and return the response. } else { p.CurState = poller.StatusFailed } diff --git a/sdk/azcore/internal/pollers/loc/loc_test.go b/sdk/azcore/internal/pollers/loc/loc_test.go index 12915dd379d2..137292999ec1 100644 --- a/sdk/azcore/internal/pollers/loc/loc_test.go +++ b/sdk/azcore/internal/pollers/loc/loc_test.go @@ -16,6 +16,7 @@ import ( "github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/exported" "github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared" + "github.com/Azure/azure-sdk-for-go/sdk/internal/mock" "github.com/Azure/azure-sdk-for-go/sdk/internal/poller" "github.com/stretchr/testify/require" ) @@ -173,3 +174,26 @@ func TestSynchronousCompletion(t *testing.T) { require.Equal(t, poller.StatusSucceeded, lp.CurState) require.True(t, lp.Done()) } + +func TestWithThrottling(t *testing.T) { + srv, close := mock.NewServer() + defer close() + srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests)) + srv.AppendResponse(mock.WithStatusCode(http.StatusAccepted)) + srv.AppendResponse(mock.WithStatusCode(http.StatusTooManyRequests)) + srv.AppendResponse(mock.WithStatusCode(http.StatusOK)) + resp := initialResponse() + resp.Header.Set(shared.HeaderLocation, srv.URL()) + lp, err := New[struct{}](exported.NewPipeline(shared.TransportFunc(func(req *http.Request) (*http.Response, error) { + return srv.Do(req) + })), resp) + require.NoError(t, err) + respCount := 0 + for !lp.Done() { + _, err = lp.Poll(context.Background()) + require.NoError(t, err) + respCount++ + } + require.EqualValues(t, 4, respCount) + require.EqualValues(t, poller.StatusSucceeded, lp.CurState) +}