Using the timeout from the context if available

- Makes PollingDuration optional

autorest/adal/sender.go

@@ -38,7 +38,7 @@ func (sf SenderFunc) Do(r *http.Request) (*http.Response, error) {
 	return sf(r)
 }
 
-// SendDecorator takes and possibily decorates, by wrapping, a Sender. Decorators may affect the
+// SendDecorator takes and possibly decorates, by wrapping, a Sender. Decorators may affect the
 // http.Request and pass it along or, first, pass the http.Request along then react to the
 // http.Response result.
 type SendDecorator func(Sender) Sender

autorest/azure/async.go

@@ -164,9 +164,17 @@ func (f Future) WaitForCompletion(ctx context.Context, client autorest.Client) e
 // 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) WaitForCompletionRef(ctx context.Context, client autorest.Client) error {
-	ctx, cancel := context.WithTimeout(ctx, client.PollingDuration)
-	defer cancel()
+func (f *Future) WaitForCompletionRef(inputCtx context.Context, client autorest.Client) error {
+
+	var ctx context.Context
+	if d := client.PollingDuration; d != nil {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, *d)
+		defer cancel()
+	} else {
+		ctx = inputCtx
+	}
+
 	done, err := f.Done(client)
 	for attempts := 0; !done; done, err = f.Done(client) {
 		if attempts >= client.RetryAttempts {

autorest/azure/async_test.go

@@ -793,9 +793,10 @@ func TestFuture_WaitForCompletion(t *testing.T) {
 	sender := mocks.NewSender()
 	sender.AppendAndRepeatResponse(r2, 2)
 	sender.AppendResponse(r3)
+	pollingDuration := autorest.DefaultPollingDuration
 	client := autorest.Client{
 		PollingDelay:    1 * time.Second,
-		PollingDuration: autorest.DefaultPollingDuration,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   autorest.DefaultRetryAttempts,
 		RetryDuration:   1 * time.Second,
 		Sender:          sender,
@@ -826,9 +827,10 @@ func TestFuture_WaitForCompletionRef(t *testing.T) {
 	sender := mocks.NewSender()
 	sender.AppendAndRepeatResponse(r2, 2)
 	sender.AppendResponse(r3)
+	pollingDuration := autorest.DefaultPollingDuration
 	client := autorest.Client{
 		PollingDelay:    1 * time.Second,
-		PollingDuration: autorest.DefaultPollingDuration,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   autorest.DefaultRetryAttempts,
 		RetryDuration:   1 * time.Second,
 		Sender:          sender,
@@ -863,9 +865,10 @@ func TestFuture_WaitForCompletionTimedOut(t *testing.T) {
 		t.Fatalf("failed to create future: %v", err)
 	}
 
+	pollingDuration := 2 * time.Second
 	client := autorest.Client{
 		PollingDelay:    autorest.DefaultPollingDelay,
-		PollingDuration: 2 * time.Second,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   autorest.DefaultRetryAttempts,
 		RetryDuration:   1 * time.Second,
 		Sender:          sender,
@@ -889,9 +892,10 @@ func TestFuture_WaitForCompletionRetriesExceeded(t *testing.T) {
 		t.Fatalf("failed to create future: %v", err)
 	}
 
+	pollingDuration := autorest.DefaultPollingDuration
 	client := autorest.Client{
 		PollingDelay:    autorest.DefaultPollingDelay,
-		PollingDuration: autorest.DefaultPollingDuration,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   autorest.DefaultRetryAttempts,
 		RetryDuration:   100 * time.Millisecond,
 		Sender:          sender,
@@ -914,9 +918,10 @@ func TestFuture_WaitForCompletionCancelled(t *testing.T) {
 		t.Fatalf("failed to create future: %v", err)
 	}
 
+	pollingDuration := autorest.DefaultPollingDuration
 	client := autorest.Client{
 		PollingDelay:    autorest.DefaultPollingDelay,
-		PollingDuration: autorest.DefaultPollingDuration,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   autorest.DefaultRetryAttempts,
 		RetryDuration:   autorest.DefaultRetryDuration,
 		Sender:          sender,

autorest/azure/rp.go

@@ -141,7 +141,7 @@ func register(client autorest.Client, originalReq *http.Request, re RequestError
 
 	// poll for registered provisioning state
 	now := time.Now()
-	for err == nil && time.Since(now) < client.PollingDuration {
+	for err == nil && (client.PollingDuration == nil || (client.PollingDuration != nil && time.Since(now) < *client.PollingDuration)) {
 		// taken from the resources SDK
 		// https://github.com/Azure/azure-sdk-for-go/blob/9f366792afa3e0ddaecdc860e793ba9d75e76c27/arm/resources/resources/providers.go#L45
 		preparer := autorest.CreatePreparer(
@@ -183,7 +183,7 @@ func register(client autorest.Client, originalReq *http.Request, re RequestError
 			return originalReq.Context().Err()
 		}
 	}
-	if !(time.Since(now) < client.PollingDuration) {
+	if client.PollingDuration != nil && !(time.Since(now) < *client.PollingDuration) {
 		return errors.New("polling for resource provider registration has exceeded the polling duration")
 	}
 	return err

autorest/azure/rp_test.go

@@ -59,10 +59,11 @@ func TestDoRetryWithRegistration(t *testing.T) {
 
 	req := mocks.NewRequestForURL("https://lol/subscriptions/rofl")
 	req.Body = mocks.NewBody("lolol")
+	pollingDuration := time.Second * 10
 	r, err := autorest.SendWithSender(client, req,
 		DoRetryWithRegistration(autorest.Client{
 			PollingDelay:    time.Second,
-			PollingDuration: time.Second * 10,
+			PollingDuration: &pollingDuration,
 			RetryAttempts:   5,
 			RetryDuration:   time.Second,
 			Sender:          client,
@@ -103,10 +104,11 @@ func TestDoRetrySkipRegistration(t *testing.T) {
 
 	req := mocks.NewRequestForURL("https://lol/subscriptions/rofl")
 	req.Body = mocks.NewBody("lolol")
+	pollingDuration := time.Second * 10
 	r, err := autorest.SendWithSender(client, req,
 		DoRetryWithRegistration(autorest.Client{
 			PollingDelay:    time.Second,
-			PollingDuration: time.Second * 10,
+			PollingDuration: &pollingDuration,
 			RetryAttempts:   5,
 			RetryDuration:   time.Second,
 			Sender:          client,
@@ -147,7 +149,7 @@ func TestDoRetryWithRegistration_CanBeCancelled(t *testing.T) {
 		_, err = autorest.SendWithSender(client, req,
 			DoRetryWithRegistration(autorest.Client{
 				PollingDelay:    time.Second,
-				PollingDuration: delay,
+				PollingDuration: &delay,
 				RetryAttempts:   5,
 				RetryDuration:   time.Second,
 				Sender:          client,

autorest/client.go

@@ -153,7 +153,8 @@ type Client struct {
 	PollingDelay time.Duration
 
 	// PollingDuration sets the maximum polling time after which an error is returned.
-	PollingDuration time.Duration
+	// if nil, the timeout from the Context is used
+	PollingDuration *time.Duration
 
 	// RetryAttempts sets the default number of retry attempts for client.
 	RetryAttempts int
@@ -174,9 +175,10 @@ type Client struct {
 // NewClientWithUserAgent returns an instance of a Client with the UserAgent set to the passed
 // string.
 func NewClientWithUserAgent(ua string) Client {
+	pollingDuration := DefaultPollingDuration
 	c := Client{
 		PollingDelay:    DefaultPollingDelay,
-		PollingDuration: DefaultPollingDuration,
+		PollingDuration: &pollingDuration,
 		RetryAttempts:   DefaultRetryAttempts,
 		RetryDuration:   DefaultRetryDuration,
 		UserAgent:       defaultUserAgent,