Add WaitForCompletion to Future type. (#194)

* Add WaitForCompletion to Future type.

The WaitForCompletion method provides a default implementation that will
wait until a future has completed, is cancelled or times out.
Fixed a bug in Future.Done() to not update polling status if the HTTP
response isn't one of the expected values.
Added support to mock framework for returning errors in responses so we
can simulate things like transient network failures.

* Remove Go 1.6 from travis as it doesn't support the context package.

* Fixes based on feedback.

* Delete Retry-After header for HTTP 500 response.

* Track mock responses and replace hard-coded counts.

Author: jhendrixMSFT

.travis.yml

@@ -6,7 +6,6 @@ go:
   - 1.9
   - 1.8
   - 1.7
-  - 1.6
 
 install:
   - go get -u github.com/golang/lint/golint

CHANGELOG.md

@@ -1,5 +1,15 @@
 # CHANGELOG
 
+## v9.4.0
+
+### New Features
+
+- Added WaitForCompletion() to Future as a default polling implementation.
+
+### Bug Fixes
+
+- Method Future.Done() shouldn't update polling status for unexpected HTTP status codes.
+
 ## v9.3.1
 
 ### Bug Fixes

autorest/autorest.go

@@ -87,6 +87,9 @@ const (
 // ResponseHasStatusCode returns true if the status code in the HTTP Response is in the passed set
 // and false otherwise.
 func ResponseHasStatusCode(resp *http.Response, codes ...int) bool {
+	if resp == nil {
+		return false
+	}
 	return containsInt(codes, resp.StatusCode)
 }
 

autorest/azure/async.go

@@ -16,6 +16,7 @@ package azure
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -38,6 +39,8 @@ const (
 	operationSucceeded  string = "Succeeded"
 )
 
+var pollingCodes = [...]int{http.StatusAccepted, http.StatusCreated, http.StatusOK}
+
 // Future provides a mechanism to access the status and results of an asynchronous request.
 // Since futures are stateful they should be passed by value to avoid race conditions.
 type Future struct {
@@ -78,7 +81,7 @@ func (f *Future) Done(sender autorest.Sender) (bool, error) {
 
 	resp, err := sender.Do(f.req)
 	f.resp = resp
-	if err != nil {
+	if err != nil || !autorest.ResponseHasStatusCode(resp, pollingCodes[:]...) {
 		return false, err
 	}
 
@@ -117,6 +120,47 @@ func (f Future) GetPollingDelay() (time.Duration, bool) {
 	return d, true
 }
 
+// WaitForCompletion will return when one of the following conditions is met: the long
+// running operation has completed, the provided context is cancelled, or the client's
+// polling duration has been exceeded.  It will retry failed polling attempts based on
+// the retry value defined in the client up to the maximum retry attempts.
+func (f Future) WaitForCompletion(ctx context.Context, client autorest.Client) error {
+	ctx, cancel := context.WithTimeout(ctx, client.PollingDuration)
+	defer cancel()
+
+	done, err := f.Done(client)
+	for attempts := 0; !done; done, err = f.Done(client) {
+		if attempts >= client.RetryAttempts {
+			return autorest.NewErrorWithError(err, "azure", "WaitForCompletion", f.resp, "the number of retries has been exceeded")
+		}
+		// we want delayAttempt to be zero in the non-error case so
+		// that DelayForBackoff doesn't perform exponential back-off
+		var delayAttempt int
+		var delay time.Duration
+		if err == nil {
+			// check for Retry-After delay, if not present use the client's polling delay
+			var ok bool
+			delay, ok = f.GetPollingDelay()
+			if !ok {
+				delay = client.PollingDelay
+			}
+		} else {
+			// there was an error polling for status so perform exponential
+			// back-off based on the number of attempts using the client's retry
+			// duration.  update attempts after delayAttempt to avoid off-by-one.
+			delayAttempt = attempts
+			delay = client.RetryDuration
+			attempts++
+		}
+		// wait until the delay elapses or the context is cancelled
+		delayElapsed := autorest.DelayForBackoff(delay, delayAttempt, ctx.Done())
+		if !delayElapsed {
+			return autorest.NewErrorWithError(ctx.Err(), "azure", "WaitForCompletion", f.resp, "context has been cancelled")
+		}
+	}
+	return err
+}
+
 // if the operation failed the polling state will contain
 // error information and implements the error interface
 func (f *Future) errorInfo() error {
@@ -152,8 +196,7 @@ func DoPollForAsynchronous(delay time.Duration) autorest.SendDecorator {
 			if err != nil {
 				return resp, err
 			}
-			pollingCodes := []int{http.StatusAccepted, http.StatusCreated, http.StatusOK}
-			if !autorest.ResponseHasStatusCode(resp, pollingCodes...) {
+			if !autorest.ResponseHasStatusCode(resp, pollingCodes[:]...) {
 				return resp, nil
 			}
 

autorest/azure/async_test.go

@@ -15,7 +15,9 @@ package azure
 //  limitations under the License.
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -634,7 +636,7 @@ func TestDoPollForAsynchronous_PollsForStatusAccepted(t *testing.T) {
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -657,7 +659,7 @@ func TestDoPollForAsynchronous_PollsForStatusCreated(t *testing.T) {
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -681,7 +683,7 @@ func TestDoPollForAsynchronous_PollsUntilProvisioningStatusTerminates(t *testing
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -705,7 +707,7 @@ func TestDoPollForAsynchronous_PollsUntilProvisioningStatusSucceeds(t *testing.T
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -726,7 +728,7 @@ func TestDoPollForAsynchronous_PollsUntilOperationResourceHasTerminated(t *testi
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -747,7 +749,7 @@ func TestDoPollForAsynchronous_PollsUntilOperationResourceHasSucceeded(t *testin
 	r, _ := autorest.SendWithSender(client, mocks.NewRequest(),
 		DoPollForAsynchronous(time.Millisecond))
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: DoPollForAsynchronous stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -876,7 +878,7 @@ func TestDoPollForAsynchronous_ReturnsErrorForLastErrorResponse(t *testing.T) {
 	r1.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
 	r2 := newProvisioningStatusResponse("busy")
 	r2.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
-	r3 := newAsynchronousResponseWithError()
+	r3 := newAsynchronousResponseWithError("400 Bad Request", http.StatusBadRequest)
 	r3.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
 
 	client := mocks.NewSender()
@@ -923,7 +925,7 @@ func TestDoPollForAsynchronous_ReturnsOperationResourceErrorForFailedOperations(
 
 func TestDoPollForAsynchronous_ReturnsErrorForFirstPutRequest(t *testing.T) {
 	// Return 400 bad response with error code and message in first put
-	r1 := newAsynchronousResponseWithError()
+	r1 := newAsynchronousResponseWithError("400 Bad Request", http.StatusBadRequest)
 	client := mocks.NewSender()
 	client.AppendResponse(r1)
 
@@ -1044,7 +1046,7 @@ func TestFuture_PollsUntilProvisioningStatusSucceeds(t *testing.T) {
 		time.Sleep(delay)
 	}
 
-	if client.Attempts() < 4 {
+	if client.Attempts() < client.NumResponses() {
 		t.Fatalf("azure: TestFuture stopped polling before receiving a terminated OperationResource")
 	}
 
@@ -1096,6 +1098,129 @@ func TestFuture_Marshalling(t *testing.T) {
 	}
 }
 
+func TestFuture_WaitForCompletion(t *testing.T) {
+	r1 := newAsynchronousResponse()
+	r1.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+	r2 := newProvisioningStatusResponse("busy")
+	r2.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+	r3 := newAsynchronousResponseWithError("Internal server error", http.StatusInternalServerError)
+	r3.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+	r3.Header.Del(http.CanonicalHeaderKey("Retry-After"))
+	r4 := newProvisioningStatusResponse(operationSucceeded)
+	r4.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+
+	sender := mocks.NewSender()
+	sender.AppendResponse(r1)
+	sender.AppendError(errors.New("transient network failure"))
+	sender.AppendAndRepeatResponse(r2, 2)
+	sender.AppendResponse(r3)
+	sender.AppendResponse(r4)
+
+	future := NewFuture(mocks.NewRequest())
+
+	client := autorest.Client{
+		PollingDelay:    1 * time.Second,
+		PollingDuration: autorest.DefaultPollingDuration,
+		RetryAttempts:   autorest.DefaultRetryAttempts,
+		RetryDuration:   1 * time.Second,
+		Sender:          sender,
+	}
+
+	err := future.WaitForCompletion(context.Background(), client)
+	if err != nil {
+		t.Fatalf("azure: WaitForCompletion returned non-nil error")
+	}
+
+	if sender.Attempts() < sender.NumResponses() {
+		t.Fatalf("azure: TestFuture stopped polling before receiving a terminated OperationResource")
+	}
+
+	autorest.Respond(future.Response(),
+		autorest.ByClosing())
+}
+
+func TestFuture_WaitForCompletionTimedOut(t *testing.T) {
+	r1 := newAsynchronousResponse()
+	r1.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+	r2 := newProvisioningStatusResponse("busy")
+	r2.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+
+	sender := mocks.NewSender()
+	sender.AppendResponse(r1)
+	sender.AppendAndRepeatResponseWithDelay(r2, 1*time.Second, 5)
+
+	future := NewFuture(mocks.NewRequest())
+
+	client := autorest.Client{
+		PollingDelay:    autorest.DefaultPollingDelay,
+		PollingDuration: 2 * time.Second,
+		RetryAttempts:   autorest.DefaultRetryAttempts,
+		RetryDuration:   1 * time.Second,
+		Sender:          sender,
+	}
+
+	err := future.WaitForCompletion(context.Background(), client)
+	if err == nil {
+		t.Fatalf("azure: WaitForCompletion returned nil error, should have timed out")
+	}
+}
+
+func TestFuture_WaitForCompletionRetriesExceeded(t *testing.T) {
+	r1 := newAsynchronousResponse()
+	r1.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+
+	sender := mocks.NewSender()
+	sender.AppendResponse(r1)
+	sender.AppendAndRepeatError(errors.New("transient network failure"), autorest.DefaultRetryAttempts+1)
+
+	future := NewFuture(mocks.NewRequest())
+
+	client := autorest.Client{
+		PollingDelay:    autorest.DefaultPollingDelay,
+		PollingDuration: autorest.DefaultPollingDuration,
+		RetryAttempts:   autorest.DefaultRetryAttempts,
+		RetryDuration:   100 * time.Millisecond,
+		Sender:          sender,
+	}
+
+	err := future.WaitForCompletion(context.Background(), client)
+	if err == nil {
+		t.Fatalf("azure: WaitForCompletion returned nil error, should have errored out")
+	}
+}
+
+func TestFuture_WaitForCompletionCancelled(t *testing.T) {
+	r1 := newAsynchronousResponse()
+	r1.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+	r2 := newProvisioningStatusResponse("busy")
+	r2.Header.Del(http.CanonicalHeaderKey(headerAsyncOperation))
+
+	sender := mocks.NewSender()
+	sender.AppendResponse(r1)
+	sender.AppendAndRepeatResponseWithDelay(r2, 1*time.Second, 5)
+
+	future := NewFuture(mocks.NewRequest())
+
+	client := autorest.Client{
+		PollingDelay:    autorest.DefaultPollingDelay,
+		PollingDuration: autorest.DefaultPollingDuration,
+		RetryAttempts:   autorest.DefaultRetryAttempts,
+		RetryDuration:   autorest.DefaultRetryDuration,
+		Sender:          sender,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(2 * time.Second)
+		cancel()
+	}()
+
+	err := future.WaitForCompletion(ctx, client)
+	if err == nil {
+		t.Fatalf("azure: WaitForCompletion returned nil error, should have been cancelled")
+	}
+}
+
 const (
 	operationResourceIllegal = `
 	This is not JSON and should fail...badly.
@@ -1171,8 +1296,8 @@ func newAsynchronousResponse() *http.Response {
 	return r
 }
 
-func newAsynchronousResponseWithError() *http.Response {
-	r := mocks.NewResponseWithStatus("400 Bad Request", http.StatusBadRequest)
+func newAsynchronousResponseWithError(response string, status int) *http.Response {
+	r := mocks.NewResponseWithStatus(response, status)
 	mocks.SetRetryHeader(r, retryDelay)
 	r.Request = mocks.NewRequestForURL(mocks.TestURL)
 	r.Body = mocks.NewBody(errorResponse)

autorest/client.go

@@ -35,6 +35,9 @@ const (
 
 	// DefaultRetryAttempts is number of attempts for retry status codes (5xx).
 	DefaultRetryAttempts = 3
+
+	// DefaultRetryDuration is the duration to wait between retries.
+	DefaultRetryDuration = 30 * time.Second
 )
 
 var (
@@ -172,7 +175,7 @@ func NewClientWithUserAgent(ua string) Client {
 		PollingDelay:    DefaultPollingDelay,
 		PollingDuration: DefaultPollingDuration,
 		RetryAttempts:   DefaultRetryAttempts,
-		RetryDuration:   30 * time.Second,
+		RetryDuration:   DefaultRetryDuration,
 		UserAgent:       defaultUserAgent,
 	}
 	c.Sender = c.sender()