Fixes to runtime.Poller (Azure#17837)

* Fixes to runtime.Poller

Fixed the behavior of Poll and Result methods and updated their doc
comments explaining the behavior.
Clarified some other Poller doc comments and improved error messages.

* clean-up

Author: jhendrixMSFT

sdk/azcore/CHANGELOG.md

@@ -22,6 +22,9 @@
 ### Bugs Fixed
 * When per-try timeouts are enabled, only cancel the context after the body has been read and closed.
 * The `Operation-Location` poller now properly handles `final-state-via` values.
+* Improvements in `runtime.Poller[T]`
+  * `Poll()` shouldn't cache errors, allowing for additional retries when in a non-terminal state.
+  * `Result()` will cache the terminal result or error but not transient errors, allowing for additional retries.
 
 ### Other Changes
 * The internal poller implementation has been refactored.

sdk/azcore/runtime/poller.go

@@ -183,6 +183,7 @@ type Poller[T any] struct {
 	resp   *http.Response
 	err    error
 	result *T
+	done   bool
 }
 
 // PollUntilDoneOptions contains the optional values for the Poller[T].PollUntilDone() method.
@@ -193,6 +194,7 @@ type PollUntilDoneOptions struct {
 }
 
 // PollUntilDone will poll the service endpoint until a terminal state is reached, an error is received, or the context expires.
+// It internally uses Poll(), Done(), and Result() in its polling loop, sleeping for the specified duration between intervals.
 // options: pass nil to accept the default values.
 // NOTE: the default polling frequency is 30 seconds which works well for most operations.  However, some operations might
 // benefit from a shorter or longer duration.
@@ -250,46 +252,69 @@ func (p *Poller[T]) PollUntilDone(ctx context.Context, options *PollUntilDoneOpt
 }
 
 // Poll fetches the latest state of the LRO.  It returns an HTTP response or error.
-// If the LRO has completed successfully, the poller's state is updated and the HTTP
-// response is returned.
-// If the LRO has completed with failure or was cancelled, the poller's state is
-// updated and the error is returned.
-// If the LRO has not reached a terminal state, the poller's state is updated and
-// the latest HTTP response is returned.
+// If Poll succeeds, the poller's state is updated and the HTTP response is returned.
 // If Poll fails, the poller's state is unmodified and the error is returned.
-// Calling Poll on an LRO that has reached a terminal state will return the final
-// HTTP response or error.
+// Calling Poll on an LRO that has reached a terminal state will return the last HTTP response.
 func (p *Poller[T]) Poll(ctx context.Context) (*http.Response, error) {
 	if p.Done() {
 		// the LRO has reached a terminal state, don't poll again
-		if p.err != nil {
-			return nil, p.err
-		}
 		return p.resp, nil
 	}
-	p.resp, p.err = p.op.Poll(ctx)
-	return p.resp, p.err
+	resp, err := p.op.Poll(ctx)
+	if err != nil {
+		return nil, err
+	}
+	p.resp = resp
+	return p.resp, nil
 }
 
 // Done returns true if the LRO has reached a terminal state.
+// Once a terminal state is reached, call Result().
 func (p *Poller[T]) Done() bool {
 	return p.op.Done()
 }
 
 // Result returns the result of the LRO and is meant to be used in conjunction with Poll and Done.
+// If the LRO completed successfully, a populated instance of T is returned.
+// If the LRO failed or was canceled, an *azcore.ResponseError error is returned.
 // Calling this on an LRO in a non-terminal state will return an error.
 func (p *Poller[T]) Result(ctx context.Context) (T, error) {
 	if !p.Done() {
-		return *new(T), errors.New("cannot return a final response from a poller in a non-terminal state")
+		return *new(T), errors.New("poller is in a non-terminal state")
+	}
+	if p.done {
+		// the result has already been retrieved, return the cached value
+		if p.err != nil {
+			return *new(T), p.err
+		}
+		return *p.result, nil
+	}
+	res, err := p.op.Result(ctx, p.result)
+	var respErr *exported.ResponseError
+	if errors.As(err, &respErr) {
+		// the LRO failed. record the error
+		p.err = err
+	} else if err != nil {
+		// the call to Result failed, don't cache anything in this case
+		return *new(T), err
+	} else {
+		// the LRO succeeded. record the result
+		p.result = &res
+	}
+	p.done = true
+	if p.err != nil {
+		return *new(T), p.err
 	}
-	return p.op.Result(ctx, p.result)
+	return *p.result, nil
 }
 
 // ResumeToken returns a value representing the poller that can be used to resume
 // the LRO at a later time. ResumeTokens are unique per service operation.
+// The token's format should be considered opaque and is subject to change.
+// Calling this on an LRO in a terminal state will return an error.
 func (p *Poller[T]) ResumeToken() (string, error) {
 	if p.Done() {
-		return "", errors.New("cannot create a ResumeToken from a poller in a terminal state")
+		return "", errors.New("poller is in a terminal state")
 	}
 	tk, err := pollers.NewResumeToken[T](p.op)
 	if err != nil {

sdk/azcore/runtime/poller_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
+	"github.com/stretchr/testify/require"
 )
 
 type none struct{}
@@ -178,6 +179,16 @@ func TestLocPollerCancelled(t *testing.T) {
 	if w.Size != 0 {
 		t.Fatalf("unexpected widget size %d", w.Size)
 	}
+	w, err = lro.Result(context.Background())
+	if err == nil {
+		t.Fatal("unexpected nil error")
+	}
+	if _, ok := err.(*exported.ResponseError); !ok {
+		t.Fatal("expected pollerError")
+	}
+	if w.Size != 0 {
+		t.Fatalf("unexpected widget size %d", w.Size)
+	}
 }
 
 func TestLocPollerWithError(t *testing.T) {
@@ -380,6 +391,68 @@ func TestOpPollerWithWidgetPUT(t *testing.T) {
 	}
 }
 
+type nonRetriableError struct {
+	Msg string
+}
+
+func (n *nonRetriableError) Error() string {
+	return n.Msg
+}
+
+func (*nonRetriableError) NonRetriable() {
+	// prevent the retry policy from masking this transient error
+}
+
+func TestOpPollerWithWidgetFinalGetError(t *testing.T) {
+	srv, close := mock.NewServer()
+	srv.AppendResponse(mock.WithStatusCode(http.StatusAccepted), mock.WithBody([]byte(`{"status": "InProgress"}`)))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK), mock.WithBody([]byte(`{"status": "Succeeded"}`)))
+	// PUT and PATCH state that a final GET will happen
+	// the first attempt at a final GET returns an error
+	srv.AppendError(&nonRetriableError{Msg: "failed attempt"})
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK), mock.WithBody([]byte(`{"size": 2}`)))
+	defer close()
+
+	reqURL, err := url.Parse(srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, closed := mock.NewTrackedCloser(http.NoBody)
+	firstResp := &http.Response{
+		Body:       body,
+		StatusCode: http.StatusAccepted,
+		Header: http.Header{
+			"Operation-Location": []string{srv.URL()},
+		},
+		Request: &http.Request{
+			Method: http.MethodPut,
+			URL:    reqURL,
+		},
+	}
+	pl := newTestPipeline(&policy.ClientOptions{Transport: srv})
+	lro, err := NewPoller[widget](firstResp, pl, nil)
+	require.Nil(t, err)
+	require.True(t, closed(), "initial response body wasn't closed")
+
+	resp, err := lro.Poll(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, http.StatusAccepted, resp.StatusCode)
+	require.False(t, lro.Done())
+
+	resp, err = lro.Poll(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	require.True(t, lro.Done())
+
+	w, err := lro.Result(context.Background())
+	require.Error(t, err)
+	require.Empty(t, w)
+
+	w, err = lro.Result(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, w.Size, 2)
+}
+
 func TestOpPollerWithWidgetPOSTLocation(t *testing.T) {
 	srv, close := mock.NewServer()
 	srv.AppendResponse(mock.WithStatusCode(http.StatusAccepted), mock.WithBody([]byte(`{"status": "InProgress"}`)))
@@ -751,6 +824,13 @@ func TestNewPollerAsync(t *testing.T) {
 	if v := *result.Field; v != "value" {
 		t.Fatalf("unexpected value %s", v)
 	}
+	result, err = poller.Result(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := *result.Field; v != "value" {
+		t.Fatalf("unexpected value %s", v)
+	}
 }
 
 func TestNewPollerBody(t *testing.T) {
@@ -831,6 +911,13 @@ func TestNewPollerARMLoc(t *testing.T) {
 	if v := *result.Field; v != "value" {
 		t.Fatalf("unexpected value %s", v)
 	}
+	result, err = poller.Result(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := *result.Field; v != "value" {
+		t.Fatalf("unexpected value %s", v)
+	}
 }
 
 func TestNewPollerInitialRetryAfter(t *testing.T) {