[WIP] Using the Context from the timeout if provided

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 != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(inputCtx, 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/rp.go

@@ -140,8 +140,8 @@ 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 {
+	registrationStartTime := time.Now()
+	for err == nil && (client.PollingDuration == 0 || (client.PollingDuration != 0 && time.Since(registrationStartTime) < 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 != 0 && !(time.Since(registrationStartTime) < client.PollingDuration) {
 		return errors.New("polling for resource provider registration has exceeded the polling duration")
 	}
 	return err

autorest/client.go

@@ -153,6 +153,7 @@ type Client struct {
 	PollingDelay time.Duration
 
 	// PollingDuration sets the maximum polling time after which an error is returned.
+	// if zero, the timeout from the Context is used
 	PollingDuration time.Duration
 
 	// RetryAttempts sets the default number of retry attempts for client.